structure LayerTree =
struct
  type square = {r: int, g: int, b: int, a: int}
  type grid = square vector vector

  fun isBlank ({a, ...}: square) = a = 0

  datatype tree =
    NODE of {key: int, value: grid, left: tree, right: tree}
  | LEAF

  fun insert (newKey, newValue, tree) =
    case tree of
      LEAF => NODE {key = newKey, value = newValue, left = LEAF, right = LEAF}
    | NODE {key, value, left, right} =>
        if newKey < key then
          let val left = insert (newKey, newValue, left)
          in NODE {key = key, value = value, left = left, right = right}
          end
        else if newKey > key then
          let val right = insert (newKey, newValue, right)
          in NODE {key = key, value = value, left = left, right = right}
          end
        else
          NODE {key = key, value = newValue, left = left, right = right}

  fun get (searchKey, tree) =
    case tree of
      LEAF => NONE
    | NODE {key, value, left, right} =>
        if searchKey < key then get (searchKey, left)
        else if searchKey > key then get (searchKey, right)
        else SOME value

  (* copies non-blank pixels in value vector into acc *)
  fun helpFlatten (value, acc) =
    Vector.mapi
      (fn (xIdx, valueYAxis) =>
         Vector.mapi
           (fn (yIdx, valuePixel) =>
              if isBlank valuePixel then
                let val accYAxis = Vector.sub (acc, xIdx)
                in Vector.sub (accYAxis, yIdx)
                end
              else
                valuePixel) valueYAxis) value

  fun flatten (tree, acc) =
    case tree of
      LEAF => acc
    | NODE {value, left, right, ...} =>
        let
          val acc = flatten (left, acc)
          val acc = helpFlatten (value, acc)
        in
          flatten (right, acc)
        end
end