# Ticket #291: QuadTree.hs

File QuadTree.hs, 14.2 kB (added by Kron, 13 months ago) |
---|

Line | |
---|---|

1 | |

2 | {-# OPTIONS_HADDOCK show-extensions #-} |

3 | |

4 | {-# LANGUAGE Safe #-} |

5 | |

6 | {-# LANGUAGE Rank2Types #-} |

7 | {-# LANGUAGE ScopedTypeVariables #-} |

8 | |

9 | {-| |

10 | Module : QuadTree |

11 | Description : Region quadtrees with Lens support. |

12 | Copyright : (c) Ashley Moni, 2014 |

13 | License : GPL-3 |

14 | Maintainer : Ashley Moni <y30099923@gmail.com> |

15 | Stability : Experimental |

16 | Portability : Rank 2 Types, Scoped Type Variables |

17 | |

18 | The purpose of this module is to provide discrete region quadtrees |

19 | that can be used as simple functional alternatives to 2D arrays, |

20 | with lens support. |

21 | |

22 | @ |

23 | test = set ('atLocation' (0,0)) \'d\' $ |

24 | set ('atLocation' (5,5)) \'c\' $ |

25 | set ('atLocation' (3,2)) \'b\' $ |

26 | set ('atLocation' (2,4)) \'a\' $ |

27 | 'makeTree' (6,6) \'.\' |

28 | @ |

29 | |

30 | >>> printTree id test |

31 | d..... |

32 | ...... |

33 | ...b.. |

34 | ...... |

35 | ..a... |

36 | .....c |

37 | -} |

38 | |

39 | module Data.QuadTree ( |

40 | -- * Data Type & Constructor |

41 | QuadTree, makeTree, |

42 | -- * Index access |

43 | -- $locations |

44 | Location, getLocation, setLocation, atLocation, |

45 | -- * Functor |

46 | fuseTree, tmap, |

47 | -- * Foldable |

48 | -- $foldables |

49 | |

50 | -- ** Tiles |

51 | -- $tiles |

52 | Region, Tile, |

53 | foldTiles, tile, expand, |

54 | -- ** Tile-based utilities |

55 | filterTree, sortTreeBy, filterTiles, sortTilesBy, |

56 | -- * Printers |

57 | showTree, printTree, |

58 | -- * Miscellaneous helpers |

59 | outOfBounds, treeDimensions, regionArea, inRegion ) where |

60 | |

61 | import Control.Lens.Type (Lens') |

62 | import Control.Lens.Lens (lens) |

63 | import Control.Lens.Setter (set) |

64 | |

65 | import Data.List (find, sortBy) |

66 | import Data.Maybe (fromJust) |

67 | import Data.Function (on) |

68 | import Data.Composition ((.:)) |

69 | |

70 | -- Foldable: |

71 | import Data.Foldable (Foldable, foldr) |

72 | import Prelude hiding (foldr) |

73 | |

74 | ---- Structures: |

75 | |

76 | -- |Tuple corresponds to (X, Y) co-ordinates. |

77 | |

78 | type Location = (Int, Int) |

79 | |

80 | -- |The eponymous data type. |

81 | -- |

82 | -- 'QuadTree' is itself a wrapper around an internal tree structure |

83 | -- along with spatial metadata about the boundaries and depth of the |

84 | -- 2D area it maps to. |

85 | |

86 | data QuadTree a = Wrapper { wrappedTree :: Quadrant a |

87 | , treeLength :: Int |

88 | , treeWidth :: Int |

89 | , treeDepth :: Int } |

90 | deriving (Show, Read) |

91 | |

92 | instance Functor QuadTree where |

93 | fmap fn = onQuads $ fmap fn |

94 | |

95 | instance Foldable QuadTree where |

96 | foldr = foldTree |

97 | |

98 | -- |

99 | |

100 | data Quadrant a = Leaf a |

101 | | Node (Quadrant a) |

102 | (Quadrant a) |

103 | (Quadrant a) |

104 | (Quadrant a) |

105 | deriving (Show, Read) |

106 | |

107 | instance Functor Quadrant where |

108 | fmap fn (Leaf x) = Leaf (fn x) |

109 | fmap fn (Node a b c d) = Node (fmap fn a) |

110 | (fmap fn b) |

111 | (fmap fn c) |

112 | (fmap fn d) |

113 | |

114 | ---- Index access: |

115 | |

116 | -- $locations |

117 | -- This provides an array-style interface to the 'QuadTree', albeit |

118 | -- with an O(log n) lookup and insertion speed. This is both faster |

119 | -- and slower than an actual array /(O(1) lookup and O(n) insertion |

120 | -- respectively)/. |

121 | -- |

122 | -- The user can imagine a two dimensional grid that can be modified |

123 | -- or queried via co-ordinate pair indices. |

124 | |

125 | -- |Lens for accessing and manipulating data at a specific |

126 | -- location. |

127 | -- |

128 | -- This is simply 'getLocation' and 'setLocation' wrapped into a lens. |

129 | atLocation :: Eq a => Location -> Lens' (QuadTree a) a |

130 | atLocation index = lens (getLocation index) (setLocation index) |

131 | |

132 | -- |Getter for the value at a given location for a 'QuadTree'. |

133 | getLocation :: Location -> QuadTree a -> a |

134 | getLocation index tree |

135 | | outOfBounds tree index = |

136 | error "Location index out of QuadTree bounds." |

137 | | otherwise = |

138 | go (offsetIndex tree index) (treeDepth tree) (wrappedTree tree) |

139 | where |

140 | go :: Location -> Int -> Quadrant a -> a |

141 | go _ _ (Leaf x) = x |

142 | go _ 0 _ = error "Wrapped tree is deeper than tree depth." |

143 | go (x,y) n (Node a b c d) = |

144 | go (x `mod` mid, y `mod` mid) (n - 1) node |

145 | where mid = 2 ^ (n - 1) |

146 | node | y < mid = if x < mid then a |

147 | else b |

148 | | otherwise = if x < mid then c |

149 | else d |

150 | |

151 | -- |Setter for the value at a given location for a 'QuadTree'. |

152 | -- |

153 | -- This automatically compresses the 'QuadTree' nodes if possible with |

154 | -- the new value. |

155 | setLocation :: forall a. Eq a => Location -> QuadTree a -> a -> QuadTree a |

156 | setLocation index tree new |

157 | | outOfBounds tree index = |

158 | error "Location index out of QuadTree bounds." |

159 | | otherwise = |

160 | onQuads (go (offsetIndex tree index) (treeDepth tree)) tree |

161 | where |

162 | go :: Eq a => Location -> Int -> Quadrant a -> Quadrant a |

163 | go (x,y) n (Leaf old) |

164 | | old == new = Leaf old |

165 | | n == 0 = Leaf new |

166 | | otherwise = go (x,y) n (Node l l l l) |

167 | where l = Leaf old |

168 | go _ 0 _ = error "Wrapped tree is deeper than tree depth." |

169 | go (x,y) n (Node a b c d) = fusedNode |

170 | where fusedNode = fuse newNode |

171 | newNode |

172 | | y < mid = if x < mid then Node (recurse a) b c d |

173 | else Node a (recurse b) c d |

174 | | otherwise = if x < mid then Node a b (recurse c) d |

175 | else Node a b c (recurse d) |

176 | recurse = go (x `mod` mid, y `mod` mid) (n - 1) |

177 | mid = 2 ^ (n - 1) |

178 | |

179 | ---- Helpers: |

180 | |

181 | -- |Checks if a 'Location' is outside the boundaries of a 'QuadTree'. |

182 | |

183 | outOfBounds :: QuadTree a -> Location -> Bool |

184 | outOfBounds tree (x,y) = x < 0 || y < 0 |

185 | || x >= treeLength tree |

186 | || y >= treeWidth tree |

187 | |

188 | -- |Dimensions of a 'QuadTree', as an Int pair. |

189 | |

190 | treeDimensions :: QuadTree a |

191 | -> (Int, Int) -- ^ (Length, Width) |

192 | treeDimensions tree = (treeLength tree, treeWidth tree) |

193 | |

194 | offsetIndex :: QuadTree a -> Location -> Location |

195 | offsetIndex tree (x,y) = (x + xOffset, y + yOffset) |

196 | where xOffset = (dimension - treeLength tree) `div` 2 |

197 | yOffset = (dimension - treeWidth tree) `div` 2 |

198 | dimension = 2 ^ treeDepth tree |

199 | |

200 | fuse :: Eq a => Quadrant a -> Quadrant a |

201 | fuse (Node (Leaf a) (Leaf b) (Leaf c) (Leaf d)) |

202 | | a == b && b == c && c == d = Leaf a |

203 | fuse oldNode = oldNode |

204 | |

205 | ---- Functor: |

206 | |

207 | onQuads :: (Quadrant a -> Quadrant b) -> QuadTree a -> QuadTree b |

208 | onQuads fn tree = tree {wrappedTree = fn (wrappedTree tree)} |

209 | |

210 | -- |Cleanup function for use after any 'Control.Monad.fmap'. |

211 | -- |

212 | -- When elements of a 'QuadTree' are modified by 'setLocation' (or |

213 | -- the 'atLocation' lens), it automatically compresses identical |

214 | -- adjacent nodes into larger ones. This keeps the 'QuadTree' from |

215 | -- bloating over constant use. |

216 | -- |

217 | -- 'Control.Monad.fmap' does not do this. If you wish to treat the |

218 | -- 'QuadTree' as a 'Control.Monad.Functor', you should compose this |

219 | -- function after to collapse it down to its minimum size. |

220 | -- |

221 | -- Example: |

222 | -- @ |

223 | -- 'fuseTree' $ 'Control.Monad.fmap' fn tree |

224 | -- @ |

225 | -- This particular example is reified in the function below. |

226 | |

227 | fuseTree :: Eq a => QuadTree a -> QuadTree a |

228 | fuseTree = onQuads fuseQuads |

229 | |

230 | fuseQuads :: Eq a => Quadrant a -> Quadrant a |

231 | fuseQuads (Node a b c d) = fuse $ Node (fuseQuads a) |

232 | (fuseQuads b) |

233 | (fuseQuads c) |

234 | (fuseQuads d) |

235 | fuseQuads leaf = leaf |

236 | |

237 | -- |tmap is simply 'Control.Monad.fmap' with 'fuseTree' applied after. |

238 | -- |

239 | -- prop> tmap fn tree == fuseTree $ fmap fn tree |

240 | tmap :: Eq b => (a -> b) -> QuadTree a -> QuadTree b |

241 | tmap = fuseTree .: fmap |

242 | |

243 | ---- Foldable: |

244 | |

245 | -- $foldable |

246 | -- 'QuadTree's can be folded just like lists. If you simply replace |

247 | -- the "Prelude" fold functions with "Data.Foldable" ones... |

248 | -- |

249 | -- @ |

250 | -- import "Data.Foldable" |

251 | -- import "Prelude" hiding (foldr, foldl, any, sum, find...) |

252 | -- @ |

253 | -- |

254 | -- ... Then you can directly call then on 'QuadTree's without |

255 | -- qualification. No list functionality will be lost since the |

256 | -- "Data.Foldable" functions also work exactly like the "Prelude" |

257 | -- folds for list processing. |

258 | -- |

259 | -- In addition you also get some extras like 'Data.Foldable.toList'. |

260 | |

261 | -- $tiles |

262 | -- Directly folding a 'QuadTree' will expand it into a sequence of |

263 | -- elements that are then folded over. For some types of operations |

264 | -- this can be incredibly inefficient; it may be faster to simply |

265 | -- manipulate a sequence of leaves and then later decompose the |

266 | -- results into a list of elements. |

267 | -- |

268 | -- For these operations, we can use 'Tile's. 'Tile's are simply |

269 | -- blocks of elements, represented by a tuple of the leaf data and |

270 | -- some information on the spatial location and dimensions of the |

271 | -- block. |

272 | |

273 | -- |Rectangular area, represented by a tuple of four Ints. |

274 | -- |

275 | -- They correspond to (X floor, Y floor, X ceiling, Y ceiling). |

276 | -- |

277 | -- The co-ordinates are inclusive of all the rows and columns in all |

278 | -- four Ints. |

279 | -- |

280 | -- prop> regionArea (x, y, x, y) == 1 |

281 | |

282 | type Region = (Int, Int, Int, Int) |

283 | |

284 | -- |Each 'Tile' is a tuple of an element from a 'QuadTree' and the |

285 | -- 'Region' it subtends. |

286 | |

287 | type Tile a = (a, Region) |

288 | |

289 | foldTree :: (a -> b -> b) -> b -> QuadTree a -> b |

290 | foldTree fn z = foldr fn z . expand . tile |

291 | |

292 | -- |Takes a list of 'Tile's and then decomposes them into a list of |

293 | -- all their elements, properly weighted by 'Tile' size. |

294 | |

295 | expand :: [Tile a] -> [a] |

296 | expand = concatMap decompose |

297 | where decompose :: Tile a -> [a] |

298 | decompose (a, r) = replicate (regionArea r) a |

299 | |

300 | -- |Returns a list of 'Tile's. The block equivalent of |

301 | -- 'Data.Foldable.toList'. |

302 | |

303 | tile :: QuadTree a -> [Tile a] |

304 | tile = foldTiles (:) [] |

305 | |

306 | -- |Decomposes a 'QuadTree' into its constituent 'Tile's, before |

307 | -- folding a 'Tile' consuming function over all of them. |

308 | |

309 | foldTiles :: forall a b. (Tile a -> b -> b) -> b -> QuadTree a -> b |

310 | foldTiles fn z tree = go (treeRegion tree) (wrappedTree tree) z |

311 | where go :: Region -> Quadrant a -> b -> b |

312 | go r (Leaf a) = fn (a, intersection) |

313 | where intersection = regionIntersection (boundaries tree) r |

314 | go (xl, yt, xr, yb) (Node a b c d) = |

315 | go (xl, yt, midx, midy) a . |

316 | go (midx + 1, yt, xr, midy) b . |

317 | go (xl, midy + 1, midx, yb) c . |

318 | go (midx + 1, midy + 1, xr, yb) d |

319 | where midx = (xr + xl) `div` 2 |

320 | midy = (yt + yb) `div` 2 |

321 | |

322 | treeRegion :: QuadTree a -> Region |

323 | treeRegion tree = (0, 0, limit, limit) |

324 | where limit = (2 ^ treeDepth tree) - 1 |

325 | |

326 | boundaries :: QuadTree a -> Region |

327 | boundaries tree = (left, top, right, bottom) |

328 | where (left, top) = offsetIndex tree (0,0) |

329 | (right, bottom) = offsetIndex tree (treeLength tree - 1, |

330 | treeWidth tree - 1) |

331 | |

332 | regionIntersection :: Region -> Region -> Region |

333 | regionIntersection (xl , yt , xr , yb ) |

334 | (xl', yt', xr', yb') = |

335 | (max xl xl', max yt yt', |

336 | min xr xr', min yb yb') |

337 | |

338 | -- |Simple helper function that lets you calculate the area of a |

339 | -- 'Region', usually for 'Data.List.replicate' purposes. |

340 | |

341 | regionArea :: Region -> Int |

342 | regionArea (xl,yt,xr,yb) = (xr + 1 - xl) * (yb + 1 - yt) |

343 | |

344 | -- |Does the region contain this location? |

345 | |

346 | inRegion :: Location -> Region -> Bool |

347 | inRegion (x,y) (xl,yt,xr,yb) = xl <= x && x <= xr && |

348 | yt <= y && y <= yb |

349 | |

350 | ---- Foldable extras: |

351 | |

352 | -- |'Data.List.filter's a list of the 'QuadTree' 's elements. |

353 | |

354 | filterTree :: (a -> Bool) -> QuadTree a -> [a] |

355 | filterTree fn = expand . filterTiles fn . tile |

356 | |

357 | -- |'Data.List.sortBy's a list of the 'QuadTree' 's elements. |

358 | |

359 | sortTreeBy :: (a -> a -> Ordering) -> QuadTree a -> [a] |

360 | sortTreeBy fn = expand . sortTilesBy fn . tile |

361 | |

362 | -- |'Data.List.filter's a list of the 'Tile's of a 'QuadTree'. |

363 | |

364 | filterTiles :: (a -> Bool) -> [Tile a] -> [Tile a] |

365 | filterTiles _ [] = [] |

366 | filterTiles fn ((a,r) : rs) |

367 | | fn a = (a,r) : filterTiles fn rs |

368 | | otherwise = filterTiles fn rs |

369 | |

370 | -- |'Data.List.sortBy's a list of the 'Tile's of a 'QuadTree'. |

371 | |

372 | sortTilesBy :: (a -> a -> Ordering) -> [Tile a] -> [Tile a] |

373 | sortTilesBy fn = sortBy (fn `on` fst) |

374 | |

375 | ---- Constructor: |

376 | |

377 | -- |Constructor that generates a 'QuadTree' of the given dimensions, |

378 | -- with all cells filled with a default value. |

379 | |

380 | makeTree :: (Int, Int) -- ^ (Length, Width) |

381 | -> a -- ^ Initial element to fill |

382 | -> QuadTree a |

383 | makeTree (x,y) a |

384 | | x <= 0 || y <= 0 = error "Invalid dimensions for tree." |

385 | | otherwise = Wrapper { wrappedTree = Leaf a |

386 | , treeLength = x |

387 | , treeWidth = y |

388 | , treeDepth = fst . fromJust $ |

389 | find ((>= max x y) . snd) $ |

390 | zip [0..] (iterate (*2) 1) } |

391 | |

392 | |

393 | ---- Sample Printers: |

394 | |

395 | -- |Generates a newline delimited string representing a 'QuadTree' as |

396 | -- a 2D block of characters. |

397 | -- |

398 | -- Note that despite the word 'show' in the function name, this does |

399 | -- not 'Text.show' the 'QuadTree'. It pretty prints it. The name |

400 | -- is simply a mnemonic for its @'QuadTree' -> String@ behaviour. |

401 | |

402 | showTree :: (a -> Char) -- ^ Function to generate characters for each |

403 | -- 'QuadTree' element. |

404 | -> QuadTree a -> String |

405 | showTree printer tree = breakString (treeLength tree) string |

406 | where string = map printer grid |

407 | grid = [getLocation (x,y) tree | |

408 | y <- [0 .. treeWidth tree - 1], |

409 | x <- [0 .. treeLength tree - 1]] |

410 | breakString :: Int -> String -> String |

411 | breakString _ [] = [] |

412 | breakString n xs = a ++ "\n" ++ breakString n b |

413 | where (a,b) = splitAt n xs |

414 | |

415 | -- |As 'showTree' above, but also prints it. |

416 | |

417 | printTree :: (a -> Char) -- ^ Function to generate characters for each |

418 | -- 'QuadTree' element. |

419 | -> QuadTree a -> IO () |

420 | printTree = putStr .: showTree |

421 | |

422 | |

423 | --------- Test: |

424 | |

425 | x' :: QuadTree Int |

426 | x' = Wrapper { treeLength = 6 |

427 | , treeWidth = 5 |

428 | , treeDepth = 3 |

429 | , wrappedTree = y' } |

430 | |

431 | y' :: Quadrant Int |

432 | y' = Node (Leaf 0) |

433 | (Node (Leaf 2) |

434 | (Leaf 3) |

435 | (Leaf 4) |

436 | (Leaf 5)) |

437 | (Leaf 1) |

438 | (Leaf 9) |

439 | |

440 | basic :: QuadTree Int |

441 | basic = Wrapper {treeLength = 4, treeWidth = 5, treeDepth = 3, |

442 | wrappedTree = Node (Leaf 0) |

443 | (Leaf 1) |

444 | (Leaf 2) |

445 | (Leaf 3)} |

446 | |

447 | x5 = set (atLocation (2,3)) 1 (makeTree (5,7) 0) |

448 | x6 = set (atLocation (2,3)) 1 (makeTree (6,7) 0) |

449 | p n = printTree (head . show) n |

450 | |

451 | test = set (atLocation (0,0)) 'd' $ |

452 | set (atLocation (5,5)) 'c' $ |

453 | set (atLocation (3,2)) 'b' $ |

454 | set (atLocation (2,4)) 'a' $ |

455 | makeTree (6,6) '.' |