gruff-0.2: src/QuadTree.hs

module QuadTree
  ( Child(..), north, south, west, east
  , Quad(..), root, child, children, parent, parents
  , filename, unsafeName
  , Square(..), square, Point(..), contains, Region(..), expand, outside
  , quads
  ) where

import Data.Bits (bit, shiftL, shiftR, testBit, (.|.))
import Data.List (unfoldr, sort)
import Data.Ratio ((%))

data Child = NorthWest | NorthEast | SouthWest | SouthEast
  deriving (Read, Show, Eq, Ord, Enum, Bounded)

north, south, west, east :: Child -> Bool
east c = fromEnum c `testBit` 0
south c = fromEnum c `testBit` 1
north = not . south
west = not . east

data Quad = Quad{ quadLevel :: !Int, quadWest, quadNorth :: !Integer }
  deriving (Read, Show, Eq, Ord)

root :: Quad
root = Quad{ quadLevel = 0, quadWest = 0, quadNorth = 0 }

child :: Child -> Quad -> Quad
child c Quad{ quadLevel = l, quadWest = x, quadNorth = y } = Quad
  { quadLevel = l + 1
  , quadWest  = x `shiftL` 1 .|. (fromIntegral . fromEnum . east ) c
  , quadNorth = y `shiftL` 1 .|. (fromIntegral . fromEnum . south) c
  }

children :: [Child] -> Quad
children = foldr child root

parent :: Quad -> Maybe (Child, Quad)
parent Quad{ quadLevel = l, quadWest = x, quadNorth = y }
  | l > 0  = Just
      ( toEnum (fromEnum (y `testBit` 0) `shiftL` 1 .|. fromEnum (x `testBit` 0))
      , Quad{ quadLevel = l - 1, quadWest = x `shiftR` 1, quadNorth = y `shiftR` 1 }
      )
  | otherwise = Nothing

parents :: Quad -> [Child]
parents = unfoldr parent

filename :: Quad -> Maybe ([FilePath], FilePath)
filename q
  | not (0 <= quadNorth q && quadNorth q < bit (quadLevel q) && 0 <= quadWest q && quadWest q < bit (quadLevel q)) = Nothing
  | null cs = Nothing
  | otherwise = Just (init cs, last cs)
  where
    cs = chunk 2 . map unsafeName . chunk 2 . reverse . parents $ q

unsafeName :: [Child] -> Char
unsafeName [c]   = ['a'..'d'] !! (fromEnum c)
unsafeName [c,d] = ['e'..'t'] !! (fromEnum c `shiftL` 2 .|. fromEnum d)
unsafeName _ = error "QuadTree.unsafeName"

chunk :: Int -> [a] -> [[a]]
chunk _ [] = []
chunk n xs = let (ys, zs) = splitAt n xs in ys : chunk n zs

data Square = Square{ squareSize, squareWest, squareNorth :: !Rational }
  deriving (Read, Show, Eq, Ord)

square :: Square -> Quad -> Square
square Square{ squareSize = s0, squareWest = x0, squareNorth = y0 } Quad{ quadLevel = l, quadWest = x, quadNorth = y } =
  Square{ squareSize = s0 / fromInteger r, squareWest = x0 + s0 * (x % r), squareNorth = y0 + s0 * (y % r) } where r = bit l

data Region = Region{ regionNorth, regionSouth, regionWest, regionEast :: !Rational }
  deriving (Read, Show, Eq, Ord)

expand :: Rational -> Region -> Region
expand f r =
  let (x,  y ) = ((regionEast r + regionWest r) / 2, (regionNorth r + regionSouth r) / 2)
      (rx, ry) = ((regionEast r - regionWest r) / 2, (regionNorth r - regionSouth r) / 2)
  in  Region{ regionNorth = y + f * ry, regionSouth = y - f * ry, regionEast = x + f * rx, regionWest = x - f * rx }

outside :: Region -> Square -> Bool
outside r s
  =  regionSouth r < squareNorth s
  || regionEast  r < squareWest  s
  || regionNorth r > squareNorth s + squareSize s
  || regionWest  r > squareWest  s + squareSize s

data Point = Point{ pointWest, pointNorth :: !Rational }
  deriving (Read, Show, Eq, Ord)

contains :: Point -> Square -> Bool
contains p s
  =  squareNorth s <= pointNorth p && pointNorth p <= squareNorth s + squareSize s
  && squareWest  s <= pointWest  p && pointWest  p <= squareWest  s + squareSize s

quads :: Square -> Region -> Int -> [Quad]
quads rootSquare region level =
  [ Quad{ quadLevel = level, quadWest = w, quadNorth = n }
  | n <- [ floor nlo' .. ceiling nhi' - 1]
  , w <- [ floor wlo' .. ceiling whi' - 1]
  ]
  where
    [nlo', nhi'] = sort [nlo, nhi]
    [wlo', whi'] = sort [wlo, whi]
    nlo = (regionSouth region - squareNorth rootSquare) / squareSize rootSquare * l
    nhi = (regionNorth region - squareNorth rootSquare) / squareSize rootSquare * l
    wlo = (regionWest  region - squareWest  rootSquare) / squareSize rootSquare * l
    whi = (regionEast  region - squareWest  rootSquare) / squareSize rootSquare * l
    l = bit level % 1