Basic proofs of balanced binary trees. Scala examples of exercises from CM 500 Logical Frameworks (see https://github.com/julianpeeters/hello-twelf/tree/main)

bbt.scala

sealed trait Nat
case object Zero extends Nat
case class Successor[N <: Nat](nat: N) extends Nat

/* Max, implemented in Twelf, then as a Scala def, and finally as a Scala match types
 * 
 * max : nat -> nat -> nat -> type.
 * %mode max +N1 +N2 -N3.
 * max/z1 : max z N N.
 * max/z2 : max N z N.
 * max/s : max (s N) (s M) (s P) <- max N M P.
 * %worlds () (max _ _ _).
 * %total [N1 N2] (max N1 N2 _).
 *
 * def max(n1: Nat, n2: Nat): Nat =
 *   (n1, n2) match
 *     case (Zero, n) => n
 *     case (n, Zero) => n
 *     case (Successor(n), Successor(m)) => max(n1, n2)
 */
type Max[M, N] = (M, N) match
  case (Zero.type, n) => n
  case (n, Zero.type) => n
  case (Successor[m], Successor[n]) => Max[m, n]

/* Tree, implemented in twelf and then as a Scala ADT
 *
 * tree : type.
 * leaf : tree.
 * node : tree -> tree -> tree.
 */
sealed trait Tree
case object Leaf extends Tree
case class Node[L <: Tree, R <: Tree](l: L, r: R) extends Tree


/* Height, implemented in Twelf, then as a Scala def, and finally Scala match types
 *
 * height : tree -> nat -> type.
 * %mode height +T -N.
 * height/leaf : height leaf z.
 * height/node 
 *   : height (node L R) (s H)
 *   <- height L HL
 *   <- height R HR
 *   <- max HL HR H.
 * %worlds () (height _ _).
 * %total T (height T _).
 *
 * def height(tree: Tree): Nat =
 *   tree match
 *     case Leaf => Zero
 *     case Node(l, r) => Successor(max(height(l), height(r)))
 */
type Height[T] = T match
  case Leaf.type => Zero.type
  case Node[l, r] => Successor[Max[Height[l], Height[r]]]

/* Balanced, implemented in Twelf, then as a Scala ADT
 *
 *  balanced : tree -> type.
 *  balanced/leaf : balanced leaf.
 *  balanced/node
 *    : balanced (node L R)
 *      <- height L H
 *      <- height R H.
 */
sealed trait Balanced[T]
case object BalancedLeaf extends Balanced[Leaf.type]
case class BalancedNode[L <: Tree, R <: Tree](node: Node[L, R])(using ev: Height[L] =:= Height[R]) extends Balanced[Node[L, R]]

// Instances of the Balanced type.
val r: Balanced[Leaf.type] = BalancedLeaf
val s: Balanced[Node[Leaf.type, Leaf.type]] = BalancedNode(Node(Leaf, Leaf))

// However, we don't want the following tree to be able to be instantiated:
//         O
//        / \
//       /   \
//      O     O
//     / \   / \
//    O   * *   O
//   / \       / \
//  *   *     *   *
val t = BalancedNode(Node(Node(Node(Leaf, Leaf), Leaf), Node(Leaf, Node(Leaf, Leaf))))

// So, try another encoding:
sealed trait Balanced2[T]
implicit object Balanced2Leaf extends Balanced2[Leaf.type]
implicit class Balanced2Node[L <: Tree, R <: Tree](node: Node[L, R])(using ev: Height[L] =:= Height[R], l: Balanced2[L], r: Balanced2[R]) extends Balanced2[Node[L, R]]

// This works better, but alas, we need to supply some of our own evidence:
given Balanced2Node[Leaf.type, Leaf.type] =
  Balanced2Node(Node(Leaf, Leaf))

given Balanced2Node[Playground.Node[Leaf.type, Leaf.type], Playground.Node[Leaf.type, Leaf.type]] =
  Balanced2Node(Node(Node(Leaf, Leaf), Node(Leaf, Leaf)))

// correctly cannot prove height are equal
//given BalancedNode[Playground.Node[Node[Leaf.type, Leaf.type], Leaf.type], Playground.Node[Leaf.type, Node[Leaf.type, Leaf.type]]] =
//  BalancedNode(Node(Node(Node(Leaf, Leaf), Leaf), Node(Leaf, Node(Leaf, Leaf))))
//val s = BalancedNode(Node(Node(Node(Leaf, Leaf), Leaf), Node(Leaf, Node(Leaf, Leaf))))



// So perhaps it's easier using match types:
type Bal[T] = T match
  case Leaf.type => EqualsZero[Zero.type, Leaf.type]
  case Node[l, r] => EqualsZero[Difference[Height[Bal[l]], Height[Bal[r]]], Node[l, r]]

type Difference[T1, T2] = (T1, T2) match
  case (Zero.type, n) => n
  case (n, Zero.type) => n
  case (Successor[m], Successor[n]) => Difference[m, n]

type EqualsZero[N, T] = (N, T) match
  case (Zero.type, t) => t

val aa: Bal[Leaf.type] = Leaf
val bb: Bal[Node[Leaf.type, Leaf.type]] = Node(Leaf, Leaf)
val cc: Bal[Node[Node[Leaf.type, Leaf.type], Node[Leaf.type, Leaf.type]]] = Node(Node(Leaf, Leaf), Node(Leaf, Leaf))
// correctly does not compile:
// val dd: Bal[Node[Node[Node[Leaf.type, Leaf.type], Leaf.type], Node[Leaf.type, Node[Leaf.type, Leaf.type]]]] = Node(Node(Node(Leaf, Leaf), Leaf), Node(Leaf, Node(Leaf, Leaf)))





/* "intrinsically typed" tree:
 *
 * bbt : nat -> type.
 * bLeaf : bbt z.
 * bNode : bbt N -> bbt N -> bbt (s N).
 */
sealed trait BBT[T]
case object BLeaf extends BBT[Zero.type]
case class BNode[N <: Nat, L <: BBT[N], R <: BBT[N]](l: L, r: R) extends BBT[Successor[N]]

/* Equality, in Twelf and then in Scala
 *
 * eqTree : bbt N -> bbt N -> type.
 * %mode eqTree +T1 +T2.
 * eqTree/leaf : eqTree bLeaf bLeaf.
 * eqTree/node : eqTree (bNode L1 R1) (bNode L2 R2)
 *                <- eqTree L1 L2
 *                <- eqTree R1 R2.
 *
 * %worlds () (eqTree _ _).
 * %total [T1 T2] (eqTree T1 T2).
 */
sealed trait EqTree[T1, T2]
case class EqLeaf(t1: BLeaf.type, t2: BLeaf.type) extends EqTree[BLeaf.type, BLeaf.type]
case class EqNode[
  N <: Nat,
  L1 <: BBT[N],
  L2 <: BBT[N],
  R1 <: BBT[N],
  R2 <: BBT[N]
](t1: BNode[N, L1, R1], t2: BNode[N, L2, R2])(using ev: ValueOf[EqTreeEv[BNode[N, L1, R1], BNode[N, L2, R2]]]) extends EqTree[BNode[N, L1, R1], BNode[N, L2, R2]]

type EqTreeEv[T1, T2] = (T1, T2) match
  case (BLeaf.type, BLeaf.type) => Unit
  case (BNode[n1, l1, r1], BNode[n2, l2, r2]) => Unite[EqTreeEv[l1, l2], EqTreeEv[r1, r2]]

type Unite[A, B] = (A, B) match
  case ((Unit, Unit), (Unit, Unit)) => Unit
  case (Unit, Unit) => Unit


// Instances of BBTs, showing unbalanced trees cannot be constructed.
val u: BBT[Zero.type] = BLeaf
val v: BBT[Successor[Zero.type]] = BNode(BLeaf, BLeaf)

println(EqLeaf(BLeaf, BLeaf))
println(EqNode(BNode(BLeaf, BLeaf), BNode(BLeaf, BLeaf)))