private record Cell[T](value  : T,
                       rank   : SizeT,
                       parent : Pointer[Cell[T]]);
  
[T]
overload Cell[T](v : T) {
  return Cell[T](v, SizeT(0), null(Cell[T]));
}
  
[T] 
overload hash(c : Cell[T]) {
  return hash(c.value);
}   
    
[T]   
overload equals?(lhs : Cell[T], rhs : Cell[T]) {
  return equals?(lhs.value, rhs.value);
}   
    
record DisjointSet[T](hash : HashMap[T, Cell[T]]);

[T]
private compress_path(c : Cell[T]) ByRef[Cell[T]] {
  if (null?(c.parent))
    return ref c;
  c.parent <-- &compress_path(c.parent^);
  return ref c.parent^;
}

[T]
overload put(dj : DisjointSet[T], t : T) {
  put(dj.hash, t, Cell[T](t));
}     
    
[T] 
merge(dj : DisjointSet[T], t1 : T, t2 : T) {
  var i = lookup(dj.hash, t1);
  if (null?(i)) 
    return false;
  var j = lookup(dj.hash, t2);
  if (null?(j))
    return false;
  println("MERGE1 ", i, " ", j);
    
  var set_i = compress_path(i^);
  var set_j = compress_path(j^);
  if (&set_i == &set_j)
    return false;
  if (set_i.rank > set_j.rank)
    set_j.parent = &set_i;
  else {
    set_i.parent = &set_j;
    println("MERGE ", set_i.parent, " ", &set_j);
    if (set_i.rank == set_j.rank)
      set_j.rank <-- set_j.rank + 1;
  }
  return true;
}         
[T]
overload lookup(dj : DisjointSet[T], t : T) {
  var i = lookup(dj.hash, t);
  if (null?(i))
    return null(T);
  else
    return &compress_path(i^).value;
}

[T]
same_set(dj : DisjointSet[T], t1 : T, t2 : T) {
  var i = lookup(dj.hash, t1);
  if (null?(i))
    return false;
  var j = lookup(dj.hash, t2);
  if (null?(j))
    return false;

  var set_i = compress_path(i^);
  var set_j = compress_path(j^);
  println("COMP ", set_i, " ", &set_i, " ", set_j, " ", &set_j);
  return &set_i == &set_j;
}

main() {
  var ok = Cell[Int32](3);
  compress_path(ok);

  var dj = DisjointSet[Int32]();
  put(dj, 3);
  put(dj, 4);
  put(dj, 5);

  assert(not same_set(dj, 3, 4));
  assert(not same_set(dj, 3, 5));
  assert(not same_set(dj, 4, 5));
  merge(dj, 3, 4);
  assert(same_set(dj, 3, 4));
  assert(not same_set(dj, 3, 5));
  assert(not same_set(dj, 4, 5));
}