// Let's implement a functional variant of a classic data structure, union-find.
// See: https://en.wikipedia.org/wiki/Disjoint-set_data_structure
// The basic idea is that we tell the data structure some known equalities among
// a set of keys, and we can query it on whether another potential equality
// is implied by the ones we know.
// We will build union-find as a nice encapsulated data structure like the ones
// we studied in the previous two classes.
// We'll *skip* some popular optimizations like path compression, but
// incorporating them could be a fun exercise for the reader!

// We'll need some preliminaries about functions and sets of equalities over
// function keys.

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
  
  // Take note of which identifiers are revealed vs. provided!
  export
    provides UnionFind, Empty, AssertEqual, CheckEqual, Valid, Represents
    provides EmptyCorrect, CheckEqualCorrect, AssertEqualCorrect
    provides F

  // 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>) {
    (forall c :: c in uf.merges ==> uf.merges[c] < c)
      && (forall k :: k in uf.classOf ==> uf.classOf[k] < uf.nextClass)
      && (forall c :: c in uf.merges ==> c in uf.classOf.Values)
      && (forall c :: c in uf.merges ==> uf.merges[c] in uf.classOf.Values)
  }

  // An empty union-find is easy to initialize: we know nothing!
  function Empty<K>(): UnionFind<K>
    ensures Valid(Empty<K>())
  {
    UF(map[], map[], 0)
  }

  // This helper function implements our idea of repeated link-following.
  // Note that termination checking would fail if we didn't know `Valid`.
  function FollowLinks<K>(uf: UnionFind<K>, cls: cls): cls
    requires Valid(uf)
    requires cls in uf.classOf.Values
    ensures FollowLinks(uf, cls) in uf.classOf.Values
  {
    if cls in uf.merges then
      FollowLinks(uf, uf.merges[cls])
    else
      cls
  }

  // This helper combines looking up a key's class with then following all
  // links from that class.  Importantly, if the key is not yet known to
  // the data structure, we assign it a fresh class.  The body is complicated
  // mostly by including a proof of a key property: the set of used classes does
  // not shrink as we add classes.
  function GetClass<K(!new)>(uf: UnionFind<K>, u: K): (UnionFind<K>, cls)
    requires Valid(uf)
    ensures var (uf', c) := GetClass(uf, u);
      Valid(uf') && c in uf'.classOf.Values && uf'.classOf.Values >= uf.classOf.Values
  {
    if u in uf.classOf then
      // Easy case: key we already knew about
      (uf, FollowLinks(uf, uf.classOf[u]))
    else
      // Hard case: a first-time key

      // It's easy to construct the new union-find structure, but we have to
      // prove it's legit.
      var r := UF(uf.classOf[u := uf.nextClass], uf.merges, uf.nextClass + 1);

      // This next fact is the one we work hardest proving.
      // The next few lines are boilerplate, for proving a quantified
      // implication.
      assert r.classOf.Values >= uf.classOf.Values by {
        forall v ensures v in uf.classOf.Values ==> v in r.classOf.Values {
          if v in uf.classOf.Values {
            // The important hint to Dafny is here: a value must have a
            // corresponding key.
            var k :| k in uf.classOf && uf.classOf[k] == v;
            // And then that key retains its value in the new structure.
            assert r.classOf[k] == v;
          }
        }
      }

      assert r.classOf[u] == uf.nextClass;
      (r, uf.nextClass)
  }

  // IMPORTANT NOTE: those last two helpers are NOT exported from the module!
  // Client code needn't fuss with them at all.
  // Indeed, it's not allowed to access them.
  
  // Our setup makes it easy to define the predicate for checking if two
  // keys are known to be equal.
  predicate CheckEqual<K(!new)>(uf: UnionFind<K>, u: K, v: K)
    requires Valid(uf)
  {
    if u in uf.classOf && v in uf.classOf then
      FollowLinks(uf, uf.classOf[u]) == FollowLinks(uf, uf.classOf[v])
    else
      // Funny corner case here: there might be a key we haven't heard of
      // yet, but we DO know it's equal to itself.  (We can't say any more than
      // that.)
      u == v
  }

  // Adding a known equality is also now quite straightforward.
  function AssertEqual<K(!new)>(uf: UnionFind<K>, u: K, v: K): UnionFind<K>
    requires Valid(uf)
    ensures Valid(AssertEqual(uf, u, v))
  {
    var (uf, uc) := GetClass(uf, u);
    var (uf, vc) := GetClass(uf, v);
    if uc == vc then
      // Already known to be equal.  No action needed.
      // (However, note that GetClass may already have modified the
      // structure just to note the addition of a new key.)
      uf
    else if uc < vc then
      // In general, we want to redirect the higher class to the lower one.
      UF(uf.classOf, uf.merges[vc := uc], uf.nextClass)
    else
      UF(uf.classOf, uf.merges[uc := vc], uf.nextClass)
  }

  // OK, now we're ready to characterize what correctness means!
  // We will think in terms of functions that assign numbers to all keys.

  // Using that last definition, we can capture what it means for a union-find
  // structure to *represent* a set of known equalities.
  ghost predicate Represents<K(!new)>(uf: UnionFind<K>, eqs: set<(K, K)>)
    requires Valid(uf)
  {
    // Boringly, we start out requiring that all keys mentioned in the
    // equalities have been assigned classes.
    && (forall u, v :: (u, v) in eqs ==> u in uf.classOf)
    && (forall u, v :: (u, v) in eqs ==> v in uf.classOf)

    // Much more interestingly, consider any two keys that are assigned
    // classes.  Two characterizations of known equality must agree.
    && (forall u, v: K :: u in uf.classOf && v in uf.classOf ==>
        // Either we get to the same place following links...
        ((FollowLinks(uf, uf.classOf[u]) == FollowLinks(uf, uf.classOf[v]))
          // ...or our idea of implied equality just above holds.
          <==> F.ImpliedEqual(eqs, u, v)))
  }

  // Now we can march through proving each public union-find function correct,
  // according to this definition.

  // Easy warm-up!
  lemma EmptyCorrect<K(!new)>()
    ensures Represents<K>(Empty(), {})
  {
  }

  // We can explain what CheckEqualCorrect should return in terms of
  // F.ImpliedEqual.
  lemma CheckEqualCorrect<K(!new)>(uf: UnionFind<K>, eqs: set<(K, K)>, u: K, v: K)
    requires Valid(uf)
    requires Represents(uf, eqs)
    ensures CheckEqual(uf, u, v) == F.ImpliedEqual(eqs, u, v)
  {
    // Dafny can handle most of this proof automatically!  However, we need to
    // help out by constructing a function that is compatible with the known
    // equalities.  Actually, the union-find structure itself provides almost
    // all of that function.  Just in case either of `u` or `v` is not assigned
    // a class, we need to conservatively add in mappings for them that are
    // (1) not already used and (2) not equal to each other.
    var f: K -> nat := (k => if k in uf.classOf then FollowLinks(uf, uf.classOf[k]) else if k == u then uf.nextClass else uf.nextClass + 1);
    assert F.Compatible(f, eqs);
  }

  // This proof is the intricate part.  We start by just stepping through the
  // structure of the function, asserting a few helpful facts.
  // Most action gets delegated to the lemma called twice at the end.
  lemma AssertEqualCorrect<K(!new)>(uf: UnionFind<K>, eqs: set<(K, K)>, u: K, v: K)
    requires Valid(uf)
    requires Represents(uf, eqs)
    ensures Represents(AssertEqual(uf, u, v), eqs + {(u, v)})
  {
    var uf' := AssertEqual(uf, u, v);
    var eqs' := eqs + {(u, v)};
    
    var (uf1, uc) := GetClass(uf, u);
    GetClassCorrect(uf, eqs, u);
    assert uc == FollowLinks(uf1, uf1.classOf[u]);
    var (uf2, vc) := GetClass(uf1, v);
    GetClassCorrect(uf1, eqs, v);
    assert uc == FollowLinks(uf2, uf2.classOf[u]) by { GetClassPreservesFollowLinks(uf1, v); }
    assert vc == FollowLinks(uf2, uf2.classOf[v]);

    if uc == vc {
    } else if uc < vc {
      AssertEqualCorrect'(uf2, eqs, u, v);
    } else {
      AssertEqualCorrect'(uf2, eqs, v, u);
      RepresentsWeaken(uf', eqs + {(v, u)}, eqs');
    }
  }

  // However, this lemma is also interesting.  It expresses that the orders of
  // pair elements in sets of equalities don't actually matter.
  // The proof requires no induction and follows quite naturally from properties
  // of sets and equality.
  lemma RepresentsWeaken<K(!new)>(uf: UnionFind<K>, eqs1: set<(K, K)>, eqs2: set<(K, K)>)
    requires Valid(uf)
    requires Represents(uf, eqs1)
    requires forall a, b :: (a, b) in eqs1 ==> (a, b) in eqs2 || (b, a) in eqs2
    requires forall a, b :: (a, b) in eqs2 ==> (a, b) in eqs1 || (b, a) in eqs1
    ensures Represents(uf, eqs2)
  {
  }

  // Here's that lemma capturing the interesting cases of `AssertEqualCorrect`.
  // It's a "without loss of generality" opportunity in the proof above, where
  // when u and v don't have equal representatives, then one representative must
  // be less than the other.
  lemma AssertEqualCorrect'<K(!new)>(uf: UnionFind<K>, eqs: set<(K, K)>, u: K, v: K)
    requires Valid(uf)
    requires Represents(uf, eqs)
    requires u in uf.classOf && v in uf.classOf
    requires FollowLinks(uf, uf.classOf[u]) < FollowLinks(uf, uf.classOf[v])
    ensures var r := UF(uf.classOf, uf.merges[FollowLinks(uf, uf.classOf[v]) := FollowLinks(uf, uf.classOf[u])], uf.nextClass);
      // Note that this definition of `r` generalizes the two final cases of `AssertEqual`.
      Valid(r) ==> Represents(r, eqs + {(u, v)})
      // A little shortcut here: we already proved validity of this union-find
      // structure elsewhere, so we can proceed here conditionally in having
      // done that proof, rather than repeat it.
  {
    var uf' := UF(uf.classOf, uf.merges[FollowLinks(uf, uf.classOf[v]) := FollowLinks(uf, uf.classOf[u])], uf.nextClass);
    var eqs' := eqs + {(u, v)};
    
    if Valid(uf') {
      forall a, b: K ensures a in uf'.classOf && b in uf'.classOf ==>
        ((FollowLinks(uf', uf'.classOf[a]) == FollowLinks(uf', uf'.classOf[b]))
        <==> F.ImpliedEqual(eqs', a, b)) {
        if a in uf'.classOf && b in uf'.classOf {
          // Everything up to this point was boilerplate, given the lemma statement.
          // The two directions of the logical equivalence are where we get a little smart.
            
          if FollowLinks(uf', uf'.classOf[a]) == FollowLinks(uf', uf'.classOf[b]) {
            // If two keys have equal representatives, we can indeed derive that
            // the new equalities imply they are equal.  They key is to use the
            // following lemma invocation to "remind" Dafny about the behavior
            // of `FollowLinks` after the merge, on the two particular classes
            // we need to prove equal.
            MergeFollowLinks(uf, uf', uf.classOf[u], uf.classOf[v], uf.classOf[a]);
            MergeFollowLinks(uf, uf', uf.classOf[u], uf.classOf[v], uf.classOf[b]);
          }

          if F.ImpliedEqual(eqs', a, b) {
            // This case of the logical equivalence is even more interesting.
            // It turns out that `FollowLinks` on `uf'` makes for a suitable
            // function that obeys `eqs'` and therefore must agree on mappings
            // of `a` and `b`!  Well, we need a little extra wrapper, to deal
            // with the precondition of `FollowLinks`.  The default value of 0
            // could be anything, and the proof would still go through.
            var f := (x) => if x in uf'.classOf then FollowLinks(uf', uf'.classOf[x]) else 0;
            assert F.Compatible(f, eqs') by {
              forall c, d ensures (c, d) in eqs' ==> f(c) == f(d) {
                if (c, d) in eqs' {
                  // Now it's these innermost-quantified variables for which we
                  // need to remind Dafny of `FollowLinks` behavior after merging.
                  MergeFollowLinks(uf, uf', uf.classOf[u], uf.classOf[v], uf.classOf[c]);
                  MergeFollowLinks(uf, uf', uf.classOf[u], uf.classOf[v], uf.classOf[d]);
                }
              }
            }
          }
        }
      }
    }
  }

  // The workhorse lemma we just invoked has a fair number of preconditions, but
  // in the end it gives exactly the characterization of `FollowLinks` after
  // merging that we need.
  lemma MergeFollowLinks<K(!new)>(uf: UnionFind<K>, uf': UnionFind<K>, u: cls, v: cls, c: cls)
    requires Valid(uf) && Valid(uf')
    requires u in uf.classOf.Values && v in uf.classOf.Values
    requires uf'.merges == uf.merges[FollowLinks(uf, v) := FollowLinks(uf, u)]
    requires c in uf.classOf.Values && c in uf'.classOf.Values
    ensures FollowLinks(uf', c) == if FollowLinks(uf, c) == FollowLinks(uf, v) then FollowLinks(uf, u) else FollowLinks(uf, c)
  {
  }

  // The primary correctness lemma for `GetClassCorrect` is pretty much all routine casework.
  lemma GetClassCorrect<K(!new)>(uf: UnionFind<K>, eqs: set<(K, K)>, u: K)
    requires Valid(uf)
    requires Represents(uf, eqs)
    ensures var (uf', c) := GetClass(uf, u); Represents(uf', eqs) && FollowLinks(uf', uf'.classOf[u]) == c
  {
    if u !in uf.classOf {
      var (uf', _) := GetClass(uf, u);

      forall v1, v2: K ensures v1 in uf'.classOf && v2 in uf'.classOf ==>
        ((FollowLinks(uf', uf'.classOf[v1]) == FollowLinks(uf', uf'.classOf[v2]))
        <==> F.ImpliedEqual(eqs, v1, v2)) {
          if v1 in uf'.classOf && v2 in uf'.classOf {
            if v1 == u {
              if v2 == u {
                // This case was easy enough to automate.
              } else {
                GetClassCorrectOneOldOneNew(uf, eqs, u, v2);
              }
            } else {
              if v2 == u {
                GetClassCorrectOneOldOneNew(uf, eqs, u, v1);
              } else {
                GetClassPreservesFollowLinks(uf, u);
              }
            }
          }
        }
    }
  }

  // We separated out this case where two keys being compared include one that
  // hasn't been seen before and one that has.
  lemma GetClassCorrectOneOldOneNew<K(!new)>(uf: UnionFind<K>, eqs: set<(K, K)>, u: K, v: K)
    requires Valid(uf)
    requires Represents(uf, eqs)
    requires u !in uf.classOf
    requires v in uf.classOf
    ensures var (uf', _) := GetClass(uf, u); ((FollowLinks(uf', uf'.classOf[v]) == FollowLinks(uf', uf'.classOf[u]))
    <==> F.ImpliedEqual(eqs, v, u))
  {
    var (uf', _) := GetClass(uf, u);
    assert FollowLinks(uf', uf'.classOf[v]) != FollowLinks(uf', uf'.classOf[u]) by {
      GetClassPreservesFollowLinks(uf, u);
    }
    assert !F.ImpliedEqual(eqs, v, u) by {
      if F.ImpliedEqual(eqs, v, u) {
        var f := (x) => if x in uf.classOf then FollowLinks(uf, uf.classOf[x]) else uf.nextClass;
        assert F.Compatible(f, eqs);
      }
    }
  }

  // A simple lemma about how `GetClass` doesn't tweak behavior of `FollowLinks`
  // on old keys
  lemma GetClassPreservesFollowLinks<K(!new)>(uf: UnionFind<K>, u: K)
    requires Valid(uf)
    ensures var (uf', _) := GetClass(uf, u);
      forall c <- uf.classOf.Values :: FollowLinks(uf', c) == FollowLinks(uf, c)
  {
    var (uf', _) := GetClass(uf, u);
    forall c <- uf.classOf.Values ensures FollowLinks(uf', c) == FollowLinks(uf, c) {
      FollowLinksDependsOnlyOnMerges(uf, uf', c);
    }
  }

  // Most of the action is pushed into automatic induction here, where `c` is a
  // lemma parameter.  In general, Dafny can't do induction on a variable that
  // only appears with a quantifier in a postcondition, which motivates a common
  // pattern of breaking out a helper lemma where that variable is a parameter.
  lemma FollowLinksDependsOnlyOnMerges<K>(uf1: UnionFind<K>, uf2: UnionFind<K>, c: cls)
    requires Valid(uf1) && Valid(uf2)
    requires uf1.merges == uf2.merges
    requires c in uf1.classOf.Values && c in uf2.classOf.Values
    ensures FollowLinks(uf1, c) == FollowLinks(uf2, c)
  {
  }
}

// Let's stop here and reflect on what we just produced.
// The UnionFind module hides the details of the data structure but exposes
// public functions that are specified in terms of Valid and Represents.
// We can go on to verify other functionality on top of UnionFind, without
// knowing its internal details... which is the essence of data abstraction!

module TheoremProver {
  import F = Functions
  import UF = UnionFind

  // Let's implement a really simple automated theorem prover, which may shed
  // just a bit of light on how Dafny solves proof goals for us.

  // We think about simple formulas that are conjunctions of equalities over
  // variables.  Note that our Variable type here becomes the type parameter to
  // the union-find structure.
  type Variable = string
  datatype Exp = Var(variable: Variable)
  datatype Pred = Equal(op1: Exp, op2: Exp) | And(conj1: Pred, conj2: Pred)

  // Let's write a simple recursive interpreter to explain when predicates are
  // true, given values for the variables.
  type Env = Variable -> nat
  ghost function EvalExp(env: Env, e: Exp): nat {
    match e
    case Var(v) => env(v)
  }
  ghost predicate EvalPred(env: Env, p: Pred) {
    match p
    case Equal(op1, op2) => EvalExp(env, op1) == EvalExp(env, op2)
    case And(conj1, conj2) => EvalPred(env, conj1) && EvalPred(env, conj2)
  }

  // Now we're on to the automated prover.  The idea is that we are going to
  // prove or disprove implications between predicates, using a union-find
  // structure to track equalities known from the lefthand side of the
  // implication.

  // This function adds those equalities for a given predicate.
  function AddAssumptions(uf: UF.UnionFind<Variable>, p: Pred): UF.UnionFind<Variable>
    requires UF.Valid(uf)
    ensures UF.Valid(AddAssumptions(uf, p))
  {
    match p
    case Equal(op1, op2) => UF.AssertEqual(uf, op1.variable, op2.variable)
    case And(conj1, conj2) =>
      var uf := AddAssumptions(uf, conj1);
      AddAssumptions(uf, conj2)
  }

  // This dual function checks that every equality in a predicate is known to
  // hold.
  predicate CheckConclusions(uf: UF.UnionFind<Variable>, p: Pred)
    requires UF.Valid(uf)
  {
    match p
    case Equal(op1, op2) => UF.CheckEqual(uf, op1.variable, op2.variable)
    case And(conj1, conj2) => CheckConclusions(uf, conj1) && CheckConclusions(uf, conj2)
  }

  // To characterize correctness of those functions, we give an alternative
  // meaning to predicates as sets of equalities.
  ghost function EqualitiesOf(p: Pred): set<(Variable, Variable)> {
    match p
    case Equal(op1, op2) => {(op1.variable, op2.variable)}
    case And(conj1, conj2) => EqualitiesOf(conj1) + EqualitiesOf(conj2)
  }

  // This lemma explains how, when we build up an assumptions "database"
  // recursively, it winds up representing exactly the right equalities.
  lemma AddAssumptionsCorrect(uf: UF.UnionFind<Variable>, eqs: set<(Variable, Variable)>, p: Pred)
    requires UF.Valid(uf)
    requires UF.Represents(uf, eqs)
    ensures UF.Represents(AddAssumptions(uf, p), eqs + EqualitiesOf(p))
    decreases p
  {
    match p
    case Equal(op1, op2) => UF.AssertEqualCorrect(uf, eqs, op1.variable, op2.variable);
    case And(conj1, conj2) =>
      AddAssumptionsCorrect(uf, eqs, conj1);
      AddAssumptionsCorrect(AddAssumptions(uf, conj1), eqs + EqualitiesOf(conj1), conj2);
      assert eqs + EqualitiesOf(conj1) + EqualitiesOf(conj2)
        == eqs + (EqualitiesOf(conj1) + EqualitiesOf(conj2));
      // Note that, with this assertion, Dafny needed a little help using
      // associativity of set union.
  }

  // A simple formulation with quantifiers explains what CheckConclusions is
  // supposed to compute, and the proof can be mostly automatic.
  lemma CheckConclusionsCorrect(uf: UF.UnionFind<Variable>, eqs: set<(Variable, Variable)>, p: Pred)
    requires UF.Valid(uf)
    requires UF.Represents(uf, eqs)
    ensures CheckConclusions(uf, p)
      == (forall u, v :: (u, v) in EqualitiesOf(p) ==> F.ImpliedEqual(eqs, u, v))
  {
    match p
    case Equal(op1, op2) => UF.CheckEqualCorrect(uf, eqs, op1.variable, op2.variable);
    case _ =>
  }

  // Now here is our final theorem prover.
  predicate Implies(hyp: Pred, conc: Pred)
  {
    var uf := UF.Empty();
    var uf := AddAssumptions(uf, hyp);
    CheckConclusions(uf, conc)
  }

  // A pair of lemmas records the connection between two ways of evaluating
  // predicates.
  
  lemma EvalPredCorrect(env: Env, p: Pred)
    ensures EvalPred(env, p) == F.Compatible(env, EqualitiesOf(p))
  {
  }

  lemma EvalPredCorrectAll()
    ensures forall env, p :: EvalPred(env, p) == F.Compatible(env, EqualitiesOf(p))
  {
    forall env, p ensures EvalPred(env, p) == F.Compatible(env, EqualitiesOf(p)) {
      EvalPredCorrect(env, p);
    }
  }

  // Here is what it means semantically for an implication to hold.
  ghost predicate EvalImplication(hyp: Pred, conc: Pred) {
    forall env :: EvalPred(env, hyp) ==> EvalPred(env, conc)
  }

  // And the big theorem!
  lemma ImpliesCorrect(hyp: Pred, conc: Pred)
    ensures Implies(hyp, conc) == EvalImplication(hyp, conc)
  {
    // First, we make the quantified fact about predicate evaluation available,
    // so Dafny can use it automatically where helpful.
    EvalPredCorrectAll();

    // Then we basically march through invoking all of the preproved lemmas
    // about the different stages in the top-level prover.
    UF.EmptyCorrect<Variable>();
    var uf := UF.Empty();
    AddAssumptionsCorrect(uf, {}, hyp);
    var uf' := AddAssumptions(uf, hyp);
    // This next assertion is the one twist, where we again need to help Dafny
    // know to use a simple set-theory fact.
    assert {} + EqualitiesOf(hyp) == EqualitiesOf(hyp);
    CheckConclusionsCorrect(uf', EqualitiesOf(hyp), conc);
  }
}

// What's the big lesson from this example?
// We can develop the two modules here completely independently, so long as
// the exported parts of UnionFind stay the same!
// That means we could swap in a completely different data structure that is
// able to answer the same kind of questions about implied equalities.

// Now let's see why it was worth defining an abstraction *relation* instead of
// an abstraction *function*.  There are a few different ways we could go about
// wrapping our prior implementation.  Confession: ChatGPT did most of the work
// building these different repetitive variants.

// Here's one that computes a finite set of equalities implied by assertions so
// far, which requires bookkeeping of a set of keys used so far.
module UnionFindSetAF {
  import F = Functions
  import UF = UnionFind

  export
    reveals Closure
    provides State, Empty, AssertEqual, CheckEqual, Valid, Keys, Abstr
    provides EmptyCorrect, AssertEqualKeys, AssertEqualCorrect, CheckEqualCorrect
    provides F

  // The fields are private to clients of this export.  'assertions' is ghost
  // bookkeeping, not part of the executable union-find implementation.
  datatype State<K(==)> = State(
    impl: UF.UnionFind<K>,
    keys: set<K>,
    ghost assertions: set<(K, K)>)

  // All keys mentioned in assertions must occur in keys; conversely, every
  // tracked key has been registered by an assertion.  We do not need the
  // concrete UF's (hidden) classOf domain to establish this invariant.
  ghost predicate Valid<K(!new)>(s: State<K>) {
    UF.Valid(s.impl) && UF.Represents(s.impl, s.assertions)
    && (forall a, b :: (a, b) in s.assertions ==> a in s.keys && b in s.keys)
  }

  // Public projection: the finite universe of keys registered so far.
  // The representation and assertion history remain opaque.
  ghost function Keys<K(!new)>(s: State<K>): set<K> {
    s.keys
  }

  function Empty<K(!new)>(): State<K>
    ensures Valid(Empty<K>())
  {
    UF.EmptyCorrect<K>();
    State(UF.Empty(), {}, {})
  }

  function AssertEqual<K(!new)>(s: State<K>, u: K, v: K): State<K>
    requires Valid(s)
    ensures Valid(AssertEqual(s, u, v))
  {
    UF.AssertEqualCorrect(s.impl, s.assertions, u, v);
    State(UF.AssertEqual(s.impl, u, v),
          s.keys + {u, v}, s.assertions + {(u, v)})
  }

  predicate CheckEqual<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
  {
    UF.CheckEqual(s.impl, u, v)
  }

  // The canonical finite abstract value.  Unlike a set of assertions, this
  // is closed under the equality consequences of those assertions.
  // Note we couldn't write this definition without using `s.keys` to ensure
  // the set is finite, since Dafny `set` is only for finite sets.
  ghost function Abstr<K(!new)>(s: State<K>): set<(K, K)>
    requires Valid(s)
  {
    set a: K, b: K | a in s.keys && b in s.keys &&
      UF.CheckEqual(s.impl, a, b) :: (a, b)
  }

  // Independent, declarative specification of the abstract closure.
  ghost function Closure<K(!new)>(keys: set<K>, eqs: set<(K, K)>): set<(K, K)>
  {
    set a: K, b: K | a in keys && b in keys &&
      F.ImpliedEqual(eqs, a, b) :: (a, b)
  }

  lemma EmptyCorrect<K(!new)>()
    ensures Keys(Empty<K>()) == {}
    ensures Abstr(Empty<K>()) == Closure<K>({}, {})
  {
  }

  // The key bridge: the abstract value is exactly the finite logical closure
  // of the assertions.  The only nontrivial proof ingredient is the original
  // CheckEqualCorrect lemma.
  lemma AbstractionCorrect<K(!new)>(s: State<K>)
    requires Valid(s)
    ensures Abstr(s) == Closure(s.keys, s.assertions)
  {
    assert forall a: K, b: K ::
      (a, b) in Abstr(s) <==> (a, b) in Closure(s.keys, s.assertions) by {
      forall a: K, b: K ensures
        (a, b) in Abstr(s) <==> (a, b) in Closure(s.keys, s.assertions) {
        UF.CheckEqualCorrect(s.impl, s.assertions, a, b);
      }
    }
  }

  lemma AssertEqualKeys<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
    ensures Keys(AssertEqual(s, u, v)) == Keys(s) + {u, v}
  {
  }

  lemma CheckEqualCorrect<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
    requires u in Keys(s) && v in Keys(s)
    ensures CheckEqual(s, u, v) == ((u, v) in Abstr(s))
  {
  }

  // This is a functional-abstraction correctness theorem, expressed in terms
  // of the OLD abstract value, not the wrapper's hidden assertion history.
  lemma AssertEqualCorrect<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
    ensures Abstr(AssertEqual(s, u, v)) ==
      Closure(Keys(s) + {u, v}, Abstr(s) + {(u, v)})
  {
    UF.AssertEqualCorrect(s.impl, s.assertions, u, v);
    AbstractionCorrect(s);
    AbstractionCorrect(AssertEqual(s, u, v));
  }
}

// Another alternative is to make the abstract representation *infinite* sets of
// key pairs, but to encode sets as Boolean functions.  Now the annoyance is
// that Dafny does *not* in general say that functions mapping equal inputs to
// equal outputs are themselves equal, so we can't explain the abstraction
// function in terms of direct equalities over sets.
module UnionFindPredicateAF {
  import F = Functions
  import UF = UnionFind

  export
    reveals Merge
    provides State, Empty, AssertEqual, CheckEqual, Valid, Abstr
    provides EmptyCorrect, AssertEqualCorrect, CheckEqualCorrect
    provides F

  // The assertions are only a ghost proof aid; they are not exposed.
  datatype State<K(==)> = State(
    impl: UF.UnionFind<K>,
    ghost assertions: set<(K, K)>)

  ghost predicate Valid<K(!new)>(s: State<K>) {
    UF.Valid(s.impl) && UF.Represents(s.impl, s.assertions)
  }

  function Empty<K(!new)>(): State<K>
    ensures Valid(Empty<K>())
  {
    UF.EmptyCorrect<K>();
    State(UF.Empty(), {})
  }

  function AssertEqual<K(!new)>(s: State<K>, u: K, v: K): State<K>
    requires Valid(s)
    ensures Valid(AssertEqual(s, u, v))
  {
    UF.AssertEqualCorrect(s.impl, s.assertions, u, v);
    State(UF.AssertEqual(s.impl, u, v), s.assertions + {(u, v)})
  }

  predicate CheckEqual<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
  {
    UF.CheckEqual(s.impl, u, v)
  }

  // Characteristic function of equality, including unregistered keys.
  ghost function Abstr<K(!new)>(s: State<K>): (K, K) -> bool
    requires Valid(s)
  {
    (a: K, b: K) => UF.CheckEqual(s.impl, a, b)
  }
  // So we just package up all possible passing `CheckEqual` calls as a set!

  lemma EmptyCorrect<K(!new)>()
    ensures forall a: K, b: K :: Abstr(Empty<K>())(a, b) == (a == b)
    // Note that unpleasant part above: we can't just write what set `Abstr(...)`
    // is equal to.  We need to `forall` quantify over possible members.
  {
    forall a: K, b: K ensures Abstr(Empty<K>())(a, b) == (a == b) {
      UF.CheckEqualCorrect(Empty().impl, Empty().assertions, a, b);
      if a == b {
      } else {
        assert !F.ImpliedEqual(Empty().assertions, a, b) by {
          if F.ImpliedEqual(Empty().assertions, a, b) {
            var f := (x) => if x == a then 0 else 1;
            assert F.Compatible(f, Empty().assertions);
          }
        }
      }
    }
  }

  // Private bridge: the original relational correctness theorem establishes
  // equality with the logical consequences of the hidden assertions.
  lemma AbstractionCorrect<K(!new)>(s: State<K>)
    requires Valid(s)
    ensures forall a: K, b: K ::
      Abstr(s)(a, b) == F.ImpliedEqual(s.assertions, a, b)
  {
    forall a: K, b: K
      ensures Abstr(s)(a, b) == F.ImpliedEqual(s.assertions, a, b)
    {
      UF.CheckEqualCorrect(s.impl, s.assertions, a, b);
    }
  }

  lemma CheckEqualCorrect<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
    ensures CheckEqual(s, u, v) == Abstr(s)(u, v)
  {
  }

  // A one-edge extension of an equivalence relation.  No finite universe.
  // This formula relies on r being an equivalence relation (as Abstr is).
  // And a critic might reasonably argue that we're now incorporating more
  // knowledge of the behavior of equality than the original version needed.
  // If you agree with that critic, then add this point to the list of rough
  // edges that show why the original version was nicer.
  ghost function Merge<K(!new)>(r: (K, K) -> bool, u: K, v: K): (K, K) -> bool {
    (a: K, b: K) => r(a, b) ||
      (r(a, u) && r(v, b)) ||
      (r(a, v) && r(u, b))
  }

  lemma AssertEqualCorrect<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
    ensures forall a: K, b: K ::
      Abstr(AssertEqual(s, u, v))(a, b) == Merge(Abstr(s), u, v)(a, b)
  {
    AbstractionCorrect(s);
    AbstractionCorrect(AssertEqual(s, u, v));
    MergeClosureCorrect(s.assertions, u, v);
  }

  // These last two lemmas get technical enough that... let's not dwell on them!
  
  lemma MergeClosureCorrect<K(!new)>(eqs: set<(K, K)>, u: K, v: K)
    ensures forall a: K, b: K ::
      F.ImpliedEqual(eqs + {(u, v)}, a, b) ==
      Merge((x: K, y: K) => F.ImpliedEqual(eqs, x, y), u, v)(a, b)
  {
    forall a: K, b: K ensures
      F.ImpliedEqual(eqs + {(u, v)}, a, b) ==
      Merge((x: K, y: K) => F.ImpliedEqual(eqs, x, y), u, v)(a, b) {
        if F.ImpliedEqual(eqs + {(u, v)}, a, b) {
          if F.ImpliedEqual(eqs, a, b) {
          } else {
            MergeClosuresCorrect'(eqs, u, v, a, b);
          }
        }
      }
  }

  lemma MergeClosuresCorrect'<K(!new)>(eqs: set<(K, K)>, u: K, v: K, a: K, b: K)
    requires F.ImpliedEqual(eqs + {(u, v)}, a, b)
    requires !F.ImpliedEqual(eqs, a, b)
    ensures (F.ImpliedEqual(eqs, a, u) && F.ImpliedEqual(eqs, v, b)) || (F.ImpliedEqual(eqs, a, v) && F.ImpliedEqual(eqs, u, b))
  {
    if !(F.ImpliedEqual(eqs, a, u) && F.ImpliedEqual(eqs, v, b))
      && !(F.ImpliedEqual(eqs, a, v) &&F.ImpliedEqual(eqs, u, b)) {
        var f: K -> nat := x =>
          if F.ImpliedEqual(eqs, x, a) ||
            (F.ImpliedEqual(eqs, x, u) && F.ImpliedEqual(eqs, v, a)) ||
            (F.ImpliedEqual(eqs, x, v) && F.ImpliedEqual(eqs, u, a))
            then 1 else 0;
            
        assert F.Compatible(f, eqs + {(u, v)});
    }
  }
}

// Now, Dafny doesn't disappoint in the end, because it turns out there *is* a
// way to use an abstraction function straightforwardly: use the `iset` type
// family, for *infinite* sets.
module UnionFindIsetAF {
  import F = Functions
  import UF = UnionFind

  export
    reveals Merge
    provides State, Empty, AssertEqual, CheckEqual, Valid, Abstr
    provides EmptyCorrect, CheckEqualCorrect, AssertEqualCorrect
    provides F

  datatype State<K(==)> = State(
    impl: UF.UnionFind<K>,
    ghost assertions: set<(K, K)>)

  ghost predicate Valid<K(!new)>(s: State<K>) {
    UF.Valid(s.impl) && UF.Represents(s.impl, s.assertions)
  }

  function Empty<K(!new)>(): State<K>
    ensures Valid(Empty<K>())
  {
    UF.EmptyCorrect<K>();
    State(UF.Empty(), {})
  }

  function AssertEqual<K(!new)>(s: State<K>, u: K, v: K): State<K>
    requires Valid(s)
    ensures Valid(AssertEqual(s, u, v))
  {
    UF.AssertEqualCorrect(s.impl, s.assertions, u, v);
    State(UF.AssertEqual(s.impl, u, v), s.assertions + {(u, v)})
  }

  predicate CheckEqual<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
  {
    UF.CheckEqual(s.impl, u, v)
  }

  // The whole equality relation, including reflexivity for unseen keys.
  ghost function Abstr<K(!new)>(s: State<K>): iset<(K, K)>
    requires Valid(s)
  {
    iset a: K, b: K | UF.CheckEqual(s.impl, a, b) :: (a, b)
  }

  // A purely logical one-edge extension, written as an infinite set.
  // The same potential complaint applies here, that we are somehow hardcoding
  // our knowledge of how equality works algebraically.
  ghost function Merge<K(!new)>(r: iset<(K, K)>, u: K, v: K): iset<(K, K)> {
    iset a: K, b: K |
      (a, b) in r ||
      ((a, u) in r && (v, b) in r) ||
      ((a, v) in r && (u, b) in r) :: (a, b)
  }

  lemma EmptyCorrect<K(!new)>()
    ensures forall a: K, b: K ::
      ((a, b) in Abstr(Empty<K>())) == (a == b)
  {
    forall a: K, b: K ensures ((a, b) in Abstr(Empty<K>())) == (a == b) {
      UF.CheckEqualCorrect(Empty().impl, Empty().assertions, a, b);
      if a == b {
      } else {
        assert !F.ImpliedEqual(Empty().assertions, a, b) by {
          if F.ImpliedEqual(Empty().assertions, a, b) {
            var f := (x) => if x == a then 0 else 1;
            assert F.Compatible(f, Empty().assertions);
          }
        }
      }
    }
  }

  lemma AbstractionCorrect<K(!new)>(s: State<K>)
    requires Valid(s)
    ensures forall a: K, b: K ::
      ((a, b) in Abstr(s)) == F.ImpliedEqual(s.assertions, a, b)
  {
    forall a: K, b: K
      ensures ((a, b) in Abstr(s)) == F.ImpliedEqual(s.assertions, a, b)
    {
      UF.CheckEqualCorrect(s.impl, s.assertions, a, b);
    }
  }

  lemma CheckEqualCorrect<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
    ensures CheckEqual(s, u, v) == ((u, v) in Abstr(s))
  {
  }

  lemma AssertEqualCorrect<K(!new)>(s: State<K>, u: K, v: K)
    requires Valid(s)
    ensures Abstr(AssertEqual(s, u, v)) == Merge(Abstr(s), u, v)
  {
    AbstractionCorrect(s);
    AbstractionCorrect(AssertEqual(s, u, v));
    MergeClosureCorrect(s.assertions, u, v);
  }

  // The final lemmas are almost the same as in the prior wrapper module.
  lemma MergeClosureCorrect<K(!new)>(eqs: set<(K, K)>, u: K, v: K)
    ensures forall a: K, b: K ::
      F.ImpliedEqual(eqs + {(u, v)}, a, b) ==
      ((a, b) in Merge(
        iset x: K, y: K | F.ImpliedEqual(eqs, x, y) :: (x, y), u, v))
  {
    forall a: K, b: K ensures
      F.ImpliedEqual(eqs + {(u, v)}, a, b) ==
      ((a, b) in Merge(iset x: K, y: K | F.ImpliedEqual(eqs, x, y) :: (x, y), u, v)) {
        if F.ImpliedEqual(eqs + {(u, v)}, a, b) {
          if F.ImpliedEqual(eqs, a, b) {
          } else {
            MergeClosuresCorrect'(eqs, u, v, a, b);
          }
        }
      }
  }

  lemma MergeClosuresCorrect'<K(!new)>(eqs: set<(K, K)>, u: K, v: K, a: K, b: K)
    requires F.ImpliedEqual(eqs + {(u, v)}, a, b)
    requires !F.ImpliedEqual(eqs, a, b)
    ensures (F.ImpliedEqual(eqs, a, u) && F.ImpliedEqual(eqs, v, b)) || (F.ImpliedEqual(eqs, a, v) && F.ImpliedEqual(eqs, u, b))
  {
    if !(F.ImpliedEqual(eqs, a, u) && F.ImpliedEqual(eqs, v, b))
      && !(F.ImpliedEqual(eqs, a, v) &&F.ImpliedEqual(eqs, u, b)) {
        var f: K -> nat := x =>
          if F.ImpliedEqual(eqs, x, a) ||
            (F.ImpliedEqual(eqs, x, u) && F.ImpliedEqual(eqs, v, a)) ||
            (F.ImpliedEqual(eqs, x, v) && F.ImpliedEqual(eqs, u, a))
            then 1 else 0;
            
        assert F.Compatible(f, eqs + {(u, v)});
    }
  }
}