A quad tree.

gistfile1.hs

module QuadTree 
  ( Point, Size, Rect(Rect), Item(Item), QuadTree, Quad
  , empty, insert, delete, filterTree, query, stats
  ) where

import qualified Data.List as L
import qualified Debug.Trace as Dbg

-------------------------------------------------------------------------------
--  TYPES
-------------------------------------------------------------------------------

type Point = (Double, Double)
type Size = (Double, Double)

data Rect = Rect Point Size deriving (Show, Eq)

data QuadTree a = QuadTree 
  { minCellSize :: Size
  , maxCellFill :: Int
  , root :: Quad a 
  } deriving (Show, Eq)


data Item a = Item Point a deriving (Show, Eq)

data Quad a = Leaf { bounds :: Rect 
           , items :: [Item a] }
      | Node { bounds :: Rect
           , topLeft :: (Quad a)
           , topRight :: (Quad a)
           , bottomLeft :: (Quad a)
           , bottomRight :: (Quad a)
           } deriving (Show, Eq)

data Stats = Stats { numLeaves :: Int, numInner :: Int } deriving (Show)

-------------------------------------------------------------------------------
--  PRIVATE RECTANGLE FUNCTIONS
-------------------------------------------------------------------------------

-- | Subdivides a rectangle into four equally sized smaller rectangles
subdivideRect :: Rect 
        -> ( Rect -- top left rect
         , Rect -- top right rect
         , Rect -- bottom left rect
         , Rect)-- bottom right rect

subdivideRect rect@(Rect (x,y) (w,h)) = ( Rect (x,y)         (w/2,h/2)
                    , Rect (x+w/2,y)     (w/2,h/2)
                    , Rect (x,y+w/2)     (w/2,h/2)
                    , Rect (x+w/2,y+w/2) (w/2,h/2))



-- | Returns whether the two rectangles are intersecting
intersects :: Rect -> Rect -> Bool

a `intersects` b = not (al > br || ar < bl || at > bb || ab < bt)  where 
  (al, at, ar, ab) = ltrb a
  (bl, bt, br, bb) = ltrb b
  -- | converts a rect defined by position and size to a rect defined by 
  -- | two corners
  ltrb :: Rect -> (Double, Double, Double, Double)
  ltrb (Rect (x,y) (w,h)) = (x,y,x+w,y+h)


-- | checks whether the point is contained within the rectangle
containedIn :: Point -> Rect -> Bool
containedIn p@(x, y) r@(Rect (rx, ry) (w,h)) = 
  x >= rx && y >= ry && x < rx+w && y < ry+h


quadIndex :: Point -> Rect -> (Int,Int)
quadIndex (x,y) (Rect (rx, ry) (w,h)) = 
  (floor $ (x - rx) / w * 2, floor $ (y - ry) / h * 2)

-------------------------------------------------------------------------------
--  PRIVATE QUADTREE FUNCTIONS
-------------------------------------------------------------------------------

-- | Collapses a node with four empty leaves into one one empty leaf

collapseEmpty :: Quad a -> Quad a
collapseEmpty (Node rect (Leaf _ []) (Leaf _ []) 
             (Leaf _ []) (Leaf _ [])) = (Leaf rect [])
collapseEmpty other = other


-- | Returns true when the quad and all sub-quads are empty

isEmpty :: Quad a -> Bool
isEmpty quad = case quad of 
  Leaf _ items -> null items
  Node _ tl tr bl br -> isEmpty tl && isEmpty tr && isEmpty bl && isEmpty br


addStats :: Stats -> Stats -> Stats
addStats (Stats a b) (Stats c d) = Stats (a+b) (c+d)

-------------------------------------------------------------------------------
--  PUBLIC QUADTREE FUNCTIONS
-------------------------------------------------------------------------------


-- | creates a QuadTree with the desired properties
empty :: Rect       -- ^ the outer bounds of the quad tree
    -> Size     -- ^ the minimum size of one quad
    -> Int        -- ^ the maximum number of elements in one quad
    -> QuadTree a -- ^ an empty QuadTree with the desired properties

empty rect minSize maxFill = QuadTree { minCellSize = minSize
                    , maxCellFill = maxFill
                    , root = Leaf rect [] }



-- | Queries all elements in the given rect from the quad tree
query :: Rect -> QuadTree a -> [Item a]
query rect tree = queryQuad rect (root tree) where 
  queryQuad :: Rect -> Quad a -> [Item a]

  queryQuad searchRect quad = case quad of
    leaf@(Leaf rect items) -> 
      filter (\(Item p _) -> p `containedIn` searchRect) items
    node@(Node rect tl tr bl br) -> 
        (recurseOnIntersection searchRect tl)
      ++  (recurseOnIntersection searchRect tr)
      ++  (recurseOnIntersection searchRect bl)
      ++  (recurseOnIntersection searchRect br)

  recurseOnIntersection :: Rect -> Quad a -> [Item a]
  recurseOnIntersection searchRect quad =
    if searchRect `intersects` (bounds quad) 
    then queryQuad searchRect quad 
    else []



-- | inserts an element into the quad tree at the given position
insert :: Item a        -- ^ the element to insert
     -> QuadTree a    -- ^ the QuadTree that is target of the insertion
     -> QuadTree a    -- ^ the resulting QuadTree

insert item tree@(QuadTree (minWidth,minHeight) _ r) =
  tree { root= (insertQuad item (expandRoot r)) } 
  where

  expandRoot :: Quad a -> Quad a
  expandRoot root = let (Item p@(x,y) _) = item 
                        rect@(Rect (rx,ry) (w,h)) = bounds root in 
    if (p `containedIn` (bounds root)) then root
    else let (dirX, dirY) = (signum $ x - rx - w/2.0, signum $ y - ry - h / 2.0) 
             (dx, dy) = ((dirX - 1) / 2, (dirY - 1) / 2)
             rootTemplate = Node (Rect (rx + dx*w,ry + dy*h) (2*w,2*h))
                                 (Leaf (Rect (rx + dx*w    ,ry + dy*h    ) (w,h)) [])
                                 (Leaf (Rect (rx + (dx+1)*w,ry + dy*h    ) (w,h)) [])
                                 (Leaf (Rect (rx + dx*w    ,ry + (dy+1)*h) (w,h)) [])
                                 (Leaf (Rect (rx + (dx+1)*w,ry + (dy+1)*h) (w,h)) [])
             newroot = case (dirX, dirY) of
                (-1.0, -1.0) -> rootTemplate { bottomRight = root }
                ( 1.0, -1.0) -> rootTemplate { bottomLeft = root }
                (-1.0,  1.0) -> rootTemplate { topRight = root }
                ( 1.0,  1.0) -> rootTemplate { topLeft = root }
         in expandRoot newroot
    

  insertQuad :: Item a -> Quad a -> Quad a

  insertQuad item @(Item p@(x,y) _) quad = case quad of
    leaf@(Leaf rect@(Rect _ (w,h)) items) -> 
      if length items < maxCellFill tree 
         || w <= minWidth || h <= minHeight
      then Leaf rect (item:items)
      else insertQuad item (subdivide leaf)

    node@(Node rect tl tr bl br) -> 
      case quadIndex (x,y) rect of
        (0,0) -> node { topLeft=(insertQuad item tl) }
        (1,0) -> node { topRight=(insertQuad item tr) }
        (0,1) -> node { bottomLeft=(insertQuad item bl) }
        (1,1) -> node { bottomRight=(insertQuad item br) }

  subdivide :: Quad a -> Quad a
  subdivide (Leaf rect items) =
    let (rtl, rtr, rbl, rbr) = subdivideRect rect 
        (itl, itr, ibl, ibr) = 
          ( filter (\(Item p _) -> p `containedIn` rtl) items
          , filter (\(Item p _) -> p `containedIn` rtr) items
          , filter (\(Item p _) -> p `containedIn` rbl) items
          , filter (\(Item p _) -> p `containedIn` rbr) items)
    in Node rect (Leaf rtl itl) (Leaf rtr itr)
           (Leaf rbl ibl) (Leaf rbr ibr)

  subdivide (Node {}) = error "Only leafs can be futher subdivided"



-- | Deletes an item from the quad tree and merges empty sub-quads.

delete :: (Eq a) 
     => Item a     -- ^ the element to insert
     -> QuadTree a -- ^ the QuadTree containing the item
     -> QuadTree a -- ^ the new QuadTree without the item

delete item tree@(QuadTree (minWidth,minHeight) _ r) = 
  tree { root=collapseEmpty $ (removeQuad item r) } 
  where
  removeQuad :: (Eq a) => Item a -> Quad a -> Quad a
  removeQuad item@(Item (x,y) _) quad = case quad of 
    Leaf rect items -> Leaf rect (L.delete item items)
    node@(Node rect tl tr bl br) ->
      case quadIndex (x,y) rect of
        (0,0) -> collapseEmpty $ node { topLeft=(removeQuad item tl) }
        (1,0) -> collapseEmpty $ node { topRight=(removeQuad item tr) }
        (0,1) -> collapseEmpty $ node { bottomLeft=(removeQuad item bl) }
        (1,1) -> collapseEmpty $ node { bottomRight=(removeQuad item br) }


-- | Returns a new QuadTree which contains all items satifying the predicate

filterTree :: (Eq a) 
     => (Item a -> Bool) -- ^ the predicate for choosing the items
     -> QuadTree a       -- ^ the original QuadTree
     -> QuadTree a       -- ^ the filtered QuadTree

filterTree predicate tree@(QuadTree (minWidth,minHeight) _ r) = 
  tree { root=collapseEmpty $ (filterQuad predicate r) } 
  where
  filterQuad :: (Eq a) => (Item a -> Bool) -> Quad a -> Quad a
  filterQuad predicate quad = case quad of 
    Leaf rect items -> Leaf rect (filter predicate items)
    node@(Node rect tl tr bl br) -> collapseEmpty $
      Node rect (filterQuad predicate tl) (filterQuad predicate tr)
            (filterQuad predicate bl) (filterQuad predicate br)


stats :: QuadTree a -> Stats
stats tree@(QuadTree {root=r}) = quadStats r (Stats 0 0) where
  quadStats :: Quad a -> Stats -> Stats
  quadStats q stats = case q of
    Leaf _ _ -> stats { numLeaves = 1 + (numLeaves stats) }
    Node _ tl tr bl br -> s4 { numInner = 1 + (numInner s4) } where
      s1 = quadStats tl stats
      s2 = quadStats tr s1
      s3 = quadStats bl s2
      s4 = quadStats br s3


-------------------------------------------------------------------------------
--  TYPECLASS INSTANCES
-------------------------------------------------------------------------------


-- | A quad tree is a functor
instance Functor QuadTree where  
  fmap f tree = tree { root = (fmap f (root tree)) }

-- | A quad is a functor
instance Functor Quad where  
  fmap f (Leaf rect items) = Leaf rect (map (fmap f) items)

  fmap f (Node rect tl tr bl br) = Node rect (fmap f tl) (fmap f tr) 
                         (fmap f bl) (fmap f br)

-- | An item is a functor
instance Functor Item where  
  fmap f (Item p a) = Item p (f a)