prepare to make cache size runtime settable
[ruff:gruff.git] / QuadTree.hs
1 module QuadTree
2   ( Child(..), north, south, west, east
3   , Quad(..), root, child, children, parent, parents
4   , filename, unsafeName
5   , Square(..), square, Point(..), contains, Region(..), expand, outside
6   , quads
7   ) where
8
9 import Data.Bits (bit, shiftL, shiftR, testBit, (.|.))
10 import Data.List (unfoldr, sort)
11 import Data.Ratio ((%))
12
13 data Child = NorthWest | NorthEast | SouthWest | SouthEast
14   deriving (Read, Show, Eq, Ord, Enum, Bounded)
15
16 north, south, west, east :: Child -> Bool
17 east c = fromEnum c `testBit` 0
18 south c = fromEnum c `testBit` 1
19 north = not . south
20 west = not . east
21
22 data Quad = Quad{ quadLevel :: !Int, quadWest, quadNorth :: !Integer }
23   deriving (Read, Show, Eq, Ord)
24
25 root :: Quad
26 root = Quad{ quadLevel = 0, quadWest = 0, quadNorth = 0 }
27
28 child :: Child -> Quad -> Quad
29 child c Quad{ quadLevel = l, quadWest = x, quadNorth = y } = Quad
30   { quadLevel = l + 1
31   , quadWest  = x `shiftL` 1 .|. (fromIntegral . fromEnum . east ) c
32   , quadNorth = y `shiftL` 1 .|. (fromIntegral . fromEnum . south) c
33   }
34
35 children :: [Child] -> Quad
36 children = foldr child root
37
38 parent :: Quad -> Maybe (Child, Quad)
39 parent Quad{ quadLevel = l, quadWest = x, quadNorth = y }
40   | l > 0  = Just
41       ( toEnum (fromEnum (y `testBit` 0) `shiftL` 1 .|. fromEnum (x `testBit` 0))
42       , Quad{ quadLevel = l - 1, quadWest = x `shiftR` 1, quadNorth = y `shiftR` 1 }
43       )
44   | otherwise = Nothing
45
46 parents :: Quad -> [Child]
47 parents = unfoldr parent
48
49 filename :: Quad -> Maybe ([FilePath], FilePath)
50 filename q
51   | not (0 <= quadNorth q && quadNorth q < bit (quadLevel q) && 0 <= quadWest q && quadWest q < bit (quadLevel q)) = Nothing
52   | null cs = Nothing
53   | otherwise = Just (init cs, last cs)
54   where
55     cs = chunk 2 . map unsafeName . chunk 2 . reverse . parents $ q
56
57 unsafeName :: [Child] -> Char
58 unsafeName [c]   = ['a'..'d'] !! (fromEnum c)
59 unsafeName [c,d] = ['e'..'t'] !! (fromEnum c `shiftL` 2 .|. fromEnum d)
60 unsafeName _ = error "QuadTree.unsafeName"
61
62 chunk :: Int -> [a] -> [[a]]
63 chunk _ [] = []
64 chunk n xs = let (ys, zs) = splitAt n xs in ys : chunk n zs
65
66 data Square = Square{ squareSize, squareWest, squareNorth :: !Rational }
67   deriving (Read, Show, Eq, Ord)
68
69 square :: Square -> Quad -> Square
70 square Square{ squareSize = s0, squareWest = x0, squareNorth = y0 } Quad{ quadLevel = l, quadWest = x, quadNorth = y } =
71   Square{ squareSize = s0 / fromInteger r, squareWest = x0 + s0 * (x % r), squareNorth = y0 + s0 * (y % r) } where r = bit l
72
73 data Region = Region{ regionNorth, regionSouth, regionWest, regionEast :: !Rational }
74   deriving (Read, Show, Eq, Ord)
75
76 expand :: Rational -> Region -> Region
77 expand f r =
78   let (x,  y ) = ((regionEast r + regionWest r) / 2, (regionNorth r + regionSouth r) / 2)
79       (rx, ry) = ((regionEast r - regionWest r) / 2, (regionNorth r - regionSouth r) / 2)
80   in  Region{ regionNorth = y + f * ry, regionSouth = y - f * ry, regionEast = x + f * rx, regionWest = x - f * rx }
81
82 outside :: Region -> Square -> Bool
83 outside r s
84   =  regionSouth r < squareNorth s
85   || regionEast  r < squareWest  s
86   || regionNorth r > squareNorth s + squareSize s
87   || regionWest  r > squareWest  s + squareSize s
88
89 data Point = Point{ pointWest, pointNorth :: !Rational }
90   deriving (Read, Show, Eq, Ord)
91
92 contains :: Point -> Square -> Bool
93 contains p s
94   =  squareNorth s <= pointNorth p && pointNorth p <= squareNorth s + squareSize s
95   && squareWest  s <= pointWest  p && pointWest  p <= squareWest  s + squareSize s
96
97 quads :: Square -> Region -> Int -> [Quad]
98 quads rootSquare region level =
99   [ Quad{ quadLevel = level, quadWest = w, quadNorth = n }
100   | n <- [ floor nlo' .. ceiling nhi' - 1]
101   , w <- [ floor wlo' .. ceiling whi' - 1]
102   ]
103   where
104     [nlo', nhi'] = sort [nlo, nhi]
105     [wlo', whi'] = sort [wlo, whi]
106     nlo = (regionSouth region - squareNorth rootSquare) / squareSize rootSquare * l
107     nhi = (regionNorth region - squareNorth rootSquare) / squareSize rootSquare * l
108     wlo = (regionWest  region - squareWest  rootSquare) / squareSize rootSquare * l
109     whi = (regionEast  region - squareWest  rootSquare) / squareSize rootSquare * l
110     l = bit level % 1