// YOUR JOB: fill in what you can to complete this proof of the theorem prover.

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

  // 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))

  // 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)
  }

  // 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)
}

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)
  }
}

// Here's an abstracted version of UnionFind.
// Your proofs up top shouldn't need to depend on anything omitted here.
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

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

  function Empty<K>(): UnionFind<K>
    ensures Valid(Empty<K>())
  predicate CheckEqual<K(!new)>(uf: UnionFind<K>, u: K, v: K)
    requires Valid(uf)
  function AssertEqual<K(!new)>(uf: UnionFind<K>, u: K, v: K): UnionFind<K>
    requires Valid(uf)
    ensures Valid(AssertEqual(uf, u, v))

  ghost predicate Represents<K(!new)>(uf: UnionFind<K>, eqs: set<(K, K)>)
    requires Valid(uf)
  lemma EmptyCorrect<K(!new)>()
    ensures Represents<K>(Empty(), {})
  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)
  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)})
}