Union Find with path compression and union-by-rank

unionfind.sml

signature UNION_FIND =
sig
    type t
    eqtype set

    val new   : t -> set
    val union : set -> set -> unit
    val repr  : set -> set
    val same  : set -> set -> bool
end

functor UnionFind(M : sig type t end) :> UNION_FIND where type t = M.t =
struct
  type t = M.t

  datatype set = Set of { guid   : unit ref
                        , height : int ref
                        , parent : set ref }
               | Empty
  fun set f (Set r) = f r
    | set _ _  = raise Fail "Impossible"

  fun new t =
      let val r = ref ()
          val s = {guid = r, height = ref 0, parent = ref Empty}
      in #parent s := Set s; Set s end

  fun repr (s as Set {guid, height, parent}) =
      if guid = set #guid (! parent)
      then s
      else (parent := repr (! parent); !parent)
    | repr _ = raise Fail "Impossible"

  fun union l r =
      let val (lp, rp) = (! (set #parent l), ! (set #parent r))
          val (hl, hr) = (set #height lp, set #height rp)
      in if !hl < !hr
         then (hr := (!hl + !hr); set #parent rp := lp)
         else (hl := (!hl + !hr); set #parent lp := rp)
      end

  fun same l r = set #guid (repr l) = set #guid (repr r)
end