gruff-0.1: QuadTree.hs

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

import Control.Monad (liftM2)
import Data.Bits (bit, shiftL, shiftR, testBit, (.|.))
import Data.List (unfoldr)
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
  | 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)

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 = quadss rootSquare region !! level

quadss :: Square -> Region -> [[Quad]]
quadss rootSquare region =
  iterate (filter (not . outside region . square rootSquare) . liftM2 child [minBound .. maxBound]) [root]