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
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)

  local
    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
  in
    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
  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 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)

  fun getItemWithDataAt (x, y, qx, qy, size, tree, data) =
    case tree of
      EMPTY => NONE
    | LEAF (item as {x = ix, y = iy, ex = iex, ey = iey, data = oldData}) =>
        if (x >= ix andalso x <= iex) andalso (y >= iy andalso y <= iey) then
          (* search coordinates are in item *)
          if data = oldData then (* data matches *) SOME item else NONE
        else
          NONE
    | NODE {tl, tr, bl, br} =>
        if size mod 2 = 0 orelse size = 1 then
          let
            val halfSize = size div 2
            val qmx = x + halfSize
            val qfx = x + size
            val qmy = y + halfSize
            val qfy = y + size
          in
            if y >= qy andalso y <= qmy then
              (* top *)
              if x >= qx andalso x <= qmx then
                (* top left *)
                getItemWithDataAt (x, y, qx, qy, halfSize, tl, data)
              else
                (* top right *)
                getItemWithDataAt (x, y, qx + halfSize, qy, halfSize, tr, data)
            else (* bottom *) if x >= qx andalso x <= qmx then
              (* bottom left *)
              getItemWithDataAt (x, y, qx, qy + halfSize, halfSize, bl, data)
            else
              (* bottom right *)
              getItemWithDataAt
                (x, y, qx + halfSize, qy + halfSize, halfSize, br, data)
          end
        else
          let
            val halfSizeBefore = size div 2
            val halfSizeAfter = (size + 1) div 2
            val qmx = x + halfSizeBefore
            val qmy = y + halfSizeBefore
          in
            if y >= qy andalso y <= qmy then
              (* top *)
              if x >= qx andalso x <= qmx then
                (* top left *)
                getItemWithDataAt (x, y, qx, qy, halfSizeAfter, tl, data)
              else
                (* top right *)
                getItemWithDataAt (x, y, qmx, qy, halfSizeAfter, tr, data)
            else (* bottom *) if x >= qx andalso x <= qmx then
              (* bottom left *)
              getItemWithDataAt (x, y, qx, qmy, halfSizeAfter, bl, data)
            else
              (* bottom right *)
              getItemWithDataAt (x, y, qmx, qmy, halfSizeAfter, br, data)
          end

  fun getLeftmostX (rootSize, rootTree, prevItem) =
    let
      val {x = prevX, y = prevY, ex = prevEx, ey = prevEy, data = prevData} =
        prevItem
    in
      case
        getItemWithDataAt (prevX - 1, prevY, 0, 0, rootSize, rootTree, prevData)
      of
        SOME (newItem as {y = newY, ey = newEy, data = newData, ...}) =>
          if prevY = newY andalso newEy = prevEy andalso prevData = newData then
            (* Y side has same edge, so is mergeable, and data is also same *)
            getLeftmostX (rootSize, rootTree, newItem)
          else
            prevX
      | NONE => prevX
    end

  fun getRightmostX (rootSize, rootTree, prevItem) =
    let
      val {x = prevX, y = prevY, ex = prevEx, ey = prevEy, data = prevData} =
        prevItem
    in
      case
        getItemWithDataAt
          (prevEx + 1, prevY, 0, 0, rootSize, rootTree, prevData)
      of
        SOME (newItem as {y = newY, ey = newEy, data = newData, ...}) =>
          if prevY = newY andalso newEy = prevEy andalso prevData = newData then
            getRightmostX (rootSize, rootTree, newItem)
          else
            prevEx
      | NONE => prevEx
    end

  fun getTopmostY (rootSize, rootTree, prevItem) =
    let
      val {x = prevX, y = prevY, ex = prevEx, ey = prevEy, data = prevData} =
        prevItem
    in
      case
        getItemWithDataAt (prevX, prevY - 1, 0, 0, rootSize, rootTree, prevData)
      of
        SOME (newItem as {x = newX, ex = newEx, data = newData, ...}) =>
          if prevX = newX andalso prevEx = newEx andalso prevData = newData then
            (* X side has same edge and data is also same, so mergeable *)
            getTopmostY (rootSize, rootTree, newItem)
          else
            prevY
      | NONE => prevY
    end

  fun getBottomY (rootSize, rootTree, prevItem) =
    let
      val {x = prevX, y = prevY, ex = prevEx, ey = prevEy, data = prevData} =
        prevItem
    in
      case
        getItemWithDataAt
          (prevX, prevEy + 1, 0, 0, rootSize, rootTree, prevData)
      of
        SOME (newItem as {x = newX, ex = newEx, data = newData, ...}) =>
          if prevX = newX andalso prevEx = newEx andalso prevData = newData then
            (* X side has same edge and data is also same, so mergeable *)
            getBottomY (rootSize, rootTree, newItem)
          else
            prevEy
      | NONE => prevEy
    end

  datatype merge_dir =
    HORIZONTAL of {left: int, right: int}
  | VERTICAL of {up: int, down: int}
  | NO_MERGE

  fun getHorizontal (ox, oex, leftX, rightX) =
    let
      val left = if leftX < ox then ox - 1 else ox
      val right = if rightX > oex then oex + 1 else oex
    in
      HORIZONTAL {left = left, right = right}
    end

  fun getVertical (oy, oey, upY, bottomY) =
    let
      val up = if upY < oy then oy - 1 else oy
      val down = if bottomY > oey then oey + 1 else oey
    in
      VERTICAL {up = up, down = down}
    end

  fun getMergeDir (rootSize, rootTree, item) =
    let
      val {x = ox, y = oy, ex = oex, ey = oey, data} = item

      val leftX = getLeftmostX (rootSize, rootTree, item)
      val rightX = getRightmostX (rootSize, rootTree, item)
      val upY = getTopmostY (rootSize, rootTree, item)
      val bottomY = getBottomY (rootSize, rootTree, item)

      val xChanged = leftX <> ox orelse rightX <> oex
      val yChanged = upY <> oy orelse bottomY <> oey
    in
      if xChanged andalso yChanged then
        let
          val diffX = abs (rightX - leftX)
          val diffY = abs (bottomY - upY)
        in
          if diffX > diffY then getHorizontal (ox, oex, leftX, rightX)
          else getVertical (oy, oey, upY, bottomY)
        end
      else if xChanged then
        getHorizontal (ox, oex, leftX, rightX)
      else if yChanged then
        getVertical (oy, oey, upY, bottomY)
      else
        NO_MERGE
    end

  fun mergeHorizontal (item, newX, newEx) =
    let val {y, ey, data, x = _, ex = _} = item
    in {x = newX, ex = newEx, y = y, ey = ey, data = data}
    end

  fun mergeVertical (item, newY, newEy) =
    let val {x, ex, data, y = _, ey = _} = item
    in {y = newY, ey = newEy, x = x, ex = ex, data = data}
    end

  local
    fun mergePass (qx, qy, size, tree, rootSize, rootTree) =
      case tree of
        EMPTY => (EMPTY, false)
      | LEAF item =>
          let
            val mergeDir = getMergeDir (rootSize, rootTree, item)
          in
            case mergeDir of
              NO_MERGE => (LEAF item, false)
            | VERTICAL {up, down} =>
                (LEAF (mergeVertical (item, up, down)), true)
            | HORIZONTAL {left, right} =>
                (LEAF (mergeHorizontal (item, left, right)), true)
          end
      | NODE {tl, tr, bl, br} =>
          if size = 1 orelse size mod 2 = 0 then
            let
              val halfSize = size div 2

              val (tl, didMergeTl) = mergePass
                (qx, qy, halfSize, tl, rootSize, rootTree)
              val (tr, didMergeTr) = mergePass
                (qx + halfSize, qy, halfSize, tr, rootSize, rootTree)
              val (bl, didMergeBl) = mergePass
                (qx, qy + halfSize, halfSize, bl, rootSize, rootTree)
              val (br, didMergeBr) = mergePass
                (qx + halfSize, qy + halfSize, halfSize, br, rootSize, rootTree)

              val didMergeAny =
                didMergeTl orelse didMergeTr orelse didMergeBl orelse didMergeBr
              val node = NODE {tl = tl, tr = tr, bl = bl, br = br}
            in
              (node, didMergeAny)
            end
          else
            let
              val halfSizeBefore = size div 2
              val halfSizeAfter = (size + 1) div 2
              val qmx = qx + halfSizeBefore
              val qmy = qy + halfSizeAfter

              val (tl, didMergeTl) = mergePass
                (qx, qy, halfSizeAfter, tl, rootSize, rootTree)
              val (tr, didMergeTr) = mergePass
                (qmx, qy, halfSizeAfter, tr, rootSize, rootTree)
              val (bl, didMergeBl) = mergePass
                (qx, qmy, halfSizeAfter, bl, rootSize, rootTree)
              val (br, didMergeBr) = mergePass
                (qmx, qmy, halfSizeAfter, br, rootSize, rootTree)

              val didMergeAny =
                didMergeTl orelse didMergeTr orelse didMergeBl orelse didMergeBr
              val node = NODE {tl = tl, tr = tr, bl = bl, br = br}
            in
              (node, didMergeAny)
            end
  in
    fun merge (rootSize, rootTree) =
      let
        val (newTree, didMerge) = mergePass
          (0, 0, rootSize, rootTree, rootSize, rootTree)
      in
        if didMerge then merge (rootSize, newTree) else newTree
      end
  end
end