QuadTree.hs

{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Eta reduce" #-}
module QuadTree (
  QuadTree (..),
  treeZero,
  treeInsert,
  treeElems,
)
where

import Brillo.Data.Point (Point)


data QuadTree a
  = -- Nil cells take up space in the world, but don't contain any elements.
    --      They can be at any depth in the tree.
    QNil
      !Point -- cell center point
      !Float -- cell size
  | -- Leaf cells are the only ones that contain elements.
    --      They are always at the bottom of the tree.
    QLeaf
      !Point -- cell center point
      !Float -- cell size
      ![a] -- elements in this cell
  | -- Node cells contain more sub-trees
    QNode
      !Point -- cell center point
      !Float -- cell size
      !(QuadTree a)
      !(QuadTree a) -- NW NE
      !(QuadTree a)
      !(QuadTree a) -- SW SE
  deriving (Eq, Show)


-- Initial ---------------------------------------------------------------------
treeZero :: Float -> QuadTree a
treeZero size =
  QNil (0, 0) size


-- Quadrant --------------------------------------------------------------------

-- | Insert an element with a bounding box into the tree
treeInsert ::
  -- | maximum depth to place a leaf
  Int ->
  -- | current depth
  Int ->
  -- | bottom left of bounding box of new element
  Point ->
  -- | top right of bounding box of new element
  Point ->
  -- | element to insert into tree
  a ->
  -- | current tree
  QuadTree a ->
  QuadTree a
treeInsert depthMax depth p0@(x0, y0) p1@(x1, y1) a tree =
  case tree of
    QNode p@(x, y) size tNW tNE tSW tSE ->
      let
        tNW'
          | y1 > y && x0 < x = treeInsert depthMax (depth + 1) p0 p1 a tNW
          | otherwise = tNW

        tNE'
          | y1 > y && x1 > x = treeInsert depthMax (depth + 1) p0 p1 a tNE
          | otherwise = tNE

        tSW'
          | y0 < y && x0 < x = treeInsert depthMax (depth + 1) p0 p1 a tSW
          | otherwise = tSW

        tSE'
          | y0 < y && x1 > x = treeInsert depthMax (depth + 1) p0 p1 a tSE
          | otherwise = tSE
      in
        QNode p size tNW' tNE' tSW' tSE'
    QLeaf p size elems
      | depth >= depthMax ->
          QLeaf p size (a : elems)
    QNil p@(x, y) size
      | depth >= depthMax ->
          QLeaf p size [a]
      | otherwise ->
          treeInsert
            depthMax
            depth
            p0
            p1
            a
            ( let s2 = size / 2
              in  QNode
                    p
                    size
                    (QNil (x - s2, y + s2) s2)
                    (QNil (x + s2, y + s2) s2)
                    (QNil (x - s2, y - s2) s2)
                    (QNil (x + s2, y - s2) s2)
            )
    QLeaf (_, _) _ _ -> error "treeInsert: QLeaf"


-- flatten a quadtree into a list of its elements.
treeElems :: QuadTree a -> [[a]]
treeElems tree =
  case tree of
    QNode _ _ tNW tNE tSW tSE ->
      treeElems tNW ++ treeElems tNE ++ treeElems tSW ++ treeElems tSE
    QLeaf _ _ elems -> [elems]
    QNil{} -> []