Quadtrees in Elm

QuadTree.elm

import Array (..)
import Text (asText)
import Graphics.Collage (..)
import Color (..)

drawBox box =
  move (box.horizontal.low + boxHalfWidth box, box.vertical.low + boxHalfHeight box) <|
     outlined (solid black) <|
        rect (boxWidth box) (boxHeight box)


drawQuadTree quadTree =
  case quadTree of
    Leaf box items -> [drawBox box] 
    Node box quadTreeNE quadTreeNW quadTreeSW quadTreeSE ->
      drawQuadTree quadTreeNE ++
      drawQuadTree quadTreeNW ++
      drawQuadTree quadTreeSW ++
      drawQuadTree quadTreeSE 

box lowX highX lowY highY = {
 horizontal = {
    low = lowX,
    high = highX },
  vertical = {
    low = lowY,
    high = highY }}

testBox = box -200 200 -200 200

item1 = box 10 20 10 20
item2 = box 30 40 0 50
item3 = box -100 -80 80 100
item4 = box 40 50 -40 0
item5 = box -50 -40 70 80
item6 = box -100 -80 130 190
item7 = box -20 -10 20 100

 
    
testQuadTree = 
 emptyQuadTree testBox |>
   insertQuadTree item1 |>
     insertQuadTree item2 |>
       insertQuadTree item3 |>
         insertQuadTree item4 |>
           insertQuadTree item5 |>
             insertQuadTree item6 |>
               insertQuadTree item7 

renderQuadTree quadTree = collage 400 400 (drawQuadTree quadTree)

main = renderQuadTree testQuadTree


---------

dropIf : (a -> Bool) -> Array a -> Array a
dropIf predicate = filter (not << predicate)

---------


type alias Interval = {
  low : Float,
  high : Float
}

type alias Box = {
  horizontal : Interval,
  vertical : Interval
}

type QuadTree a =
  Leaf Box (Array a) |
  Node Box (QuadTree a) (QuadTree a) (QuadTree a) (QuadTree a)

emptyQuadTree : Box -> QuadTree a
emptyQuadTree box = Leaf box empty

lengthQuadTree : QuadTree a -> Int
lengthQuadTree quadTree =
  case quadTree of
    Leaf box items -> length items
    Node box quadTreeNE quadTreeNW quadTreeSW quadTreeSE ->
      lengthQuadTree quadTreeNE +
      lengthQuadTree quadTreeNW +
      lengthQuadTree quadTreeSW +
      lengthQuadTree quadTreeSE

insertQuadTree : Box -> QuadTree Box -> QuadTree Box
insertQuadTree item quadTree =
  case quadTree of
    Leaf box items ->
      if intersectBoxes item box then
        let allItems = push item items
            insertNew quadrant =
              foldr (\item quadTree -> insertQuadTree item quadTree)
                    (emptyQuadTree quadrant)
                    allItems
            quadTreeNE = subdivideNE box
            quadTreeNW = subdivideNW box
            quadTreeSW = subdivideSW box
            quadTreeSE = subdivideSE box
        in
          if length items < 2 then Leaf box (push item items)
          else
            Node box (insertNew quadTreeNE)
                     (insertNew quadTreeNW)
                     (insertNew quadTreeSW)
                     (insertNew quadTreeSE)
      else
        quadTree

    Node box quadTreeNE quadTreeNW quadTreeSW quadTreeSE ->
      if intersectBoxes item box then
        Node box (insertQuadTree item quadTreeNE)
                 (insertQuadTree item quadTreeNW)
                 (insertQuadTree item quadTreeSW)
                 (insertQuadTree item quadTreeSE)
      else
        quadTree

removeFromQuadTree : a -> QuadTree a -> QuadTree a
removeFromQuadTree item quadTree =
  case quadTree of
    Leaf box items -> Leaf box (dropIf (\it -> it == item) items)
    Node box quadTreeNE quadTreeNW quadTreeSW quadTreeSE ->
      Node box (removeFromQuadTree item quadTreeNE)
               (removeFromQuadTree item quadTreeNW)
               (removeFromQuadTree item quadTreeSW)
               (removeFromQuadTree item quadTreeSE)


boxWidth : Box -> Float
boxWidth box =
  box.horizontal.high - box.horizontal.low

boxHeight : Box -> Float
boxHeight box =
  box.vertical.high - box.vertical.low

boxHalfWidth : Box -> Float
boxHalfWidth box =
  boxWidth box / 2

boxHalfHeight : Box -> Float
boxHalfHeight box =
  boxHeight box / 2


pointInInterval : Float -> Interval -> Bool
pointInInterval point interval =
  point < interval.high && point > interval.low

intersectIntervals : Interval -> Interval -> Bool
intersectIntervals interval1 interval2 =
  pointInInterval interval1.low interval2

intersectBoxes : Box -> Box -> Bool
intersectBoxes box1 box2 =
  intersectIntervals box1.horizontal box2.horizontal &&
  intersectIntervals box1.vertical box2.vertical

subdivideNE : Box -> Box
subdivideNE box =
  let vlow = box.vertical.low + boxHalfHeight box
      hlow = box.horizontal.high - boxHalfWidth box
  in {
    horizontal = {
      low = hlow,
      high = box.horizontal.high
    },
    vertical = {
      low = vlow,
      high = box.vertical.high
    }
  }

subdivideNW : Box -> Box
subdivideNW box =
  let vlow = box.vertical.low + boxHalfHeight box
      hhigh = box.horizontal.high - boxHalfWidth box
  in {
    horizontal = {
      low = box.horizontal.low,
      high = hhigh
    },
    vertical = {
      low = vlow,
      high = box.vertical.high
    }
  }

subdivideSW : Box -> Box
subdivideSW box =
  let vhigh = box.vertical.high - boxHalfHeight box
      hhigh = box.horizontal.high - boxHalfWidth box
  in {
    horizontal = {
      low = box.horizontal.low,
      high = hhigh
    },
    vertical = {
      low = box.vertical.low,
      high = vhigh
    }
  }

subdivideSE : Box -> Box
subdivideSE box =
  let vhigh = box.vertical.high - boxHalfHeight box
      hlow = box.horizontal.low + boxHalfWidth box
  in {
    horizontal = {
      low = hlow,
      high = box.horizontal.high
    },
    vertical = {
      low = box.vertical.low,
      high = vhigh
    }
  }