# UnionFind — functional disjoint-set structure for U32 elements.
#
# The parent map is keyed by the canonical String form of each U32. An absent
# key denotes a singleton root. Union links one root to the other; find is a
# structurally terminating, map-preserving walk fueled by the map size.
import Base

type UnionFind is Data:
  UF{parent: Map<&2, U32>}

def UnionFind.empty() -> UnionFind:
  UF{Map.new(&2, U32)}

def UnionFind.new() -> UnionFind:
  UnionFind.empty()

def UnionFind.key(x: U32) -> String:
  U32.show(x)

def UnionFind.parent_of(r: Map<&2, U32> & U32) -> U32:
  (parent, p) = r
  p

def UnionFind.find.go(fuel: Nat, +parent: Map<&2, U32>, +x: U32) -> U32:
  match fuel:
    case 0n:
      x
    case 1n+rest:
      +p = UnionFind.parent_of(Map.get(U32, x, parent, UnionFind.key(x)))
      Bool.pick(U32, U32.is_eq(x, p), x, UnionFind.find.go(rest, parent, p))

def UnionFind.find(+uf: UnionFind, x: U32) -> U32:
  match uf:
    case UF{+parent}:
      UnionFind.find.go(1n+Map.size(&2, U32, parent), parent, x)

def UnionFind.union(uf: UnionFind, +x: U32, +y: U32) -> UnionFind:
  match uf:
    case UF{+parent}:
      +rx = UnionFind.find.go(1n+Map.size(&2, U32, parent), parent, x)
      +ry = UnionFind.find.go(1n+Map.size(&2, U32, parent), parent, y)
      Bool.pick(UnionFind, U32.is_eq(rx, ry), UF{parent},
        UF{Map.set(&2, U32, parent, UnionFind.key(ry), rx)})

def UnionFind.connected(+uf: UnionFind, x: U32, y: U32) -> Bool:
  U32.is_eq(UnionFind.find(uf, x), UnionFind.find(uf, y))