brillo-algorithms-2.0.0: Brillo/Data/QuadTree.hs

{-# LANGUAGE RankNTypes #-}

{-| A QuadTree can be used to recursively divide up 2D space into quadrants.
  The smallest division corresponds to an unit `Extent`, so the total depth
  of the tree will depend on what sized `Extent` you start with.
-}
module Brillo.Data.QuadTree (
  QuadTree (..),
  emptyNode,
  takeQuadOfTree,
  liftToQuad,
  insertByPath,
  insertByCoord,
  lookupNodeByPath,
  lookupByPath,
  lookupByCoord,
  flattenQuadTree,
  flattenQuadTreeWithExtents,
)
where

import Brillo.Data.Extent
import Brillo.Data.Quad


-- | The quad tree structure.
data QuadTree a
  = -- | An empty node.
    TNil
  | -- | A leaf containint some value.
    TLeaf a
  | -- | A node with four children.
    TNode
      (QuadTree a)
      (QuadTree a) -- NW NE
      (QuadTree a)
      (QuadTree a) -- SW SE
  deriving (Show)


-- | A node with `TNil`. for all its branches.
emptyNode :: QuadTree a
emptyNode = TNode TNil TNil TNil TNil


{-| Get a quadrant from a node.
  If the tree does not have an outer node then `Nothing`.
-}
takeQuadOfTree ::
  Quad ->
  QuadTree a ->
  Maybe (QuadTree a)
takeQuadOfTree quad tree =
  case tree of
    TNil -> Nothing
    TLeaf{} -> Nothing
    TNode nw ne sw se ->
      case quad of
        NW -> Just nw
        NE -> Just ne
        SW -> Just sw
        SE -> Just se


{-| Apply a function to a quadrant of a node.
  If the tree does not have an outer node then return the original tree.
-}
liftToQuad ::
  Quad ->
  (QuadTree a -> QuadTree a) ->
  QuadTree a ->
  QuadTree a
liftToQuad quad f tree =
  case tree of
    TNil -> tree
    TLeaf{} -> tree
    TNode nw ne sw se ->
      case quad of
        NW -> TNode (f nw) ne sw se
        NE -> TNode nw (f ne) sw se
        SW -> TNode nw ne (f sw) se
        SE -> TNode nw ne sw (f se)


{-| Insert a value into the tree at the position given by a path.
  If the path intersects an existing `TLeaf` then return the original tree.
-}
insertByPath :: [Quad] -> a -> QuadTree a -> QuadTree a
insertByPath [] x _ =
  TLeaf x
insertByPath (q : qs) x tree =
  case tree of
    TNil -> liftToQuad q (insertByPath qs x) emptyNode
    TLeaf{} -> tree
    TNode{} -> liftToQuad q (insertByPath qs x) tree


{-| Insert a value into the node containing this coordinate.
  The node is created at maximum depth, corresponding to an unit `Extent`.
-}
insertByCoord :: Extent -> Coord -> a -> QuadTree a -> Maybe (QuadTree a)
insertByCoord extent coord x tree =
  do
    path <- pathToCoord extent coord
    return $ insertByPath path x tree


-- | Lookup a node based on a path to it.
lookupNodeByPath ::
  [Quad] ->
  QuadTree a ->
  Maybe (QuadTree a)
lookupNodeByPath [] tree =
  Just tree
lookupNodeByPath (q : qs) tree =
  case tree of
    TNil -> Nothing
    TLeaf{} -> Nothing
    TNode{} ->
      case takeQuadOfTree q tree of
        Nothing -> Nothing
        Just quad -> lookupNodeByPath qs quad


-- | Lookup an element based given a path to it.
lookupByPath :: [Quad] -> QuadTree a -> Maybe a
lookupByPath path tree =
  case lookupNodeByPath path tree of
    Just (TLeaf x) -> Just x
    _ -> Nothing


-- | Lookup a node if a tree given a coordinate which it contains.
lookupByCoord ::
  forall a.
  -- | Extent that covers the whole tree.
  Extent ->
  -- | Coordinate of the value of interest.
  Coord ->
  QuadTree a ->
  Maybe a
lookupByCoord extent coord tree =
  do
    path <- pathToCoord extent coord
    lookupByPath path tree


-- | Flatten a QuadTree into a list of its contained values, with coordinates.
flattenQuadTree ::
  forall a.
  -- | Extent that covers the whole tree.
  Extent ->
  QuadTree a ->
  [(Coord, a)]
flattenQuadTree extentInit treeInit =
  flatten' extentInit treeInit
  where
    flatten' extent tree =
      case tree of
        TNil -> []
        TLeaf x ->
          let (_, s, _, w) = takeExtent extent
          in  [((w, s), x)]
        TNode{} -> concat $ map (flattenQuad extent tree) allQuads

    flattenQuad extent tree quad = do
      let extent' = cutQuadOfExtent quad extent
      case takeQuadOfTree quad tree of
        Nothing -> []
        Just tree' -> flatten' extent' tree'


-- | Flatten a QuadTree into a list of its contained values, with extents.
flattenQuadTreeWithExtents ::
  forall a.
  -- | Extent that covers the whole tree.
  Extent ->
  QuadTree a ->
  [(Extent, a)]
flattenQuadTreeWithExtents extentInit treeInit =
  flatten' extentInit treeInit
  where
    flatten' extent tree =
      case tree of
        TNil -> []
        TLeaf x ->
          [(extent, x)]
        TNode{} -> concat $ map (flattenQuad extent tree) allQuads

    flattenQuad extent tree quad = do
      let extent' = cutQuadOfExtent quad extent
      case takeQuadOfTree quad tree of
        Nothing -> []
        Just tree' -> flatten' extent' tree'