include "pset07_definitions.dfy"
include "pset07.dfy"

module Pset7Validator {
  import Defs
  import Pset7

  lemma {:induction false} EmptyCorrect()
    ensures Pset7.Valid(Defs.Empty())
    ensures Defs.Elements(Defs.Empty()) == {} {
    Pset7.EmptyCorrect();
  }

  lemma {:induction false} MemberCorrect(x: int, s: Defs.Tree)
    requires Pset7.Valid(s)
    ensures Defs.Member(x, s) == (x in Defs.Elements(s))
  {
    Pset7.MemberCorrect(x, s);
  }

  lemma {:induction false} InsertCorrect(x: int, s: Defs.Tree)
    requires Pset7.Valid(s)
    ensures Pset7.Valid(Defs.Insert(x, s))
    ensures Defs.Elements(Defs.Insert(x, s)) == Defs.Elements(s) + {x}
  {
    Pset7.InsertCorrect(x, s);
  }

  lemma {:induction false} ValidBalanced(s: Defs.Tree)
    requires Pset7.Valid(s)
    ensures Defs.Balanced(s)
  {
    Pset7.ValidBalanced(s);
  }

  // As a bonus, here is a module that uses your proofs as a black box, to prove
  // correctness of graph reachability.  We won't bother to explain most of the
  // details, but the point is that they are independent of how you chose to
  // phrase Valid.
  module GraphReachability {
    import Defs
    import Pset7
    
    // A graph is sequence of edges, each from a source node to a destination node.
    type Graph = seq<(int, int)>

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

    // We can implement reachability by recursive search, with a set of visited
    // nodes to avoid infinite recursion.  We use red-black trees to maintain
    // those sets.
    predicate Reaches(g: Graph, src: int, dst: int)
      requires src in Nodes(g)
    {
      Pset7.EmptyCorrect();
      Reaches'(g, src, dst, Defs.Empty())
    }

    // The next two mutually recursive functions are the workhorses of the
    // algorithm.  Note that there is a nontrivial termination argument that
    // depends on the invariant and lemmas you wrote!  Specifically, we need to
    // know that:
    // (1) The visited set only grows as the recursion proceeds.
    // (2) It only ever contains nodes from the graph, which form a finite set.
    // Hence, the visited set can't grow indefinitely.
    predicate Reaches'(g: Graph, src: int, dst: int, visited: Defs.Tree)
      requires Pset7.Valid(visited)
      requires src in Nodes(g)
      decreases Nodes(g) - Defs.Elements(visited)
    {
      Pset7.MemberCorrect(src, visited);
      !Defs.Member(src, visited) && (src == dst || Reaches''(g, g, src, dst, visited))
    }

    predicate Reaches''(g: Graph, gCopy: Graph, src: int, dst: int, visited: Defs.Tree)
      requires Pset7.Valid(visited)
      requires src !in Defs.Elements(visited)
      requires src in Nodes(g)
      requires forall p :: p in gCopy ==> p.1 in Nodes(g)
      decreases Nodes(g) - Defs.Elements(visited), gCopy
    {
      |gCopy| > 0
        && (Pset7.InsertCorrect(src, visited);
           assert Nodes(g) - (Defs.Elements(visited) + {src}) == (Nodes(g) - Defs.Elements(visited)) - {src};
           (gCopy[0].0 == src && Reaches'(g, gCopy[0].1, dst, Defs.Insert(src, visited)))
        || Reaches''(g, gCopy[1..], src, dst, visited))
    }

    // This definition is important to understand what is proved overall.  It
    // characterizes graph reachability in terms of existence of duplicate-free
    // paths.
    ghost predicate Reachable(g: Graph, src: int, dst: int) {
      exists path : seq<int> ::
        |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])
    }
    // Now it's safe to skip down to lemma ReachesCorrect, which has proof
    // details before and after it.
    
    ghost predicate ReachableWithout(g: Graph, src: int, dst: int, visited: Defs.Tree) {
      exists path : seq<int> ::
        |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 Defs.Elements(visited))
    }

    ghost predicate ReachableWithout2(g0: Graph, g: Graph, src: int, dst: int, visited: Defs.Tree) {
      exists path : seq<int> ::
        |path| >= 2
        && path[0] == src
        && path[|path|-1] == dst
        && (path[0], path[1]) in g0
        && (forall i :: 1 <= 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 Defs.Elements(visited))
    }

    lemma ReachableWithoutTwo(g: Graph, src: int, dst: int, visited: Defs.Tree)
      requires src != dst
      ensures ReachableWithout2(g, g, src, dst, visited) == ReachableWithout(g, src, dst, visited)
    {
    }

    lemma ReachableWithoutRefl(g: Graph, src: int, visited: Defs.Tree)
      requires src !in Defs.Elements(visited)
      ensures ReachableWithout(g, src, src, visited)
    {
      var path := [src];
    }

    lemma ReachableWithoutExtend(g: Graph, gCopy: Graph, src: int, dst: int, visited: Defs.Tree)
      requires |gCopy| > 0
      requires gCopy[0].0 == src
      requires ReachableWithout(g, gCopy[0].1, dst, Defs.Insert(src, visited))
      requires Pset7.Valid(visited)
      requires src !in Defs.Elements(visited)
      ensures ReachableWithout2(gCopy, g, src, dst, visited)
    {
      var path : seq<int> :| |path| >= 1
        && path[0] == gCopy[0].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 Defs.Elements(Defs.Insert(src, visited)));
      Pset7.InsertCorrect(src, visited);
      var path' := [src] + path;

      assert |path'| >= 2
        && path'[0] == src
        && path'[|path'|-1] == dst
        && (path'[0], path'[1]) in gCopy
        && (forall i :: 1 <= 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 Defs.Elements(visited));
    }
    
    lemma ReachesCorrect(g: Graph, src: int, dst: int)
      requires src in Nodes(g)
      ensures Reaches(g, src, dst) == Reachable(g, src, dst)
    {
      Pset7.EmptyCorrect();
      Reaches'Correct(g, src, dst, Defs.Empty());
    }

    lemma Reaches'Correct(g: Graph, src: int, dst: int, visited: Defs.Tree)
      requires Pset7.Valid(visited)
      requires src in Nodes(g)
      ensures Reaches'(g, src, dst, visited) == ReachableWithout(g, src, dst, visited)
      decreases Nodes(g) - Defs.Elements(visited)
    {
      Pset7.MemberCorrect(src, visited);
      if Defs.Member(src, visited) {
      } else if src == dst {
        ReachableWithoutRefl(g, src, visited);
      } else {
        Reaches''Correct(g, g, src, dst, visited);
      }
    }

    lemma Reaches''Step(g: Graph, gCopy: Graph, src: int, dst: int, visited: Defs.Tree)
      requires |gCopy| > 0
      requires forall p :: p in gCopy ==> p in g
      requires !(gCopy[0].0 == src && ReachableWithout(g, gCopy[0].1, dst, Defs.Insert(src, visited)))
      requires Pset7.Valid(visited)
      ensures ReachableWithout2(gCopy[1..], g, src, dst, visited) == ReachableWithout2(gCopy, g, src, dst, visited)
    {
      if ReachableWithout2(gCopy, g, src, dst, visited) {
        
        var path : seq<int> :| |path| >= 2
          && path[0] == src
          && path[|path|-1] == dst
          && (path[0], path[1]) in gCopy
          && (forall i :: 1 <= 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 Defs.Elements(visited));

        if (path[0], path[1]) == gCopy[0] {
          if gCopy[0].0 == src {
            Pset7.InsertCorrect(src, visited);
            assert ReachableWithout(g, gCopy[0].1, dst, Defs.Insert(src, visited)) by {
              var path' := path[1..];
            }
          } else {
          }
        } else {
        }
      }
    }

    lemma Reaches''Correct(g: Graph, gCopy: Graph, src: int, dst: int, visited: Defs.Tree)
      requires Pset7.Valid(visited)
      requires src !in Defs.Elements(visited)
      requires src in Nodes(g)
      requires src != dst
      requires forall p :: p in gCopy ==> p in g
      ensures Reaches''(g, gCopy, src, dst, visited) == ReachableWithout2(gCopy, g, src, dst, visited)
      decreases Nodes(g) - Defs.Elements(visited), gCopy
    {
      if |gCopy| == 0 {
      } else {
        Pset7.InsertCorrect(src, visited);
        assert Nodes(g) - (Defs.Elements(visited) + {src}) == (Nodes(g) - Defs.Elements(visited)) - {src};
        if gCopy[0].0 == src && Reaches'(g, gCopy[0].1, dst, Defs.Insert(src, visited)) {
          Reaches'Correct(g, gCopy[0].1, dst, Defs.Insert(src, visited));
          ReachableWithoutExtend(g, gCopy, src, dst, visited);
        } else {
          Reaches'Correct(g, gCopy[0].1, dst, Defs.Insert(src, visited));
          Reaches''Step(g, gCopy, src, dst, visited);
        }
      }
    }
  }
}