# Tree — foundational binary tree for Bend.
# Publish entry for this package. Depends only on Base.
#
# Encoding: Leaf{} (empty) | Node{left, value, right}. Values live at nodes.
# Quantity: Tree<a, A> is Kind(a), same convention as List/Maybe.
# Map uses parallel fork on subtrees (Array.map style) when both sides recurse.
import Base

type Tree<a, -A: Kind(a)> is Kind(a):
  Leaf{}
  Node{left: Tree<a, A>, value: A, right: Tree<a, A>}

def Tree.empty(a, -A: Kind(a)) -> Tree<a, A>:
  Leaf{}

# Alias for the empty leaf constructor (values live on Node).
def Tree.leaf(a, -A: Kind(a)) -> Tree<a, A>:
  Leaf{}

def Tree.singleton(a, -A: Kind(a), x: A) -> Tree<a, A>:
  Node{Leaf{}, x, Leaf{}}

def Tree.is_empty(a, -A: Kind(a), t: Tree<a, A>) -> Bool:
  match t:
    case Leaf{}:
      True{}
    case Node{l, v, r}:
      False{}

def Tree.size(a, -A: Kind(a), t: Tree<a, A>) -> Nat:
  match t:
    case Leaf{}:
      0n
    case Node{l, v, r}:
      1n+Nat.add(Tree.size(a, A, l), Tree.size(a, A, r))

# Template map, same shape as List.map (affine Tree<A> = Tree<&1, A>).
# Subtrees mapped in parallel via `nl nr = ... ...`.
def Tree.map(~A: Type, ~B: Type, ~f: A -> B, t: Tree<A>) -> Tree<B>:
  match t:
    case Leaf{}:
      Leaf{}
    case Node{l, v, r}:
      nl nr = Tree.map(~A, ~B, ~f, l) Tree.map(~A, ~B, ~f, r)
      Node{nl, f(v), nr}

# Catamorphism: ~leaf/~node so both branches can reuse them (tree fan-out).
def Tree.fold(
  ~a: Quant, ~A: Kind(a), ~B: Type, ~leaf: B, ~node: B -> A -> B -> B, t: Tree<a, A>
) -> B:
  match t:
    case Leaf{}:
      leaf
    case Node{l, v, r}:
      node(
        Tree.fold(~a, ~A, ~B, ~leaf, ~node, l),
        v,
        Tree.fold(~a, ~A, ~B, ~leaf, ~node, r)
      )

# Combiner for reduce: (left ⊕ value) ⊕ right.
def Tree.reduce.node(~A: Type, ~f: A -> A -> A, lv: A, x: A, rv: A) -> A:
  f(f(lv, x), rv)

# Monoid-style reduce with zero (~z duplicated across branches).
def Tree.reduce(
  ~a: Quant, ~A: Kind(a), ~f: A -> A -> A, ~z: A, t: Tree<a, A>
) -> A:
  Tree.fold(~a, ~A, ~A, ~z, ~Tree.reduce.node(~A, ~f), t)

# In-order flattening.
def Tree.to_list(a, -A: Kind(a), t: Tree<a, A>) -> List<a, A>:
  match t:
    case Leaf{}:
      Nil{}
    case Node{l, v, r}:
      List.append(a, A, Tree.to_list(a, A, l), v <> Tree.to_list(a, A, r))