// In this file, you can find hints and suggestions for...
//  - proving EmptyCorrect
//  - Reaches (executable version) vs. Reachable (path-based definition)
//  - proving AssertContradiction (and EntailedCorrect)
//  - proving AssertExtended
//  - proving EntailedCorrect


// ╭──────────────────────╮
// │ proving EmptyCorrect │
// ╰──────────────────────╯

// To prove EmptyCorrect(), you want to assume there exists a node that can
// reach itself and derive a contradiction.
// Use "var path : seq<Variable> :| ..." to materialize the path that
// demonstrates reachability, then explain to Dafny the contradiction inherent
// in that path.


// ╭───────────────────────╮
// │ Reaches vs. Reachable │
// ╰───────────────────────╯

// You'll almost certainly want to prove logical equivalence of Reaches
// (executable version) and Reachable (easier to reason about based on paths).
// Since that's not the main point of the assignment, we share our proof here
// (though we think we've presented enough examples that you should also be able
// to prove this one on your own).
  lemma ReachesCorrect(kb: KnowledgeBase, src: Variable, dst: Variable)
    requires Valid(kb)
    decreases ReachTo(kb, dst)
    ensures Reaches(kb, src, dst) == Reachable(kb, src, dst)
  {
    if Reaches(kb, src, dst) {
      var p :| p in kb && p.1 == dst && (p.0 == src || (ReachToDecrease(kb, dst, p.0); Reaches(kb, src, p.0)));
      if p.0 == src {
        var path := [src, dst];
        assert Reachable(kb, src, dst);
      } else {
        ReachToDecrease(kb, dst, p.0);
        assert Reaches(kb, src, p.0);
        assert Reachable(kb, src, p.0) by {
          ReachesCorrect(kb, src, p.0);
        }
        var path : seq<Variable> :| |path| >= 2
          && path[0] == src
          && path[|path|-1] == p.0
          && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in kb);
        var path' := path + [dst];
        assert path'[|path'|-1] == dst;
        assert Reachable(kb, src, dst);
      }
    }

    if Reachable(kb, src, dst) {
      var path : seq<Variable> :| |path| >= 2
        && path[0] == src
        && path[|path|-1] == dst
        && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in kb);

      if |path| == 2 && path[1] == dst {
        assert Reaches(kb, src, dst);
      } else {
        assert (path[|path|-2], dst) in kb;
        ReachToDecrease(kb, dst, path[|path|-2]);
        ReachesCorrect(kb, src, path[|path|-2]);
        var path' := path[..|path|-1];
        assert Reachable(kb, src, path[|path|-2]);
        assert Reaches(kb, src, path[|path|-2]);
        assert Reaches(kb, src, dst);
      }
    }
  }


// ╭─────────────────────────────╮
// │ proving AssertContradiction │
// ╰─────────────────────────────╯

// To prove AssertContradiction and EntailedCorrect, we rely on a lemma to show that
// graph reachability implies less-than in any compatible variable assignment.  Here are
// the statement of the lemma and its key helper.
  lemma ReachesLt(kb: KnowledgeBase, src: Variable, dst: Variable)
    requires Valid(kb)
    requires Reaches(kb, src, dst)
    ensures forall asgn :: InterpPreds(asgn, Elements(kb)) ==> asgn(src) < asgn(dst)

  lemma ReachableLt(path: seq<Variable>, kb: KnowledgeBase, src: Variable, dst: Variable)
    requires Valid(kb)
    requires |path| >= 2
    requires path[0] == src
    requires path[|path|-1] == dst
    requires forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in kb
    ensures forall asgn :: InterpPreds(asgn, Elements(kb)) ==> asgn(src) < asgn(dst)


// ╭────────────────────────╮
// │ proving AssertExtended │
// ╰────────────────────────╯

// The proof of AssertExtended additionally requires showing that Valid is
// preserved when we add a new fact to the knowledge base.  We must show that
// any cycle in the extended graph implies a contradiction, either by finding a
// cycle in the original graph or by contradicting the reachability check that
// AssertExtended performs.  We found the following lemmas (the latter a helper
// for the former) handy.
  lemma RemoveNewEdgeFromCycle(path: seq<Variable>, kb: KnowledgeBase, a: Variable, b: Variable, x: Variable)
    requires Valid(kb)
    requires a != b
    requires !Reachable(kb, b, a)
    requires |path| >= 2
    requires path[0] == x
    requires path[|path|-1] == x
    requires forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in kb + {(a, b)}
    ensures false

  lemma RemoveNewEdgeFromCycle'(path: seq<Variable>, kb: KnowledgeBase, a: Variable, b: Variable)
    requires a != b
    requires |path| >= 2
    requires path[0] == b
    requires path[|path|-1] == a
    requires forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in kb + {(a, b)}
    ensures Reachable(kb, b, a)
    // Note the reachability conclusion in a graph smaller by one edge,
    // vs. the premises.

// Our proof of RemoveNewEdgeFromCycle begins with the following case split:
    if exists i :: 0 <= i < |path|-1 && path[i] == a && path[i+1] == b {
      // The path uses the new edge.  Get the index where it appears.
      var i :| 0 <= i < |path|-1 && path[i] == a && path[i+1] == b;
      // ...
    } else {
      // The path doesn't use the new edge.  Relatively easy case to handle....
      // ...
    }
// Our proof of RemoveNewEdgeFromCycle' begins with the same case split.


// ╭─────────────────────────╮
// │ proving EntailedCorrect │
// ╰─────────────────────────╯

// The proof of EntailedCorrect involves two directions of logical implication:
  lemma EntailedCorrect(kb: KnowledgeBase, p: Pred)
    requires Valid(kb)
    ensures Entailed(kb, p) == Implies(Elements(kb), p)
  {
    if Reaches(kb, p.a.x, p.b.x) {
      // ...
      assert Implies(Elements(kb), p);
    }

    if !Reaches(kb, p.a.x, p.b.x) && Implies(Elements(kb), p) {
      // ...
      assert false;
    }
  }

// The first direction is easy by one of the lemmas we suggested above.

// The second direction is much more interesting.  To derive a contradiction, we
// need to demonstrate a variable assignment that satisfies Elements(kb) but not
// p.  The key trick is to find the assignment by computing a *topological sort*
// of the graph.  See:
//   https://en.wikipedia.org/wiki/Topological_sorting
// (It shouldn't actually be necessary to digest the content of that Wikipedia
// page, but you might find it helpful.)  The idea is that we compute an
// ordering of the graph nodes, where the ordering respects edges.  Cyclic
// graphs don't have such orderings in general, but luckily we are enforcing
// acyclicity with Valid.
// We also need the topological sort to satisfy a particular property: as we
// have just failed to show a < b with a reachability check, we want to make
// sure that a appears later than b in the ordering.  Then the ordering
// satisfies Elements(kb) but not p, and we have our contradiction.
// Proving existence of such a topological sort requires a bit of heavy lifting,
// so we'll just give you our proof below.  You should be able to use the first
// lemma as a black box, with no need to understand its proof or the other
// helper definitions we've included (though it may be edifying to do so).
  lemma TopologicalSort(kb: KnowledgeBase, a: Variable, b: Variable) returns (ts: seq<Variable>)
    requires Valid(kb)
    requires !Reachable(kb, a, b)

    ensures forall x :: x in ts <==> x in Nodes(kb)
    // The topological sort includes exactly the graph nodes.

    ensures forall i, j :: 0 <= i < |ts| && 0 <= j < |ts| && (ts[i], ts[j]) in kb ==> i < j
    // The topological sort respects graph edges.

    ensures forall i, j :: 0 <= i < |ts| && 0 <= j < |ts| && ts[i] == b && ts[j] == a ==> j >= i
    // We succeeded in placing a no earlier than b in the topological sort.
  {
    ts := TopologicalSort'(kb, [], a, b);
  }

  ghost function TsNodes(ts: seq<Variable>): set<Variable> {
    set x | x in ts
  }

  ghost function ChooseNode(kb: KnowledgeBase, ts: seq<Variable>, x: Variable): (r: Variable)
    requires Valid(kb)
    requires x in Nodes(kb) - TsNodes(ts)
    ensures r in Nodes(kb) - TsNodes(ts)
    ensures r in ReachTo(kb, x)
    ensures forall p :: p in kb && p.1 == r ==> p.0 in ts
    decreases ReachTo(kb, x)
  {
    if forall p :: p in kb && p.1 == x ==> p.0 in ts then
      x
    else
      var p :| p in kb && p.1 == x && p.0 !in ts;
      ReachToDecrease(kb, x, p.0);
      ChooseNode(kb, ts, p.0)
  }

  lemma TopologicalSort'(kb: KnowledgeBase, ts: seq<Variable>, a: Variable, b: Variable) returns (ts': seq<Variable>)
    requires Valid(kb)
    requires !Reachable(kb, a, b)
    requires forall x :: x in ts ==> x in Nodes(kb)
    requires forall i, j :: 0 <= i < j < |ts| ==> ts[i] != ts[j]
    requires forall i, j :: 0 <= i < |ts| && 0 <= j < |ts| && (ts[i], ts[j]) in kb ==> i < j
    requires forall i, j :: 0 <= i < |ts| && 0 <= j < |ts| && ts[i] == b && ts[j] == a ==> j >= i
    requires a in ts ==> b in ts || b !in Nodes(kb)
    requires forall p :: p in kb && p.1 in ts ==> p.0 in ts
    ensures forall x :: x in ts' <==> x in Nodes(kb)
    ensures forall i, j :: 0 <= i < j < |ts'| ==> ts'[i] != ts'[j]
    ensures forall i, j :: 0 <= i < |ts'| && 0 <= j < |ts'| && (ts'[i], ts'[j]) in kb ==> i < j
    ensures forall i, j :: 0 <= i < |ts'| && 0 <= j < |ts'| && ts'[i] == b && ts'[j] == a ==> j >= i
    decreases Nodes(kb) - TsNodes(ts)
  {
    if |Nodes(kb) - TsNodes(ts)| == 0 {
      return ts;
    } else if b in ts || b !in Nodes(kb) {
      var x :| x in Nodes(kb) - TsNodes(ts);
      var y := ChooseNode(kb, ts, x);
      var ts'' := ts + [y];

      ts' := TopologicalSort'(kb, ts'', a, b);
    } else {
      var y := ChooseNode(kb, ts, b);
      var ts'' := ts + [y];

      ts' := TopologicalSort'(kb, ts'', a, b);
    }
  }

// So your proof of EntailedCorrect can call TopologicalSort to get a variable
// ordering suitable for deriving a contradiction.  But how?  The key idea is to
// turn the topological sort into a variable assignment.  We used this helper to
// produce such assignments:
  ghost function AssignmentOfTopoSort(ts: seq<Variable>): Assgn {
    (x) => if x in ts then
             (var i :| 0 <= i < |ts| && ts[i] == x; i)
           else |ts|
  }
// Note how we simply assign all variables that are not nodes to an integer
// greater than all the others we assign.

// There is one more complication: in trying to derive a contradiction around an
// inequality a < b that we confirmed doesn't hold, what if a or b aren't even
// nodes in the graph?  Then the assignment described just above may not give
// the right relative ordering.  Our proof does case splits on scenarios where
// one or both of these variables are not in Nodes(kb).  These cases deserve a
// slightly fancier assignment builder:
  ghost function AssignmentOfTopoSortWithOverride(ts: seq<Variable>, k: Variable, v: int): Assgn {
    (x) => if x in ts then
             (var i :| 0 <= i < |ts| && ts[i] == x; i)
           else if x == k then v
           else |ts|
  }
// This function is like the last, except we single out one variable k for
// special treatment with a value v.  In every scenario where we don't have both
// variables as graph nodes, you should be able to pick one of the variables to
// get special treatment and thus both stay compatible with Elements(kb) and
// also falsify a < b.