Ticket #291: QuadTree.hs

File QuadTree.hs, 14.2 kB (added by Kron, 4 months ago)
Line 
1
2{-# OPTIONS_HADDOCK show-extensions #-}
3
4{-# LANGUAGE Safe #-}
5
6{-# LANGUAGE Rank2Types #-}
7{-# LANGUAGE ScopedTypeVariables #-}
8
9{-|
10Module      : QuadTree
11Description : Region quadtrees with Lens support.
12Copyright   : (c) Ashley Moni, 2014
13License     : GPL-3
14Maintainer  : Ashley Moni <y30099923@gmail.com>
15Stability   : Experimental
16Portability : Rank 2 Types, Scoped Type Variables
17
18The purpose of this module is to provide discrete region quadtrees
19that can be used as simple functional alternatives to 2D arrays,
20with lens support.
21
22@
23test = 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
31d.....
32......
33...b..
34......
35..a...
36.....c
37-}
38
39module 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
61import Control.Lens.Type (Lens')
62import Control.Lens.Lens (lens)
63import Control.Lens.Setter (set)
64
65import Data.List (find, sortBy)
66import Data.Maybe (fromJust)
67import Data.Function (on)
68import Data.Composition ((.:))
69
70-- Foldable:
71import Data.Foldable (Foldable, foldr)
72import Prelude hiding (foldr)
73
74---- Structures:
75
76-- |Tuple corresponds to (X, Y) co-ordinates.
77
78type 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
86data QuadTree a = Wrapper { wrappedTree :: Quadrant a
87                          , treeLength :: Int
88                          , treeWidth  :: Int
89                          , treeDepth :: Int }
90  deriving (Show, Read)
91
92instance Functor QuadTree where
93  fmap fn = onQuads $ fmap fn
94
95instance Foldable QuadTree where
96  foldr = foldTree
97
98--
99
100data Quadrant a = Leaf a
101                | Node (Quadrant a)
102                       (Quadrant a)
103                       (Quadrant a)
104                       (Quadrant a)
105  deriving (Show, Read)
106
107instance 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.
129atLocation :: Eq a => Location -> Lens' (QuadTree a) a
130atLocation index = lens (getLocation index) (setLocation index)
131
132-- |Getter for the value at a given location for a 'QuadTree'.
133getLocation :: Location -> QuadTree a -> a
134getLocation 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.
155setLocation :: forall a. Eq a => Location -> QuadTree a -> a -> QuadTree a
156setLocation 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
183outOfBounds :: QuadTree a -> Location -> Bool
184outOfBounds 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
190treeDimensions :: QuadTree a
191               -> (Int, Int) -- ^ (Length, Width)
192treeDimensions tree = (treeLength tree, treeWidth tree)
193
194offsetIndex :: QuadTree a -> Location -> Location
195offsetIndex 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
200fuse :: Eq a => Quadrant a -> Quadrant a
201fuse (Node (Leaf a) (Leaf b) (Leaf c) (Leaf d))
202  | a == b && b == c && c == d = Leaf a
203fuse oldNode                   = oldNode
204
205---- Functor:
206
207onQuads :: (Quadrant a -> Quadrant b) -> QuadTree a -> QuadTree b
208onQuads 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
227fuseTree :: Eq a => QuadTree a -> QuadTree a
228fuseTree = onQuads fuseQuads
229
230fuseQuads :: Eq a => Quadrant a -> Quadrant a
231fuseQuads (Node a b c d) = fuse $ Node (fuseQuads a)
232                                       (fuseQuads b)
233                                       (fuseQuads c)
234                                       (fuseQuads d)
235fuseQuads leaf           = leaf
236
237-- |tmap is simply 'Control.Monad.fmap' with 'fuseTree' applied after.
238--
239-- prop> tmap fn tree == fuseTree $ fmap fn tree
240tmap :: Eq b => (a -> b) -> QuadTree a -> QuadTree b
241tmap = 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
282type 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
287type Tile a = (a, Region)
288
289foldTree :: (a -> b -> b) -> b -> QuadTree a -> b
290foldTree 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
295expand :: [Tile a] -> [a]
296expand = 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
303tile :: QuadTree a -> [Tile a]
304tile = foldTiles (:) []
305
306-- |Decomposes a 'QuadTree' into its constituent 'Tile's, before
307-- folding a 'Tile' consuming function over all of them.
308
309foldTiles :: forall a b. (Tile a -> b -> b) -> b -> QuadTree a -> b
310foldTiles 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
322treeRegion :: QuadTree a -> Region
323treeRegion tree = (0, 0, limit, limit)
324  where limit = (2 ^ treeDepth tree) - 1
325
326boundaries :: QuadTree a -> Region
327boundaries 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
332regionIntersection :: Region -> Region -> Region
333regionIntersection (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
341regionArea :: Region -> Int
342regionArea (xl,yt,xr,yb) = (xr + 1 - xl) * (yb + 1 - yt)
343
344-- |Does the region contain this location?
345
346inRegion :: Location -> Region -> Bool
347inRegion (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
354filterTree :: (a -> Bool) -> QuadTree a -> [a]
355filterTree fn = expand . filterTiles fn . tile
356
357-- |'Data.List.sortBy's a list of the 'QuadTree' 's elements.
358
359sortTreeBy :: (a -> a -> Ordering) -> QuadTree a -> [a]
360sortTreeBy fn = expand . sortTilesBy fn . tile
361
362-- |'Data.List.filter's a list of the 'Tile's of a 'QuadTree'.
363
364filterTiles :: (a -> Bool) -> [Tile a] -> [Tile a]
365filterTiles _  [] = []
366filterTiles 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
372sortTilesBy :: (a -> a -> Ordering) -> [Tile a] -> [Tile a]
373sortTilesBy 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
380makeTree :: (Int, Int) -- ^ (Length, Width)
381                  -> a -- ^ Initial element to fill
382                  -> QuadTree a
383makeTree (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
402showTree :: (a -> Char) -- ^ Function to generate characters for each
403                        -- 'QuadTree' element.
404         -> QuadTree a -> String
405showTree 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
417printTree :: (a -> Char) -- ^ Function to generate characters for each
418                         -- 'QuadTree' element.
419          -> QuadTree a -> IO ()
420printTree = putStr .: showTree
421
422
423--------- Test:
424
425x' :: QuadTree Int
426x' = Wrapper { treeLength = 6
427            , treeWidth = 5
428            , treeDepth = 3
429            , wrappedTree = y' }
430
431y' :: Quadrant Int
432y' = Node (Leaf 0)
433          (Node (Leaf 2)
434                (Leaf 3)
435                (Leaf 4)
436                (Leaf 5))
437          (Leaf 1)
438          (Leaf 9)
439
440basic :: QuadTree Int
441basic = Wrapper {treeLength = 4, treeWidth = 5, treeDepth = 3,
442                 wrappedTree = Node (Leaf 0)
443                                    (Leaf 1)
444                                    (Leaf 2)
445                                    (Leaf 3)}
446
447x5 = set (atLocation (2,3)) 1 (makeTree (5,7) 0)
448x6 = set (atLocation (2,3)) 1 (makeTree (6,7) 0)
449p n = printTree (head . show) n
450
451test = 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) '.'