// In this assignment, we continue the theme from class 15 and implement a
// simple automated theorem prover.  The exercise is valuable both for practice
// with data abstraction and for revealing a bit of what Dafny is doing under
// the hood on our behalf.  While class 15 presented a prover for equalities
// between variables, we now focus instead on simple inequalities.

// Please check out the included definitions before proceeding.
include "pset08_definitions.dfy"

module Pset8 {
  import opened Defs

  export
    reveals AssertResult
    provides Defs, KnowledgeBase, Valid, Elements, Empty, Assert, Entailed
    provides EmptyCorrect, AssertContradiction, AssertExtended, EntailedCorrect

  // You are actually free to modify the following definitions arbitrarily, so
  // long as our validator (pset08_validator.dfy) accepts your implementation!
  // However, you will probably find it valuable to start with our skeleton
  // code below.  We also provide further hints (spoiler alert!) in
  // pset08_hints.dfy.

  // This is the type used to represent sets of hypotheses: pairs of variables
  // known to be less than each other.  You're free to change to some more
  // clever data structure here to improve performance, though you would
  // probably then need to work harder on the proofs.
  type KnowledgeBase = set<(Variable, Variable)>

  // When is a knowledge base valid?  When it's not contradictory, that is, when
  // there exists an assignment that satisfies all the hypotheses.  It turns out
  // that property is equivalent to lack of any variable that has a path to
  // itself in the *graph* induced by treating the known less-than relationships
  // as edges, where nodes are variables.
  ghost predicate Valid(kb: KnowledgeBase) {
    forall x :: !Reachable(kb, x, x)
  }

  // We define reachability in the style from class 16, in terms of existence
  // of appropriate paths.
  ghost predicate Reachable(kb: KnowledgeBase, src: Variable, dst: Variable) {
    exists path: seq<Variable> ::
      |path| >= 2 // Important: with ">= 1", every node could reach itself!
      && path[0] == src
      && path[|path|-1] == dst
      && forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in kb
  }

  // We also need to know which predicates are represented by a knowledge base.
  ghost function Elements(kb: KnowledgeBase): Preds {
    set p | p in kb :: LT(Var(p.0), Var(p.1))
  }

  // Now on to the three key operations exported by this module.

  // First, the easy one: initializing an empty knowledge base.
  function Empty(): KnowledgeBase {
    {}
  }

  // Next, a method to add a predicate to a knowledge base.
  // The return type is a little fancy.  Two cases:
  datatype AssertResult =
    | Contradiction
    // The new predicate *contradicts* the facts we already knew.
    // There is no point in maintaining the knowledge base further, since a
    // contradictory knowledge base implies any less-than fact!
    | Extended(kb: KnowledgeBase)
    // The new predicate is consistent with what we alrady knew.
    // Here is an extended knowledge base incorporating the new info.
  function Assert(kb: KnowledgeBase, p: Pred): AssertResult
    requires Valid(kb)
  {
    // Two cases indicate a contradiction:
    // 1. We are asserting a variable is less than itself.
    // 2. We assert a < b, but the knowledge base already implies b < a.
    //    We check the latter by calling a graph-reachability checker.
    if p.a.x == p.b.x || Reaches(kb, p.b.x, p.a.x) then
      Contradiction
    else
      Extended(kb + {(p.a.x, p.b.x)})
  }

  // Finally, graph reachability also determines if a predicate is implied by a
  // knowledge base.
  predicate Entailed(kb: KnowledgeBase, p: Pred)
    requires Valid(kb)
  {
    Reaches(kb, p.a.x, p.b.x)
  }

  // Here's our implementation of reachability, which makes nifty use of a
  // quantifier to express the property succinctly.  Dafny determines that the
  // domain of the quantifier is finite, so it knows how to execute this
  // function body as executable code!
  predicate Reaches(kb: KnowledgeBase, src: Variable, dst: Variable)
    requires Valid(kb)
    decreases ReachTo(kb, dst)
  {
    exists p :: p in kb
      // There exists some edge that will form part of a path from src to dst.
      && p.1 == dst
      // Indeed, the edge is heading into dst.
      && (p.0 == src
        // Either the edge comes out of src, and we are done...
        || (ReachToDecrease(kb, dst, p.0); Reaches(kb, src, p.0)))
        // ...or we can find a path from the source to the origin of the edge.
      // Note that we prove termination using the lack of cycles in the graph, as
      // formalized by Valid, relying on a helper function ReachTo to capture
      // the idea of a set of nodes that can reach dst, which shrinks as we
      // traverse edges backwards.
  }

  ghost function ReachTo(kb: KnowledgeBase, dst: Variable): set<Variable> {
    {dst} + (set p | p in kb && Reachable(kb, p.0, dst) :: p.0)
  }

  // It will also be helpful to have a quick way to refer to the full set of
  // nodes in one of these graphs.
  ghost function Nodes(kb: KnowledgeBase): set<Variable> {
    (set p | p in kb :: p.0) + (set p | p in kb :: p.1)
  }

  // Here is a lemma we used to prove termination of Reaches.
  lemma ReachToDecrease(kb: KnowledgeBase, x: Variable, y: Variable)
    requires Valid(kb)
    requires (y, x) in kb
    ensures ReachTo(kb, y) < ReachTo(kb, x)
  {
    if x == y {
      var path := [x, x];
      assert Reachable(kb, x, x);
      assert false;
    }

    forall z ensures z in ReachTo(kb, y) ==> z in ReachTo(kb, x) {
      if z in ReachTo(kb, y) {
        if z == y {
          var path := [y, x];
          assert z in ReachTo(kb, x);
        } else {
          var path: seq<Variable> :| |path| >= 2
            && path[0] == z
            && path[|path|-1] == y
            && forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in kb;
          var path' := path + [x];
          assert path'[|path'|-1] == x;
          assert z in ReachTo(kb, x);
        }
      }
    }

    assert x !in ReachTo(kb, y) by {
      if x in ReachTo(kb, y) {
        var path: seq<Variable> :| |path| >= 2
          && path[0] == x
          && path[|path|-1] == y
          && forall i :: 0 <= i < |path|-1 ==> (path[i], path[i+1]) in kb;
        var path' := path + [x];
        assert path'[|path'|-1] == x;

        assert Reachable(kb, x, x);
        assert false;
      }
    }
    assert x in ReachTo(kb, x);
  }

  // Now, your first challenge (should you choose to keep our skeleton code
  // above unchanged): prove the natural correctness property of Empty.
  lemma EmptyCorrect()
    ensures Valid(Empty())
    ensures Elements(Empty()) == {}

  // What does it mean when Assert returns a Contradiction?
  lemma AssertContradiction(kb: KnowledgeBase, p: Pred)
    requires Valid(kb)
    requires Assert(kb, p).Contradiction?
    ensures Contradictory(Elements(kb) + {p})

  // And how about when Assert returns an Extended?
  lemma AssertExtended(kb: KnowledgeBase, p: Pred)
    requires Valid(kb)
    requires Assert(kb, p).Extended?
    ensures Valid(Assert(kb, p).kb)
    ensures Elements(Assert(kb, p).kb) == Elements(kb) + {p}

  // Finally, how do we know Entailed returns correct answers?
  lemma EntailedCorrect(kb: KnowledgeBase, p: Pred)
    requires Valid(kb)
    ensures Entailed(kb, p) == Implies(Elements(kb), p)
}