# Grid: a square quadtree, and the parallel walks over it.
# ========================================================
#
# A Grid of depth d holds 2^d * 2^d values. The four fields keep the
# screen order that Bend's Image uses, so a board maps onto a frame
# without a transpose.
#
# The grid is Data, so a `+` copy costs one reference count. That is
# what lets every leaf of a parallel walk read the whole previous
# board at the same time.

import Base

type Quad is Data:
  QTL{}
  QTR{}
  QBL{}
  QBR{}

type Grid<a, -A: Kind(a)> is Kind(a):
  GLeaf{val: A}
  GQuad{tl: Grid<a, A>, tr: Grid<a, A>, bl: Grid<a, A>, br: Grid<a, A>}

# Picks a quadrant from one bit of x and one bit of y.
def Quad.of(bx: Bool, by: Bool) -> Quad:
  match bx by:
    case False{} False{}:
      QTL{}
    case True{} False{}:
      QTR{}
    case False{} True{}:
      QBL{}
    case True{} True{}:
      QBR{}

# The route from the root of a depth-d grid down to cell (x, y).
# Bit d-1 of x and y picks the first turn, bit 0 the last.
def Grid.path(d: Nat, +x: U32, +y: U32) -> List<&2, Quad>:
  match d:
    case 0n:
      Nil{}
    case Succ{+p}:
      bx = U32.is_ne(U32.and(U32.shrn(x, p), 1), 0)
      by = U32.is_ne(U32.and(U32.shrn(y, p), 1), 0)
      Quad.of(bx, by) <> Grid.path(p, x, y)

# Reads the value at the end of a path.
# A path shorter than the grid stops at the top-left corner below it.
def Grid.get(-A: Data, g: Grid<&2, A>, path: List<&2, Quad>) -> A:
  match g path:
    case GLeaf{v} _:
      v
    case GQuad{tl, tr, bl, br} Nil{}:
      Grid.get(A, tl, Nil{})
    case GQuad{tl, tr, bl, br} Con{QTL{}, rest}:
      Grid.get(A, tl, rest)
    case GQuad{tl, tr, bl, br} Con{QTR{}, rest}:
      Grid.get(A, tr, rest)
    case GQuad{tl, tr, bl, br} Con{QBL{}, rest}:
      Grid.get(A, bl, rest)
    case GQuad{tl, tr, bl, br} Con{QBR{}, rest}:
      Grid.get(A, br, rest)

# Writes v at the end of a path. Only the nodes on the path change.
# A path shorter than the grid stops where Grid.get stops, at the
# top-left corner below it. That is what makes `read_back` hold for
# every path, not only for one of the right length.
def Grid.set(-A: Data, g: Grid<&2, A>, path: List<&2, Quad>, v: A) -> Grid<&2, A>:
  match g path:
    case GLeaf{old} _:
      GLeaf{v}
    case GQuad{tl, tr, bl, br} Nil{}:
      GQuad{Grid.set(A, tl, Nil{}, v), tr, bl, br}
    case GQuad{tl, tr, bl, br} Con{QTL{}, rest}:
      GQuad{Grid.set(A, tl, rest, v), tr, bl, br}
    case GQuad{tl, tr, bl, br} Con{QTR{}, rest}:
      GQuad{tl, Grid.set(A, tr, rest, v), bl, br}
    case GQuad{tl, tr, bl, br} Con{QBL{}, rest}:
      GQuad{tl, tr, Grid.set(A, bl, rest, v), br}
    case GQuad{tl, tr, bl, br} Con{QBR{}, rest}:
      GQuad{tl, tr, bl, Grid.set(A, br, rest, v)}

# Builds a depth-d grid from a function of the payload and the cell
# position. The four quadrants run in parallel and are the same size,
# which is what Bend's fork-join scheduler wants.
def Grid.init(
  ~A: Data, ~P: Data, ~f: @+q: P -> @+u: U32 -> @+v: U32 -> A,
  d: Nat, +p: P, +x: U32, +y: U32
) -> Grid<&2, A>:
  match d:
    case 0n:
      GLeaf{f(p, x, y)}
    case Succ{+k}:
      +h = U32.shln(1, k)
      a b c e =
        Grid.init(~A, ~P, ~f, k, p, x, y)
        Grid.init(~A, ~P, ~f, k, p, (x + h : U32), y)
        Grid.init(~A, ~P, ~f, k, p, x, (y + h : U32))
        Grid.init(~A, ~P, ~f, k, p, (x + h : U32), (y + h : U32))
      GQuad{a, b, c, e}

# Maps every value. The four quadrants run in parallel.
def Grid.map(~A: Data, ~B: Data, ~f: @c: A -> B, g: Grid<&2, A>) -> Grid<&2, B>:
  match g:
    case GLeaf{v}:
      GLeaf{f(v)}
    case GQuad{tl, tr, bl, br}:
      a b c e =
        Grid.map(~A, ~B, ~f, tl)
        Grid.map(~A, ~B, ~f, tr)
        Grid.map(~A, ~B, ~f, bl)
        Grid.map(~A, ~B, ~f, br)
      GQuad{a, b, c, e}

# Adds a number up over every value. The four quadrants run in parallel.
def Grid.sum(~A: Data, ~f: @c: A -> U32, g: Grid<&2, A>) -> U32:
  match g:
    case GLeaf{v}:
      f(v)
    case GQuad{tl, tr, bl, br}:
      a b c e =
        Grid.sum(~A, ~f, tl)
        Grid.sum(~A, ~f, tr)
        Grid.sum(~A, ~f, bl)
        Grid.sum(~A, ~f, br)
      (a + b + c + e : U32)