// YOUR JOB: fill in the last predicate definition.
// It should help to copy in the definition of Valid, too, maybe from what we
// just worked out together.

module Functions {
  // First, here is what it means for such a function to agree with a set of
  // required equalities.
  ghost predicate Compatible<K(!new)>(f: K -> nat, eqs: set<(K, K)>) {
    forall u, v :: (u, v) in eqs ==> f(u) == f(v)
  }

  // Next, here is what it means for a set of equalities to imply that two
  // keys are equal *under all compatible functions*.
  ghost predicate ImpliedEqual<K(!new)>(eqs: set<(K, K)>, u: K, v: K) {
    forall f :: Compatible(f, eqs) ==> f(u) == f(v)
  }
}

module UnionFind {
  import F = Functions
  
  // As the data structure is informed of new keys, we assign them
  // *classes* which we represent as natural numbers.
  type cls = nat

  // Here is our internal data structure, which works for any key type that
  // supports equality.
  datatype UnionFind<K(==)> = UF(classOf: map<K, cls>, merges: map<cls, cls>, nextClass: cls)
  // `classOf`: maps known keys to their classes.
  // `merges`: includes a mapping from `u` to `v` to indicate that `u` is known
  //           to be equal to `v`.
  // `nextClass`: a number higher than any used so far as a class.

  // The basic idea will be that, when we want to check if two keys are
  // known as equal, we look up their classes and then, for each class, keep
  // following `Merges` until there are no more links to follow.  The keys
  // are known as equal exactly when this process yields the same class for
  // each one.

  // Problem: how do we know that link-following terminates?
  // We'll adopt a simple heuristic here of only allowing a class to be linked
  // to a *lower* one.
  // The following predicate records enough information to maintain that
  // property.

  ghost predicate Valid<K>(uf: UnionFind<K>)

  ghost predicate Represents<K(!new)>(uf: UnionFind<K>, eqs: set<(K, K)>)
    requires Valid(uf)

  // It's probably helpful to mention FollowLinks in Represents, so here is its type.
  function FollowLinks<K>(uf: UnionFind<K>, cls: cls): cls
}