// Definition: all pairs of nodes that reach each other in a graph.
ghost function ReachingPairs(g: Graph): Graph {
  set src, dst |
    && src in Nodes(g)
    && dst in Nodes(g)
    && Reachable(g, src, dst) :: (src, dst)
}

// Executable method to compute those pairs.
// Add annotations so that Dafny proves correctness:
method AllPairsReachability(g: Graph) returns (g': Graph)
  ensures g' == ReachingPairs(g)
{
  var longestPathLen := |Nodes(g)|;
  var allPairs := g + ReflGraph(g);
  var lenSoFar := 2;

  while lenSoFar < longestPathLen
  {
    lenSoFar := lenSoFar * 2 - 1;
    allPairs := SelfComposition(allPairs);
  }

  return allPairs;
}

// === Graph definitions ===

type Node = int
type Edge = (Node, Node)
type Graph = set<Edge>

function Nodes(graph: Graph): set<Node> {
  (set edge | edge in graph :: edge.0) + (set edge | edge in graph :: edge.1)
}

// === Reachability ===

ghost predicate Reachable(graph: 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 graph)
}

// === Bounded reachability ===

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

ghost predicate ReachableBounded(g: Graph, src: Node, dst: Node) {
  ReachableBoundedBy(g, src, dst, |Nodes(g)|)
}

lemma ReachableBoundedEquiv(g: Graph, src: Node, dst: Node)
  requires src in Nodes(g)
  ensures ReachableBounded(g, src, dst) == Reachable(g, src, dst)

lemma ReachableBoundedBySizeEquiv(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)

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|

// === Reaching pairs ===

function SelfComposition(g: Graph): Graph {
  g + (set e1, e2 | e1 in g && e2 in g && e1.1 == e2.0 :: (e1.0, e2.1))
}

function ReflGraph(g: Graph): Graph {
  set n | n in Nodes(g) :: (n, n)
}

function ReachesUpTo(g: Graph, step: nat): Graph {
  if step == 0 then
    g + ReflGraph(g)
  else
    SelfComposition(ReachesUpTo(g, step-1))
}

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

lemma ReachesUpToCorrect(g: Graph, step: nat)
  ensures ReachesUpTo(g, step) == ReachingPairsBounded(g, LengthFor(step))

lemma ReachesUpToBase(g: Graph)
  ensures g + ReflGraph(g) == ReachingPairsBounded(g, 2)

lemma ReachesUpToStep(g: Graph, g': Graph, bound: nat)
  requires bound > 0
  requires g' == ReachingPairsBounded(g, bound)
  ensures SelfComposition(g') == ReachingPairsBounded(g, 2 * bound - 1)