// Let's look at an example, rather involved in its details, of verifying a
// loop-based solver for systems of linear equations.

datatype Option<T> = None | Some(value: T)

// Here is the syntax of systems of equations.  The only unusual part is that we
// build in, to syntax nodes for variables, multiplication by a constant.  Also
// note that we use Dafny's real-number type, which for our purposes is more
// accurately considered to be a type of rational numbers, where the operations
// we use can all be computed exactly.
type Var = string
datatype Exp = Const(n: real) | Var(n: real, x: Var) | Plus(e1: Exp, e2: Exp)
type Eqn = (Exp, Exp)
type Eqns = set<Eqn>

// Our goal in solving a system of equations is to find a variable assignment
// that satisfies it.  A more general solver would find a compact description of
// *all* such assignments rather than just one, but we kept the goal more modest
// to simplify verification.
type Asgn = Var -> real

// Then evaluation of different kinds of syntax w.r.t. assignments is
// straightforward.

function EvalExp(asgn: Asgn, e: Exp): real {
  match e
  case Const(n) => n
  case Var(n, x) => n * asgn(x)
  case Plus(e1, e2) => EvalExp(asgn, e1) + EvalExp(asgn, e2)
}

ghost predicate EvalEqn(asgn: Asgn, eq: Eqn) {
  EvalExp(asgn, eq.0) == EvalExp(asgn, eq.1)
}

ghost predicate EvalEqns(asgn: Asgn, eqs: Eqns) {
  forall eq :: eq in eqs ==> EvalEqn(asgn, eq)
}

// We also want to explain which variables appear in a piece of syntax.

function ExpVars(e: Exp): set<Var> {
  match e
  case Const(_) => {}
  case Var(_, x) => {x}
  case Plus(e1, e2) => ExpVars(e1) + ExpVars(e2)
}

function EqnVars(eq: Eqn): set<Var> {
  ExpVars(eq.0) + ExpVars(eq.1)
}

// Now some key helper functions over syntax, used in our solver.
// We build specifications into these definitions and don't bother to narrate
// the details, since the overall proof is quite involved (and taxes Dafny's
// ability to find the proofs without timing out!).
// First, we multiply by a constant, which manifests as multiplying similarly in
// every spot where a constant appears.

function Multiply(k: real, e: Exp): Exp
  ensures ExpVars(Multiply(k, e)) == ExpVars(e)
  ensures forall asgn :: EvalExp(asgn, Multiply(k, e)) == k * EvalExp(asgn, e)
{
  match e
  case Const(n) => Const(k * n)
  case Var(n, x) => Var(k * n, x)
  case Plus(e1, e2) => Plus(Multiply(k, e1), Multiply(k, e2))
}

// Next, we need to be able to substitute equations of the form "x = e", for "e"
// an expression not mentioning "x".  We need different operations for different
// kinds of syntax to substitute.

// This next variant operates in a system of equations, building in the removal
// of one equation before substituting.  (In intended usage, that equation will
// be one that we manipulated into the form "x = e" before substituting.)
function SubstituteEqns(eqs: Eqns, eq: Eqn, x: Var, e: Exp): Eqns {
  set eq' | eq' in eqs - {eq} :: SubstituteEqn(x, e, eq')
}

ghost function OverrideAsgn(asgn: Asgn, x: Var, v: real): Asgn {
  y => if y == x then v else asgn(y)
}

function SubstituteEqn(x: Var, e: Exp, eq: Eqn): Eqn
  ensures EqnVars(SubstituteEqn(x, e, eq)) <= EqnVars(eq) - {x} + ExpVars(e)
  ensures forall asgn :: EvalEqn(asgn, SubstituteEqn(x, e, eq)) == EvalEqn(OverrideAsgn(asgn, x, EvalExp(asgn, e)), eq)
{
  (SubstituteExp(x, e, eq.0), SubstituteExp(x, e, eq.1))
}

function SubstituteExp(x: Var, e1: Exp, e2: Exp): Exp
  ensures ExpVars(SubstituteExp(x, e1, e2)) <= ExpVars(e2) - {x} + ExpVars(e1)
  ensures forall asgn :: EvalExp(asgn, SubstituteExp(x, e1, e2)) == EvalExp(OverrideAsgn(asgn, x, EvalExp(asgn, e1)), e2)
{
  match e2
  case Const(_) => e2
  case Var(k, y) => if y == x then Multiply(k, e1) else e2
  case Plus(e21, e22) => Plus(SubstituteExp(x, e1, e21), SubstituteExp(x, e1, e22))
}

// Note that this style of substitution leaves rather large expressions, because
// it doesn't simplify by combining like terms.  This shortcut is another
// compromise to simplify verification.  Adding simplification is left as an
// exercise to the reader!  It certainly improves performance by a lot.

// Here is our overall solver.  It first *isolates* by rewriting each equation
// into some "x = e", where "x" does not appear in "e".  From that point, it is
// relatively easy to choose values, by evaluating the rewritten equations in
// reverse order.  Note that this method returns an Option type, allowing it to
// give up and return None in any situation.  We tried to write the code to
// return Some whenever there exists a solution to the equations, though we know
// one place where we cut a corner and that isn't true.  (Can you spot it?)
method SolveEqns(eqs: Eqns) returns (r: Option<Asgn>)
  ensures r.Some? ==> EvalEqns(r.value, eqs)
{
  var isolated := IsolateEqns(eqs);
  if isolated.None? {
    return None;
  }
  var asgn := ChooseValues(isolated.value, eqs);
  return Some(asgn);
}

// Here is the first main helper method.  We'll present its key loop invariant
// later.
method IsolateEqns(eqs: Eqns) returns (r: Option<seq<(Var, Exp)>>)
  ensures r.Some? ==> forall i, j :: 0 <= i < j < |r.value| ==> r.value[i].0 != r.value[j].0
  ensures r.Some? ==> forall i, j :: 0 <= i <= j < |r.value| ==> r.value[i].0 !in ExpVars(r.value[j].1)
  ensures r.Some? ==> forall asgn : Asgn :: (forall p :: p in r.value ==> asgn(p.0) == EvalExp(asgn, p.1)) ==> EvalEqns(asgn, eqs)
{
  var isolated : seq<(Var, Exp)> := [];
  // These are the equations already rewritten as "x = e".
  var not_isolated := eqs;
  // These are the equations we didn't yet rewrite.

  while |not_isolated| > 0
    invariant IsolateEqnsInvariant(eqs, isolated, not_isolated)
  {
    var eq :| eq in not_isolated;
    if x :| x in EqnVars(eq) {
      // There is a variable appearing in the current equation.
      var eo := IsolateVarEqn(x, eq);
      match eo
      case None => return None;
      case Some(e) =>
        isolated, not_isolated := IsolateEqnsBody(eqs, isolated, not_isolated, eq, x, e);
    } else {
      // Weird: an equation with no variables.  It's most likely contradictory.
      return None;
    }
  }

  return Some(isolated);
}

// A helper to that helper: isolate a given variable within an equation.
function IsolateVarEqn(x: Var, eq: Eqn): Option<Exp>
  ensures match IsolateVarEqn(x, eq)
    case None => true
    case Some(e) => (ExpVars(e) == EqnVars(eq) - {x} &&
      forall asgn : Asgn :: EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, e)))
{
  var (n1, e1) := IsolateVarExp(x, eq.0);
  var (n2, e2) := IsolateVarExp(x, eq.1);
  var n := n1 - n2;
  if n == 0.0 then
    None
    // This is one fundanental place where the solver gives up on finding a
    // solution!  To proceed, we would need to divide by zero.
  else
    var e := Plus(Multiply(1.0/n, e2), Multiply(-1.0/n, e1));
    IsolateVarEqn'(eq, x, n1, e1, n2, e2);
      
    Some(e)
}

// This helper is for isolating a variable out of an expression (one side of an
// equation).
function IsolateVarExp(x: Var, e: Exp): (real, Exp)
  ensures var (n, e') := IsolateVarExp(x, e);
    ExpVars(e') == ExpVars(e) - {x} &&
    forall asgn :: EvalExp(asgn, e) == n * asgn(x) + EvalExp(asgn, e')
{
  match e
  case Const(_) => (0.0, e)
  case Var(n, y) => if y == x then (n, Const(0.0)) else (0.0, e)
  case Plus(e1, e2) =>
    var (n1, e1') := IsolateVarExp(x, e1);
    var (n2, e2') := IsolateVarExp(x, e2);
    (n1 + n2, Plus(e1', e2'))
}

// We have to do more explicit lemma work than we're used to in Dafny, because
// ironically, though we are verifying a *linear* equation solver, the fact we
// need to prove here uses *nonlinear* arithmetic, which Dafny struggles to
// handle efficiently and predictably.
lemma IsolateVarEquation(x: real, e1: real, n1: real, e1': real, e2: real, n2: real, e2': real)
  requires e1 == n1 * x + e1'
  requires e2 == n2 * x + e2'
  requires n1 - n2 != 0.0
  ensures (e1 == e2) == (x == (1.0 / (n1 - n2)) * e2' + (-1.0 / (n1 - n2)) * e1')
{
}

lemma IsolateVarEqn'(eq: Eqn, x: Var, n1: real, e1: Exp, n2: real, e2: Exp)
  requires forall asgn :: EvalExp(asgn, eq.0) == n1 * asgn(x) + EvalExp(asgn, e1)
  requires forall asgn :: EvalExp(asgn, eq.1) == n2 * asgn(x) + EvalExp(asgn, e2)
  requires n1 - n2 != 0.0
  ensures forall asgn: Asgn :: EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, Plus(Multiply(1.0/(n1-n2), e2), Multiply(-1.0/(n1-n2), e1))))
{
  forall asgn: Asgn ensures EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, Plus(Multiply(1.0/(n1-n2), e2), Multiply(-1.0/(n1-n2), e1)))) {
    IsolateVarEquation(asgn(x), EvalExp(asgn, eq.0), n1, EvalExp(asgn, e1),
      EvalExp(asgn, eq.1), n2, EvalExp(asgn, e2));
  }
}

// Here's that central loop invariant.
ghost predicate IsolateEqnsInvariant(eqs: Eqns, isolated: seq<(Var, Exp)>, not_isolated: set<Eqn>) {
  && (forall i, j :: 0 <= i < j < |isolated| ==> isolated[i].0 != isolated[j].0)
  // The isolated equations don't assign the same variable twice.
  && (forall i, j :: 0 <= i <= j < |isolated| ==> isolated[i].0 !in ExpVars(isolated[j].1))
  // Whenever an isolated equation assigns a variable, that variable doesn't
  // appear later, including in its own equation's righthand side.
  && (forall p, eq :: p in isolated && eq in not_isolated ==> p.0 !in EqnVars(eq))
  // No isolated variable appears anymore in the equations we have yet to
  // process (we substituted it away).
  && (forall asgn : Asgn :: (forall p :: p in isolated ==> asgn(p.0) == EvalExp(asgn, p.1)) && (forall eq :: eq in not_isolated ==> EvalEqn(asgn, eq)) ==> EvalEqns(asgn, eqs))
  // Finally, tying back into the point of the solver, any solution to the
  // current equations (isolated and not) is also a solution to the original
  // equations.  (This implication could actually be an equivalence, but we
  // aren't bothering to prove that here.)
}

// We pushed most of the action of the loop body into this method.
method IsolateEqnsBody(eqs: Eqns, old_isolated: seq<(Var, Exp)>, old_not_isolated: set<Eqn>,
  eq: Eqn, x: Var, e: Exp)
  returns (isolated: seq<(Var, Exp)>, not_isolated: set<Eqn>)
  requires eq in old_not_isolated
  requires x in EqnVars(eq)
  requires ExpVars(e) == EqnVars(eq) - {x}
  requires forall asgn : Asgn :: EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, e))
  requires IsolateEqnsInvariant(eqs, old_isolated, old_not_isolated)
  ensures IsolateEqnsInvariant(eqs, isolated, not_isolated)
  ensures |not_isolated| < |old_not_isolated|
{
  isolated := old_isolated + [(x, e)];
  not_isolated := SubstituteEqns(old_not_isolated, eq, x, e);
  IsolateEqnsDecrease(eqs, old_isolated, old_not_isolated, eq, x, e);
  IsolateEqnsStep(eqs, old_isolated, old_not_isolated, eq, x, e);
}

lemma IsolateEqnsDecrease(eqs: Eqns, isolated: seq<(Var, Exp)>, not_isolated: set<Eqn>, eq: Eqn, x: Var, e: Exp)
  requires eq in not_isolated
  requires ExpVars(e) == EqnVars(eq) - {x}
  ensures |SubstituteEqns(not_isolated, eq, x, e)| < |not_isolated|
{
  SetMapDecrease(not_isolated - {eq}, x, e);
}

lemma SetMapDecrease(s: set<Eqn>, x: Var, e: Exp)
  requires x !in ExpVars(e)
  ensures |(set eq' | eq' in s :: SubstituteEqn(x, e, eq'))| <= |s|
{
  if y :| y in s {
    assert (set eq' | eq' in s :: SubstituteEqn(x, e, eq')) == (set eq' | eq' in s - {y} :: SubstituteEqn(x, e, eq')) + {SubstituteEqn(x, e, y)};
    SetMapDecrease(s - {y}, x, e);
  }
}

lemma IsolateEqnsStep(eqs: Eqns, old_isolated: seq<(Var, Exp)>, old_not_isolated: set<Eqn>, eq: Eqn, x: Var, e: Exp)
  requires eq in old_not_isolated
  requires x in EqnVars(eq)
  requires ExpVars(e) == EqnVars(eq) - {x}
  requires forall asgn : Asgn :: EvalEqn(asgn, eq) == (asgn(x) == EvalExp(asgn, e))
  requires IsolateEqnsInvariant(eqs, old_isolated, old_not_isolated)
  ensures IsolateEqnsInvariant(eqs, old_isolated + [(x, e)], SubstituteEqns(old_not_isolated, eq, x, e))
{
  var isolated := old_isolated + [(x, e)];
  var not_isolated := SubstituteEqns(old_not_isolated, eq, x, e);
  forall asgn : Asgn | (forall p :: p in isolated ==> asgn(p.0) == EvalExp(asgn, p.1)) && (forall eq :: eq in not_isolated ==> EvalEqn(asgn, eq)) ensures EvalEqns(asgn, eqs) {
    forall eq' | eq' in old_not_isolated ensures EvalEqn(asgn, eq') {
      if eq' != eq {
        assert EvalEqn(asgn, SubstituteEqn(x, e, eq'));
        assert EvalEqn(OverrideAsgn(asgn, x, EvalExp(asgn, e)), eq');
        EvalEqnExt(asgn, OverrideAsgn(asgn, x, EvalExp(asgn, e)), eq');
      }
    }
  }
}

// After that proof detour, let's look at how to build the assignment for a
// system of isolated equations.

// We'll build up an assignment as a map and then transform to an Asgn at the
// end.  Note how every variable not in the map gets a zero value.
function AsgnOf(xs: map<Var, real>): Asgn {
  x => if x in xs then xs[x] else 0.0
}

method ChooseValues(isolated: seq<(Var, Exp)>, ghost eqs: Eqns) returns (r: Asgn)
  // Note the preconditions here that follow directly from the previous loop
  // invariant.
  requires forall i, j :: 0 <= i < j < |isolated| ==> isolated[i].0 != isolated[j].0
  requires forall i, j :: 0 <= i <= j < |isolated| ==> isolated[i].0 !in ExpVars(isolated[j].1)
  requires forall asgn : Asgn :: (forall p :: p in isolated ==> asgn(p.0) == EvalExp(asgn, p.1)) ==> EvalEqns(asgn, eqs)
  ensures EvalEqns(r, eqs)
{
  var isol := isolated;
  var xs: map<Var, real> := map[];
  // The loop will keep removing the last isolated equation and using it to add
  // an entry to xs.
  
  while |isol| > 0
    invariant forall i, j :: 0 <= i < j < |isol| ==> isol[i].0 != isol[j].0
    invariant forall i, j :: 0 <= i <= j < |isol| ==> isol[i].0 !in ExpVars(isol[j].1)
    // The first two invariants remember the method precondition (though now
    // applied to a prefix of the original equation sequence).
    invariant forall p :: p in isol ==> p.0 !in xs
    // We also need to remember that any remaining equations assign variables
    // not yet covered in xs (otherwise, we could find we have contradictory
    // requirements).

    // Finally, an extension of the final method precondition: one way to know
    // an assignment satisfies the original equations is to confirm that:
    invariant forall asgn : Asgn ::
      (forall p :: p in isol ==> asgn(p.0) == EvalExp(asgn, p.1))
      // (1) It agrees with all remaining equations.
      && (forall x :: x in xs ==> asgn(x) == xs[x])
      // (2) It agrees with all mappings of xs.
      && (forall x :: x !in xs && (forall p :: p in isol ==> x != p.0) ==> asgn(x) == 0.0)
      // (3) All other variables are zeroed out.
      // These variables are underconstrained and can in fact receive any
      // values.
      ==> EvalEqns(asgn, eqs)
  {
    ghost var old_xs := xs;
    var next := isol[|isol|-1];
    xs := xs[next.0 := EvalExp(AsgnOf(xs), next.1)];
    isol := isol[..|isol|-1];

    FinalSolveStep(isol, next, old_xs, eqs);
  }

  return AsgnOf(xs);
}

ghost function VarsOf(isolated: seq<(Var, Exp)>): set<Var> {
  set p | p in isolated :: p.0
}

lemma FinalSolveStep(isolated: seq<(Var, Exp)>, next: (Var, Exp), xs: map<Var, real>, eqs: Eqns)
  requires next.0 !in xs
  requires next.0 !in ExpVars(next.1)
  requires forall p :: p in isolated ==> p.0 !in ExpVars(next.1)
  requires forall asgn : Asgn :: (forall p :: p in isolated + [next] ==> asgn(p.0) == EvalExp(asgn, p.1)) && (forall x :: x in xs ==> asgn(x) == xs[x]) && (forall x :: x !in xs.Keys + VarsOf(isolated + [next]) ==> asgn(x) == 0.0) ==> EvalEqns(asgn, eqs)
  ensures forall asgn : Asgn :: (forall p :: p in isolated ==> asgn(p.0) == EvalExp(asgn, p.1)) && (forall x :: x in xs[next.0 := EvalExp(AsgnOf(xs), next.1)] ==> asgn(x) == xs[next.0 := EvalExp(AsgnOf(xs), next.1)][x]) && (forall x :: x !in xs[next.0 := EvalExp(AsgnOf(xs), next.1)].Keys + VarsOf(isolated) ==> asgn(x) == 0.0) ==> EvalEqns(asgn, eqs)
{
  forall asgn : Asgn | (forall p :: p in isolated ==> asgn(p.0) == EvalExp(asgn, p.1)) && (forall x :: x in xs[next.0 := EvalExp(AsgnOf(xs), next.1)] ==> asgn(x) == xs[next.0 := EvalExp(AsgnOf(xs), next.1)][x]) && (forall x :: x !in xs[next.0 := EvalExp(AsgnOf(xs), next.1)].Keys + VarsOf(isolated) ==> asgn(x) == 0.0) ensures EvalEqns(asgn, eqs) {
    FinalSolveStep'(isolated, next, xs, asgn);
  }
}

lemma FinalSolveStep'(isolated: seq<(Var, Exp)>, next: (Var, Exp), xs: map<Var, real>, asgn: Asgn)
  requires next.0 !in xs
  requires next.0 !in ExpVars(next.1)
  requires forall p :: p in isolated ==> p.0 !in ExpVars(next.1)
  requires forall p :: p in isolated ==> asgn(p.0) == EvalExp(asgn, p.1)
  requires forall x :: x in xs[next.0 := EvalExp(AsgnOf(xs), next.1)] ==> asgn(x) == xs[next.0 := EvalExp(AsgnOf(xs), next.1)][x]
  requires forall x :: x !in xs.Keys + {next.0} + VarsOf(isolated) ==> asgn(x) == 0.0
  ensures forall p :: p in isolated + [next] ==> asgn(p.0) == EvalExp(asgn, p.1)
  ensures forall x :: x in xs ==> asgn(x) == xs[x]
  ensures forall x :: x !in xs.Keys + VarsOf(isolated + [next]) ==> asgn(x) == 0.0
{
  forall p | p in isolated + [next] ensures asgn(p.0) == EvalExp(asgn, p.1) {
    if p == next {
      EvalExpExt(asgn, AsgnOf(xs), next.1);
    }
  }
}

// This pair of lemmas is critical at two points above.
lemma EvalExpExt(asgn1: Asgn, asgn2: Asgn, e: Exp)
  requires forall x :: x in ExpVars(e) ==> asgn1(x) == asgn2(x)
  ensures EvalExp(asgn1, e) == EvalExp(asgn2, e)
{
}

lemma EvalEqnExt(asgn1: Asgn, asgn2: Asgn, eq: Eqn)
  requires forall x :: x in EqnVars(eq) ==> asgn1(x) == asgn2(x)
  ensures EvalEqn(asgn1, eq) == EvalEqn(asgn2, eq)
{
  EvalExpExt(asgn1, asgn2, eq.0);
  EvalExpExt(asgn1, asgn2, eq.1);
}

// Here is executable testing code.  In general, when a Dafny program runs, it
// starts with a Main() method.  To run this file, try at the command line:
//   dafny run class17_linear_solver.dfy
method Main()
{
  // A randomly generated example in 5 variables
  var eqs := {
    (Plus(Var(2.0/3.0, "x1"), Plus(Var(0.5, "x2"), Plus(Var(-1.0, "x3"), Plus(Var(1.0, "x4"), Var(0.75, "x5"))))), Const(1.0)),
    (Plus(Var(1.0, "x1"), Plus(Var(-5.0/6.0, "x2"), Plus(Var(0.4, "x3"), Plus(Var(-0.5, "x4"), Var(1.0, "x5"))))), Const(-3.0)),
    (Plus(Var(0.75, "x1"), Plus(Var(1.0, "x2"), Plus(Var(-1.0, "x3"), Plus(Var(1.0/3.0, "x4"), Var(-2.0/7.0, "x5"))))), Const(2.0)),
    (Plus(Var(0.2, "x1"), Plus(Var(-1.0, "x2"), Plus(Var(1.0, "x3"), Plus(Var(1.0, "x4"), Var(1.0/6.0, "x5"))))), Const(0.0)),
    (Plus(Var(1.0, "x1"), Plus(Var(1.0, "x2"), Plus(Var(1.0, "x3"), Plus(Var(1.0, "x4"), Var(1.0, "x5"))))), Const(4.0))
  };
  // More standard syntax for the above:
  //   (2/3)*x1 + (1/2)*x2 - x3 + x4 + (3/4)*x5 = 1
  //   x1 - (5/6)*x2 + (2/5)*x3 - (1/2)*x4 + x5 = -3
  //   (3/4)*x1 + x2 - x3 + (1/3)*x4 - (2/7)*x5 = 2
  //   (1/5)*x1 - x2 + x3 + x4 + (1/6)*x5 = 0
  //   x1 + x2 + x3 + x4 + x5 = 4
  // Thanks, ChatGPT, for the random instance in two forms!

  var asgno := SolveEqns(eqs);
  match asgno
  case None => print "No solution\n";
  case Some(asgn) =>
    print "Solution:\n";
    print "x1: ", asgn("x1"), "\n";
    print "x2: ", asgn("x2"), "\n";
    print "x3: ", asgn("x3"), "\n";
    print "x4: ", asgn("x4"), "\n";
    print "x5: ", asgn("x5"), "\n";
    // Note that we see here that Dafny's representation for rational numbers
    // comically avoids reducing to lowest terms, printing very long renderings
    // of numbers, but they do seem to be the right answers!
}