sml-projects/fcore/quad-tree.sml

structure BinTree =
struct
  datatype 'a bintree =
    NODE of
      { x: int
      , y: int
      , ex: int
      , ey: int
      , data: 'a
      , left: 'a bintree
      , right: 'a bintree
      }
  | LEAF

  val empty = LEAF

  fun insert (newItem as {x, y, ex, ey, data}, tree) =
    case tree of
      LEAF =>
        NODE
          { x = x
          , y = y
          , ex = ex
          , ey = ey
          , data = data
          , left = LEAF
          , right = LEAF
          }
    | NODE {x = ox, y = oy, ex = oex, ey = oey, data = oldData, left, right} =>
        let
          val dir =
            if x < ox then
              LESS
            else if x > ox then
              GREATER
            else
              (if y < oy then
                 LESS
               else if y > oy then
                 GREATER
               else
                 (if ex < oex then
                    LESS
                  else if ex > oex then
                    GREATER
                  else
                    (if ey < oey then LESS
                     else if ey > oey then GREATER
                     else EQUAL)))
        in
          case dir of
            LESS =>
              NODE
                { left = insert (newItem, left)
                , right = right
                , x = ox
                , y = oy
                , ex = oex
                , ey = oey
                , data = oldData
                }
          | GREATER =>
              NODE
                { right = insert (newItem, right)
                , left = left
                , x = ox
                , y = oy
                , ex = oex
                , ey = oey
                , data = oldData
                }
          | EQUAL =>
              NODE
                { left = left
                , right = right
                , x = x
                , y = y
                , ex = ex
                , ey = ey
                , data = data
                }
        end

  fun foldr (f, tree, acc) =
    case tree of
      NODE {x, y, ex, ey, data, left, right} =>
        let
          val acc = foldr (f, right, acc)
          val item = {x = x, y = y, ex = ex, ey = ey, data = data}
          val acc = f (item, acc)
        in
          foldr (f, left, acc)
        end
    | LEAF => acc

  fun toList tree = foldr (fn (item, acc) => item ::acc, tree, [])
end

structure CollisionTree =
struct
  (* ignoreData = any data we find in grid but which doesn't concern us
   * so we omit from tree, to reduce number of items in tree,
   * and therefore decrease the constant in O(n) *)
  val ignoreData = 0

  local
    fun loopYAxis (x, y, eX, eY, yAxis, col) =
      if y > eY orelse y >= Vector.length yAxis then
        true
      else
        let
          val newCol = Vector.sub (yAxis, y)
        in
          if col = newCol then loopYAxis (x, y + 1, eX, eY, yAxis, col)
          else false
        end

    fun loopColour (x, y, eX, eY, grid, col) =
      if x > eX orelse x >= Vector.length grid then
        true
      else
        let
          val yAxis = Vector.sub (grid, x)
        in
          if loopYAxis (x, y, eX, eY, yAxis, col) then
            loopColour (x + 1, y, eX, eY, grid, col)
          else
            false
        end
  in
    fun quadHasSameColour (startX, startY, endX, endY, grid) =
      let
        val yAxis = Vector.sub (grid, startX)
        val col = Vector.sub (yAxis, startY)
      in
        loopColour (startX, startY, endX, endY, grid, col)
      end
  end

  (* tree creation *)
  fun build (x, y, size, grid, bintree) =
    if x >= Vector.length grid orelse y >= Vector.length grid then
      bintree
    else if quadHasSameColour (x, y, x + size, y + size, grid) then
      let
        val yAxis = Vector.sub (grid, x)
        val data = Vector.sub (yAxis, y)
      in
        if data = ignoreData then
          bintree
        else
          let
            val ex = x + size
            val ey = y + size
            val item = {x = x, y = y, ex = ex, ey = ey, data = data}
          in
            BinTree.insert (item, bintree)
          end
      end
    else
      (if size mod 2 = 0 orelse size = 1 then
         let
           val halfSize = size div 2
           val bintree = build (x, y, halfSize, grid, bintree)
           val bintree = build (x + halfSize, y, halfSize, grid, bintree)
           val bintree = build (x, y + halfSize, halfSize, grid, bintree)
         in
           build (x + halfSize, y + halfSize, halfSize, grid, bintree)
         end
       else
         (* handles odd-number divisions.
          * For example, `7 div 2` is 3 because of integer division. 
          * We would not cover every pixel unless we handle odd numbers specially. *)
         let
           val halfSizeBefore = size div 2
           val halfSizeAfter = (size + 1) div 2
           val bintree = build (x, y, halfSizeAfter, grid, bintree)
           val bintree = build
             (x + halfSizeBefore, y, halfSizeAfter, grid, bintree)
           val bintree = build
             (x, y + halfSizeBefore, halfSizeAfter, grid, bintree)
         in
           build
             ( x + halfSizeBefore
             , y + halfSizeBefore
             , halfSizeAfter
             , grid
             , bintree
             )
         end)

    fun getClickPoint (clickPoints, pos) =
      let val idx = Int.min (pos, Vector.length clickPoints - 1)
      in Vector.sub (clickPoints, idx)
      end

  fun folder
    ( windowWidth
    , windowHeight
    , canvasWidth
    , canvasHeight
    , xClickPoints
    , yClickPoints
    ) ({x, ex, y, ey, data}, acc) =
    let
      val ex = if ex = x then x + 1 else ex
      val ey = if ey = y then y + 1 else ey

      val x = getClickPoint (xClickPoints, x)
      val y = getClickPoint (yClickPoints, y)
      val ex = getClickPoint (xClickPoints, ex)
      val ey = getClickPoint (yClickPoints, ey)

      val startX = Ndc.fromPixelX (x, windowWidth, windowHeight)
      val endX = Ndc.fromPixelX (ex, windowWidth, windowHeight)
      val startY = Ndc.fromPixelY (y, windowWidth, windowHeight)
      val endY = Ndc.fromPixelY (ey, windowWidth, windowHeight)
    in
      Ndc.ltrbToVertexRgb (startX, startY, endX, endY, 0.0, 0.0, 0.0) :: acc
    end

  fun toTriangles
    ( windowWidth
    , windowHeight
    , squares
    , size
    , canvasWidth
    , canvasHeight
    , xClickPoints
    , yClickPoints
    ) =
    let
      val bintree = build (0, 0, size, squares, BinTree.empty)

      val f = folder
        ( windowWidth
        , windowHeight
        , canvasWidth
        , canvasHeight
        , xClickPoints
        , yClickPoints
        )
      val vec = BinTree.foldr (f, bintree, [])
    in
      Vector.concat vec
    end

  (* building and querying quad tree, plus compression *)
  datatype quad_tree =
    LEAF of {x: int, y: int, ex: int, ey: int, data: int}
  | NODE of {tl: quad_tree, tr: quad_tree, bl: quad_tree, br: quad_tree}
  | EMPTY

  fun foldWithDuplicates (f, tree, acc) =
    case tree of
      EMPTY => acc
    | LEAF item => f (item, acc)
    | NODE {tl, tr, bl, br} =>
        let
          val acc = foldWithDuplicates (f, tl, acc)
          val acc = foldWithDuplicates (f, tr, acc)
          val acc = foldWithDuplicates (f, bl, acc)
        in
          foldWithDuplicates (f, br, acc)
        end

  fun toBintree qtree =
    foldWithDuplicates (BinTree.insert, qtree, BinTree.empty)

  fun buildTree (x, y, size, grid) =
    if x >= Vector.length grid orelse y >= Vector.length grid then
      EMPTY
    else if quadHasSameColour (x, y, x + size, y + size, grid) then
      let
        val yAxis = Vector.sub (grid, x)
        val data = Vector.sub (yAxis, y)
      in
        if data = ignoreData then
          EMPTY
        else
          let
            val ex = x + size
            val ey = y + size
          in
            LEAF {x = x, y = y, ex = ex, ey = ey, data = data}
          end
      end
    else
      (if size mod 2 = 0 orelse size = 1 then
         let
           val halfSize = size div 2
           val tl = buildTree (x, y, halfSize, grid)
           val tr = buildTree (x + halfSize, y, halfSize, grid)
           val bl = buildTree (x, y + halfSize, halfSize, grid)
           val br = buildTree (x + halfSize, y + halfSize, halfSize, grid)
         in
           NODE {tl = tl, tr = tr, bl = bl, br = br}
         end
       else
         (* handles odd-number divisions.
          * For example, `7 div 2` is 3 because of integer division. 
          * We would not cover every pixel unless we handle odd numbers specially. *)
         let
           val halfSizeBefore = size div 2
           val halfSizeAfter = (size + 1) div 2
           val tl = buildTree (x, y, halfSizeAfter, grid)
           val tr = buildTree (x + halfSizeBefore, y, halfSizeAfter, grid)
           val bl = buildTree (x, y + halfSizeBefore, halfSizeAfter, grid)
           val br =
             buildTree
               (x + halfSizeBefore, y + halfSizeBefore, halfSizeAfter, grid)
         in
           NODE {tl = tl, tr = tr, bl = bl, br = br}
         end)

    local
      fun loop (x, y, ex, ey, grid) =
        if x > 0 then
          if quadHasSameColour (x, y, ex, ey, grid) then
            loop (x - 1, y, x, ey, grid)
          else
            ex
        else
          0
    in
      fun getLeftmostX ({x, y, ex, ey, data}, grid) =
        loop (x - 1, y, x, ey, grid)
    end

    local
      fun loop (x, y, ex, ey, grid) =
        if x < Vector.length grid andalso ex < Vector.length grid then
          if quadHasSameColour (x, y, ex, ey, grid) then
            loop (ex, y, ex + 1, ey, grid)
          else
            x
        else
          Vector.length grid - 1
    in
      fun getRightmostX ({x, y, ex, ey, data}, grid) =
        loop (ex, y, ex + 1, ey, grid)
    end

    local
      fun loop (x, y, ex, ey, grid) =
        if y < 0 then 0
        else
          if quadHasSameColour (x, y, ex, ey, grid) then
            loop (x, y - 1, ex, y, grid)
          else
            ey
    in
      fun getTopmostY ({x, y, ex, ey, data}, grid) =
        loop (x, y - 1, ex, y, grid)
    end

    local
      fun loop (x, y, ex, ey, grid) =
        if y < Vector.length grid andalso ey < Vector.length grid then
          if quadHasSameColour (x, y, ex, ey, grid) then
            loop (x, ey, ex, ey + 1, grid)
          else
            y
        else
          Vector.length grid - 1
    in
      fun getBottomY ({x, y, ex, ey, data}, grid) =
        loop (x, ey, ex, ey + 1, grid)
    end
end