// In this class, we'll work our way up to some first examples of verified
// loops, but we'll do it using reachability in directed graphs as a running
// example.

// A graph is set of edges, each from a source node to a destination node.
type Node = int
type Edge = (Node, Node)
type Graph = set<Edge>

// This function just computes the set of nodes of a graph (all nodes
// appearing in edges).
function Nodes(g: Graph): set<Node> {
  (set edge | edge in g :: edge.0) + (set edge | edge in g :: edge.1)
}

// The following is our central definition of reachability of one node from
// another.
ghost predicate Reachable(g: Graph, src: Node, dst: Node) {
  // There exists a path connecting the source and destination.
  exists path: seq<Node> ::
    && |path| >= 1           // The path is nonempty.
    && path[0] == src        // It starts at the right place...
    && path[|path|-1] == dst // ... and finishes at the right place.
    && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
    // And, of course, it only includes edges that actually appear in the graph.
}

// We will prove equivalence with the following directly executable version.
function Reaches(g: Graph, src: Node, dst: Node): bool
  requires src in Nodes(g)
{
  Reaches'(g, src, dst, {})
}

function Reaches'(g: Graph, src: Node, dst: Node, visited: set<Node>): bool
  requires src in Nodes(g)
  decreases Nodes(g) - visited
{
  && src !in visited
  && (src == dst || exists e :: e in g && e.0 == src && Reaches'(g, e.1, dst, visited + {src}))
}
// It might not be obvious that Reaches' is executable given its use of a
// quantifier.  However, the quantifier clearly only considers elements of a
// given finite set, so Dafny has no trouble compiling it into a loop.


// The definition of Reachable allowed duplicate nodes in paths, which we know
// is not important to characterize reachability, as duplicates represent
// detours that can be removed.  Let's start with the exercise of proving
// equivalence with the following definition.
ghost predicate ReachableNoDup(g: Graph, src: Node, dst: Node) {
  exists path: seq<Node> ::
    && |path| >= 1
    && path[0] == src
    && path[|path|-1] == dst
    && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
    && (forall i, j :: 0 <= i < j < |path| ==> path[i] != path[j])
    // This new clause enforces lack of duplicates.
}

// Let's prove the equivalence.
lemma ReachableNoDupEquiv(g: Graph, src: Node, dst: Node)
  ensures ReachableNoDup(g, src, dst) == Reachable(g, src, dst)
{
  if Reachable(g, src, dst) {
    // Dafny just needs help via the following lemma.
    ReachableRemoveDups(g, src, dst);
  }
}

// We go from simple reachability to duplicate-free reachability.
lemma ReachableRemoveDups(g: Graph, src: Node, dst: Node)
  requires Reachable(g, src, dst)
  ensures ReachableNoDup(g, src, dst)
{
  // We materialize the path that exists based on known reachability.
  var path: seq<Node> :|
    && |path| >= 1
    && path[0] == src
    && path[|path|-1] == dst
    && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g);

  // The following helper then does most of the work.
  ReachableRemoveDups'(path, g, src, dst);
}

lemma ReachableRemoveDups'(path: seq<Node>, g: Graph, src: Node, dst: Node)
  requires |path| >= 1
  requires path[0] == src
  requires path[|path|-1] == dst
  requires forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g
  ensures ReachableNoDup(g, src, dst)
{
  // Check out this cool combined version of "if" and "var x :| ...".
  if i, j :| 0 <= i < j < |path| && path[i] == path[j] {
    // If we get here, then there exists a duplicate on the path.
    // No problem: we cut out the nodes between the duplicates to make a new
    // path.
    var path' := path[..i] + path[j..];
    // Proceed through a recursive call, to remove any more duplicates.
    ReachableRemoveDups'(path', g, src, dst);
    // Luckily for termination checking, every removed duplicate shrinks the
    // path.
  }
}

// OK, another important step on the path to proving Reaches correct is to
// characterize reachability *avoiding a set of disallowed nodes*.  This set
// will canonically be those already visited in a graph traversal.
ghost predicate ReachableNoDupWithout(g: Graph, src: Node, dst: Node, visited: set<Node>) {
  exists path: seq<Node> ::
    && |path| >= 1
    && path[0] == src
    && path[|path|-1] == dst
    && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
    && (forall i, j :: 0 <= i < j < |path| ==> path[i] != path[j])
    && (forall n :: n in path ==> n !in visited)
    // This last clause is the only addition vs. ReachableNoDup.
}

// The next two lemms are quite straightforward and push the real work onto
// Reaches'Correct.

lemma ReachesCorrect(g: Graph, src: Node, dst: Node)
  requires src in Nodes(g)
  ensures Reaches(g, src, dst) == Reachable(g, src, dst)
{
  ReachableNoDupWithoutEmpty(g, src, dst);
  Reaches'Correct(g, src, dst, {});
}

lemma ReachableNoDupWithoutEmpty(g: Graph, src: Node, dst: Node)
  ensures ReachableNoDupWithout(g, src, dst, {}) == Reachable(g, src, dst)
{
  ReachableNoDupEquiv(g, src, dst);
}

lemma Reaches'Correct(g: Graph, src: Node, dst: Node, visited: set<Node>)
  requires src in Nodes(g)
  ensures Reaches'(g, src, dst, visited) == ReachableNoDupWithout(g, src, dst, visited)
  decreases Nodes(g) - visited
{
  if src in visited {
    // We ran into a cycle in the graph traversal.
    // There is no point in exploring further.
  } else if src == dst {
    // Nice, we found the node we're looking for.
    // A trivial path witnesses reachability.
    // Declaring a variable like this one gives Dafny ideas for what to try
    // filling in for existential quantifiers.
    var path := [src];
  } else if e :| e in g && e.0 == src && Reaches'(g, e.1, dst, visited + {src}) {
    // There is an edge from the source to a node that goes on to reach the
    // destination.  First, invoke the IH.
    Reaches'Correct(g, e.1, dst, visited + {src});
    // Then materialize the path that exists.
    var path: seq<Node> :|
      && |path| >= 1
      && path[0] == e.1
      && path[|path|-1] == dst
      && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
      && (forall i, j :: 0 <= i < j < |path| ==> path[i] != path[j])
      && (forall n :: n in path ==> n !in visited + {src});
    // The extended path we want simply adds the original source to the front of
    // the path.
    var path' := [src] + path;

    // Declaring the following one of the conditions on paths helps nudge Dafny
    // to use path' to instantiate the "exists" quantifier.
    assert path'[|path'|-1] == dst;
  } else if ReachableNoDupWithout(g, src, dst, visited) {
    // The last case to consider is when the destination is indeed reachable
    // without using the visited nodes.  This path explains how.
    var path: seq<Node> :|
      && |path| >= 1
      && path[0] == src
      && path[|path|-1] == dst
      && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
      && (forall i, j :: 0 <= i < j < |path| ==> path[i] != path[j])
      && (forall n :: n in path ==> n !in visited);

    // Let's show that Reaches' also returns "true" in this case.
    // Length-one paths are easy to handle.
    if |path| > 1 {
      // Otherwise, we first declare a fact that explains which edge Reaches'
      // will succeed in trying.
      assert (path[0], path[1]) in g;
      // Then, to know that the recursive call to Reaches' succeeds, we let
      // Dafny know which path shows reachability for it.
      var path' := path[1..];
    }
  }
}

// Now, let's reintroduce methods, which we played with briefly at the start of
// the semester.  These are imperative "functions" that are allowed to cause
// side effects, like via reassigning variables.

// Here is a simple recursive example of a square-root finder.  It mostly defers
// to a helper method (and we later show an alternative helper that might be
// faster to execute).
method SquareRoot(N: nat) returns (r: nat)
  ensures r * r <= N < (r + 1) * (r + 1)
  // We want to show that N is sandwiched between the right two consecutive
  // perfect squares.
{
  r := SquareRoot'(0, N);
  // r := SquareRoot'2(0, N, 1);
}

// Note how we're back to the same kind of syntax we've used in lemmas for "if"
// and so on.  Ironically, our proofs have been following more of the syntax for
// imperative code than for functional!
method SquareRoot'(n: nat, N: nat) returns (r: nat)
  requires n * n <= N
  ensures r * r <= N < (r + 1) * (r + 1)
  decreases N - n
{
  if N < (n + 1) * (n + 1) {
    return n;
  } else {
    r := SquareRoot'(n + 1, N);
  }
}

// Alternatively, we could keep around a precomputed "n plus one squared" and
// update it with a simple formula.
method SquareRoot'2(n: nat, N: nat, c: nat) returns (r: nat)
  requires n * n <= N && c == (n + 1) * (n + 1)
  ensures r * r <= N < (r + 1) * (r + 1)
  decreases N - n
{
  if N < c {
    return n;
  } else {
    r := SquareRoot'2(n + 1, N, c + 2*n + 3);
  }
}

// OK, now let's take a big jump up in method complexity.  Think of this example
// as a preview to motivate where we're headed in the next few weeks.  We won't
// expect you yet to be able to write this kind of code/proof!

// First, we create yet one more variant of reachability: it's like the
// originals except that paths start from a *set* of nodes, not a single known
// source.
ghost predicate OneReachable(g: Graph, frontier: set<Node>, dst: Node, visited: set<Node>) {
  exists n :: n in frontier && ReachableNoDupWithout(g, n, dst, visited)
}

// Now here's a nice idiomatic imperative implementation of reachability
// checking.
method CanReach(g: Graph, src: Node, dst: Node) returns (b: bool)
  requires src in Nodes(g)
  ensures b == Reachable(g, src, dst)
{
  var frontier := {src};
  // The *frontier* holds all of the nodes that we still need to check for
  // whether they have paths to dst.
  var visited := {};
  // However, we have already processed all nodes in the *visited* set, so we
  // need not consider paths that go through this set.

  // This lemma invocation helps us establish the loop invariant at the start.
  ReachableNoDupWithoutEmpty(g, src, dst);

  // Recall from other examples we blasted past that a loop *invariant* is a
  // fact that holds at the start of every loop iteration.  We also need to
  // provide a "decreases" clause to explain why the loop terminates: the
  // visited set keeps growing, but it can't grow larger than the graph's full
  // set of nodes.
  while |frontier| > 0
    invariant frontier <= Nodes(g)
    invariant visited <= Nodes(g)
    invariant frontier !! visited // I.e., these two sets don't overlap.
    invariant Reachable(g, src, dst) == OneReachable(g, frontier, dst, visited)
    // That last invariant is the tricky one: we have confirmed that the
    // original reachability problem can be reexpressed in terms of reachability
    // from the frontier, ignoring visited nodes.
    decreases Nodes(g) - visited
  {
    // Pick a frontier element to work on next.
    var next :| next in frontier;

    if next == dst {
      // We reached the destination!  A trivial path exhibits reachability.
      ghost var path := [next];
      return true;
    }

    // A lemma does most of the work in showing that one loop iteration
    // preserves the invariants.
    CanReachStep(g, frontier, dst, visited, next);
    frontier := frontier - {next};
    visited := visited + {next};
    frontier := frontier + (set e | e in g && e.0 == next && e.1 !in visited :: e.1);
    // We expand the frontier by adding all neighbors of "next" that are not
    // already visited.
  }

  return false;
}

// This lemma just restates exactly the proof challenge for invariant
// preservation above, dropping a few known facts that turn out not to be
// needed.
lemma CanReachStep(g: Graph, frontier: set<Node>, dst: Node, visited: set<Node>, next: Node)
  requires next in frontier
  requires next != dst
  requires frontier !! visited
  ensures OneReachable(g, frontier, dst, visited)
    == OneReachable(g, frontier - {next} + (set e | e in g && e.0 == next && e.1 !in visited + {next} :: e.1), dst, visited + {next})
{
  // We prove that each condition implies the other.
  
  if OneReachable(g, frontier, dst, visited) {
    // OK, if the destination is reachable via the frontier, there must be a
    // specific frontier node that has a path.
    var n0 :| n0 in frontier && ReachableNoDupWithout(g, n0, dst, visited);
    var path: seq<Node> :|
      && |path| >= 1
      && path[0] == n0
      && path[|path|-1] == dst
      && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
      && (forall i, j :: 0 <= i < j < |path| ==> path[i] != path[j])
      && (forall n :: n in path ==> n !in visited);

    if n0 == next {
      // This path happens to start at precisely the next node chosen by the
      // algorithm.  So now there will be a path from the new frontier: it's
      // this one minus the first step.  We remind Dafny that the new first step
      // really is in the new frontier.
      assert path[1] in frontier - {next} + (set e | e in g && e.0 == next && e.1 !in visited :: e.1);
      var path' := path[1..];
    } else {
      // The first step of our path from the frontier doesn't use the next node
      // chosen by the algorithm.
      if next in path {
        // The next node we chose is being added to the visited set, so we can't
        // obviously use a path that includes it.  Luckily, the node coming
        // *after* next in the path is a member of the new frontier!  We can cut
        // off the earlier path content.
        var i :| 0 <= i < |path| && path[i] == next;
        var path' := path[i+1..];
        assert path[i+1] in frontier - {next} + (set e | e in g && e.0 == next && e.1 !in visited :: e.1);
      } else {
        // If next is not on the path, we can reuse it unchanged, as the path
        // must begin at a node that was on the frontier before and *still* is.
        assert n0 in frontier - {next} + (set e | e in g && e.0 == next && e.1 !in visited :: e.1);
      }
    }
  }

  // The other direction is a little more direct to prove.
  if OneReachable(g, frontier - {next} + (set e | e in g && e.0 == next && e.1 !in visited :: e.1), dst, visited + {next}) {
    // Choose the frontier node we start from, and materialize the path from it.
    var n0 :| n0 in frontier - {next} + (set e | e in g && e.0 == next && e.1 !in visited :: e.1) && ReachableNoDupWithout(g, n0, dst, visited + {next});
    var path: seq<Node> :|
      && |path| >= 1
      && path[0] == n0
      && path[|path|-1] == dst
      && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
      && (forall i, j :: 0 <= i < j < |path| ==> path[i] != path[j])
      && (forall n :: n in path ==> n !in visited + {next});

    // When the node we start from belongs to both frontiers (before and after),
    // the proof is straightforward, but otherwise:
    if n0 !in frontier {
      // The only option left is starting from one of the new frontier nodes, so
      // materialize the specific graph edge that led us to add n0 to the
      // frontier.
      var e :| e in g && e.0 == next && e.1 !in visited && e.1 == n0;
      // That is, n0 can be reached in one step from next, so we get a valid
      // path just by adding next in front.
      var path' := [next] + path;
      assert path'[|path'|-1] == dst;
    }
  }
}

// Next, let's look at computing all pairs of nodes that reach each other in a
// graph.

// One more variation of reachability will be helpful: reachability with an
// upper bound on path length.
ghost predicate ReachableBoundedBy(g: Graph, src: Node, dst: Node, len: nat) {
  exists path: seq<Node> ::
    && 1 <= |path| <= len
    && path[0] == src
    && path[|path|-1] == dst
    && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
}

// We can reexpress original reachability using the number of nodes as a bound
// on path length.
ghost predicate ReachableBounded(g: Graph, src: Node, dst: Node) {
  ReachableBoundedBy(g, src, dst, |Nodes(g)|)
}

// Two quick lemmas prove equivalence.
lemma ReachableBoundedEquiv(g: Graph, src: Node, dst: Node)
  requires src in Nodes(g)
  ensures ReachableBounded(g, src, dst) == Reachable(g, src, dst)
{
  ReachableBoundedEquiv'(g, src, dst, |Nodes(g)|);
}

lemma ReachableBoundedEquiv'(g: Graph, src: Node, dst: Node, bound: nat)
  requires src in Nodes(g)
  requires bound >= |Nodes(g)|
  ensures ReachableBoundedBy(g, src, dst, bound) == Reachable(g, src, dst)
{
  // Invoke this old lemma to reveal that reachability can be considered via
  // only duplicate-free paths.
  if Reachable(g, src, dst) {
    ReachableNoDupEquiv(g, src, dst);
    var path: seq<Node> :|
      && |path| >= 1
      && path[0] == src
      && path[|path|-1] == dst
      && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g)
      && (forall i, j :: 0 <= i < j < |path| ==> path[i] != path[j]);

    // It's not quite obvious to Dafny that every path element is a node, as we
    // defined node sets to include all sources and targets of edges.
    forall n | n in path
      ensures n in Nodes(g)
    {
      // Every element of the path appears at some index i.
      var i :| 0 <= i < |path| && path[i] == n;
      if i == 0 {
        // The first node we actually assumed to be in the node set.
      } else if i == |path|-1 {
        // If we are looking at the last step in the path, it is in the node set
        // because it is the *target* of an edge.
        assert (path[i-1], path[i]) in g;
      } else {
        // Otherwise, this node is the *origin* of an edge.
        assert (path[i], path[i+1]) in g;
      }
    }

    // Now we can use a lemma to conclude a bound on path length.
    LengthNoDup(path, Nodes(g));
    assert ReachableBounded(g, src, dst);
  }
}

// A duplicate-free list whose elements all come from a set can't be longer than
// the set is large.
lemma LengthNoDup(path: seq<Node>, s: set<Node>)
  requires forall n :: n in path ==> n in s
  requires forall i, j :: 0 <= i < j < |path| ==> path[i] != path[j]
  ensures |path| <= |s|
{
  if |path| > 0 {
    LengthNoDup(path[1..], s - {path[0]});
  }
}

// Here is a classic graph/relation operation: considering the graph as a
// relation, compose it with itself to learn new edges.
function SelfComposition(g: Graph): Graph {
  g + (set e1, e2 | e1 in g && e2 in g && e1.1 == e2.0 :: (e1.0, e2.1))
}

// Also create all self-edges for a graph.
function ReflGraph(g: Graph): Graph {
  set n | n in Nodes(g) :: (n, n)
}

// Now here is a recursive recipe to compute many pairs that reach each other.
function ReachesUpTo(g: Graph, step: nat): Graph {
  if step == 0 then
    g + ReflGraph(g)
  else
    SelfComposition(ReachesUpTo(g, step-1))
}

// To express what lengths of paths are covered above, we use powers of 2.
function LengthFor(step: nat): nat {
  if step == 0 then
    2
  else
    2 * LengthFor(step-1) - 1
}

ghost function ReachingPairsBounded(g: Graph, bound: nat): Graph {
  set src, dst | src in Nodes(g) && dst in Nodes(g) && ReachableBoundedBy(g, src, dst, bound) :: (src, dst)
}

// This lemma characterizes exactly which paths ReachesUpTo is finding.
lemma ReachesUpToCorrect(g: Graph, step: nat)
  ensures ReachesUpTo(g, step) == ReachingPairsBounded(g, LengthFor(step))
  decreases step
{
  if step == 0 {
    ReachesUpToBase(g);
  } else {
    var g' := ReachesUpTo(g, step-1);
    ReachesUpToCorrect(g', step-1);
    ReachesUpToStep(g, g', LengthFor(step-1));
  }
}

lemma ReachesUpToBase(g: Graph)
  ensures g + ReflGraph(g) == ReachingPairsBounded(g, 2)
{
  forall src, dst | (src, dst) in g + ReflGraph(g)
    ensures (src, dst) in ReachingPairsBounded(g, 2)
  {
    if src == dst {
      var path := [src];
    } else {
      var path := [src, dst];
    }
  }
}

lemma ReachesUpToStep(g: Graph, g': Graph, bound: nat)
  requires bound > 0
  requires g' == ReachingPairsBounded(g, bound)
  ensures SelfComposition(g') == ReachingPairsBounded(g, 2 * bound - 1)
{
  forall src, dst | (src, dst) in SelfComposition(g')
    ensures src in Nodes(g) && dst in Nodes(g) && ReachableBoundedBy(g, src, dst, 2 * bound - 1)
  {
    if (src, dst) !in g' {
      var e1, e2 :| e1 in g' && e2 in g' && e1.1 == e2.0 && e1.0 == src && e2.1 == dst;
      var mid := e1.1;
      var path1: seq<Node> :|
        && 1 <= |path1| <= bound
        && path1[0] == src
        && path1[|path1|-1] == mid
        && (forall i :: 0 <= i < |path1|-1 ==> (path1[i], path1[i+1]) in g);
      var path2: seq<Node> :|
        && 1 <= |path2| <= bound
        && path2[0] == mid
        && path2[|path2|-1] == dst
        && (forall i :: 0 <= i < |path2|-1 ==> (path2[i], path2[i+1]) in g);
      var path := path1 + path2[1..];
      assert path[|path|-1] == dst;
    }
  }

  forall src, dst | src in Nodes(g) && dst in Nodes(g) && ReachableBoundedBy(g, src, dst, 2 * bound - 1)
    ensures (src, dst) in SelfComposition(g')
  {
    var path: seq<Node> :|
      && 1 <= |path| <= 2 * bound - 1
      && path[0] == src
      && path[|path|-1] == dst
      && (forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in g);
    if |path| > 1 {
      var mid := |path|/2;
      var path1 := path[..mid+1];
      var path2 := path[mid..];
      assert |path1| <= bound;
      assert ReachableBoundedBy(g, src, path[mid], bound);
      assert ReachableBoundedBy(g, path[mid], dst, bound);
      assert (path[mid-1], path[mid]) in g;
      assert path[mid] in Nodes(g);
      assert (src, dst) in SelfComposition(g');
    }
  }
}

ghost function ReachingPairs(g: Graph): Graph {
  set src, dst | src in Nodes(g) && dst in Nodes(g) && Reachable(g, src, dst) :: (src, dst)
}

method AllPairsReachability(g: Graph) returns (g': Graph)
  ensures g' == ReachingPairs(g)
{
  var longestPathLen := |Nodes(g)|;
  var allPairs := g + ReflGraph(g);
  var lenSoFar := 2;

  ReachesUpToBase(g);
  while lenSoFar < longestPathLen
    invariant allPairs == ReachingPairsBounded(g, lenSoFar)
  {
    ReachesUpToStep(g, allPairs, lenSoFar);
    lenSoFar := lenSoFar * 2 - 1;
    allPairs := SelfComposition(allPairs);
  }

  // We need to push the equivalence of ReachableBoundedEquiv' inside
  // quantifiers that appear in reasoning about the correctness of the final
  // answer.
  forall src, dst | src in Nodes(g) && dst in Nodes(g)
    ensures Reachable(g, src, dst) == ReachableBoundedBy(g, src, dst, lenSoFar)
  {
    ReachableBoundedEquiv'(g, src, dst, lenSoFar);
  }
    
  return allPairs;
}