// Let's bring back this simple example method from Class 1.
method Triple(x: int) returns (r: int)
  ensures r == 3 * x
{
  var y := 2 * x;
  r := x + y;
}

// Here is another method that calls Triple.
// Note that Dafny is able to prove the requested property about the result of
// the method call.
method Caller() {
  var t := Triple(18);
  assert t < 100;
}

// To help explain how Dafny reasons about method calls, let's show how to model
// the method call with more primitive commands.
method Caller_model() {
  // First, we make up fresh versions of the variable names involved in Triple's
  // specification.  That way, we avoid getting confused across different calls
  // to the same method (or just multiple methods that share parameter/return
  // names).

  // So, first, we assign to a fresh variable storing the parameter to the
  // method.
  var x' := 18;

  // Then we create another variable for the (in-principle) unknown return
  // value.
  var r' : int;

  // Here's the trippy part: we *assume* the postcondition of the method.
  // That means we add an extra true fact without proving it!
  // This modeling trick is *only* sound given that we really did verify the
  // method against the claimed postcondition.  IDEs highlight 'assume' in a
  // scary way, so you think twice before using it!
  assume r' == 3 * x';

  // Finally, we assign the return value to the original local variable where it
  // was intended to end up.
  var t := r';

  // Now the assertion can be proved straightforwardly.
  assert t < 100;
}

// Let's repeat the kind of exercise from class 2 and show how we calculate
// what Dafny knows at each line of code, first in the forward direction.
method Caller_forward() {
  var x' := 18;
  assert x' == 18;

  var r' : int;
  assert x' == 18;

  assume r' == 3 * x';
  assert x' == 18 && r' == 3 * x';
  // Yup, the fact that we assume just magically pops into the formula!

  var t := r';
  assert x' == 18 && r' == 3 * x' && t == r';

  assert t < 100;
}

// We can also calculate known facts backward.
// Like last time, read the assertions here from bottom to top as we calculate.
method Caller_backward() {
  // This final fact is straightforward to prove by arithmetic.
  assert forall r' :: r' == 3 * 18 ==> r' < 100;
  var x' := 18;

  // We glossed over this point last time: we can calculate backward past a
  // variable using a "for all" quantifier.
  assert forall r' :: r' == 3 * x' ==> r' < 100;
  var r' : int;

  // Aha: backwards reasoning treats 'assume' by adding an implication!
  assert r' == 3 * x' ==> r' < 100;
  assume r' == 3 * x';

  assert r' < 100;
  var t := r';

  assert t < 100;
}

// Here's another method for us to use in working through examples.
method Divide(a: int, b: int) returns (r: int)
  requires b != 0
  ensures r == a / b
{
  return a / b;
}

// A simple call that manages to establish Divide's precondition:
method DivideCaller(x: int)
{
  var t := Divide(100, 2 * x + 1);
  assert t <= 100;
}

// How do we extend our modeling of method calls via more primitive constructs?
// This time, we already have the required tool in our arsenal: 'assert'.
method DivideCaller_model(x: int)
{
  var a', b' := 100, 2 * x + 1;
  assert b' != 0;
  var r': int;
  assume r' == a' / b';
  var t := r';
  assert t <= 100;
}

method DivideCaller_forward(x: int)
{
  var a', b' := 100, 2 * x + 1;
  assert a' == 100 && b' == 2 * x + 1;

  // An 'assert' doesn't actually change what we know; it just triggers a check
  // that a particular fact is provable, before we proceed.
  assert b' != 0;
  assert a' == 100 && b' == 2 * x + 1;

  var r': int;
  assume r' == a' / b';
  assert a' == 100 && b' == 2 * x + 1 && r' == a' / b';

  var t := r';
  assert a' == 100 && b' == 2 * x + 1 && r' == a' / b' && t == r';

  assert t <= 100;
}

method DivideCaller_backward(x: int)
{
  assert (forall r' :: r' == 100 / (2 * x + 1) ==> r' <= 100) && 2 * x + 1 != 0;
  var a', b' := 100, 2 * x + 1;

  assert (forall r' :: r' == a' / b' ==> r' <= 100) && b' != 0;
  assert b' != 0;

  assert forall r' :: r' == a' / b' ==> r' <= 100;
  var r': int;

  assert r' == a' / b' ==> r' <= 100;
  assume r' == a' / b';

  assert r' <= 100;
  var t := r';

  assert t <= 100;
}

// Now that we understand how Danfy calculates through method calls, we're in a
// good position to verify some recursive functions.  This one checks that a
// sequence is sorted.
method Sorted(s: seq<int>) returns (r: bool)
  ensures r == (forall i :: 0 <= i < |s|-1 ==> s[i] <= s[i+1])
{
  if |s| <= 1 {
    return true;
  } else {
    var t := Sorted(s[1..]);
    return s[0] <= s[1] && t;
  }
}
// That wasn't so bad at all!  Beyond normal programming, we just needed to
// explain what correctness means for this function.
// We aren't always so lucky, though.

// Maybe we want to generate a sequence of consecutive integers, sort of like
// Python range().  Note that the postcondition only records that the final
// sequence is sorted.
method Upto(n: int) returns (r: seq<int>)
  requires n >= 0
  ensures forall i :: 0 <= i < |r|-1 ==> r[i] <= r[i+1]
{
  if n == 0 {
    return [];
  } else {
    var t := Upto(n-1);
    return t + [n-1];
  }
}
// Yikes: Dafny says it can't prove the postcondition for the 'else' branch!
// To see why, let's rewrite Upto using the modeling we just practiced.

method Upto_model(n: int) returns (r: seq<int>)
  requires n >= 0
  ensures forall i :: 0 <= i < |r|-1 ==> r[i] <= r[i+1]
{
  if n == 0 {
    return [];
  } else {
    // Now we can see very clearly why it is important to use fresh renamings of
    // the variables of the method we're calling.  In a recursive call, we would
    // confusingly overwrite the original variable values, without this strategy.
    var n' := n - 1;
    assert n' >= 0;
    var r' : seq<int>;
    assume forall i :: 0 <= i < |r'|-1 ==> r'[i] <= r'[i+1];
    var t := r';
    return t + [n-1];
  }
}

// Let's do a forward calculation to figure out why Dafny can't prove the
// postcondition.
method Upto_forward(n: int) returns (r: seq<int>)
  requires n >= 0
  ensures forall i :: 0 <= i < |r|-1 ==> r[i] <= r[i+1]
{
  if n == 0 {
    return [];
  } else {
    assert n > 0;

    var n' := n - 1;
    assert n > 0 && n' == n - 1;

    assert n' >= 0;
    var r' : seq<int>;
    assume forall i :: 0 <= i < |r'|-1 ==> r'[i] <= r'[i+1];
    assert n > 0 && n' == n - 1 && (forall i :: 0 <= i < |r'|-1 ==> r'[i] <= r'[i+1]);

    var t := r';
    assert n > 0 && n' == n - 1 && (forall i :: 0 <= i < |r'|-1 ==> r'[i] <= r'[i+1]) && t == r';

    // The postcondition is asking us to prove the following:
    assert forall i :: 0 <= i < |t + [n-1]|-1 ==> (t + [n-1])[i] <= (t + [n-1])[i+1];
    // But it genuinely isn't provable here!  Why?
    // From what we know about the recursive call, we can't be sure that n-1 is
    // really greater than or equal to all elements of t!
    return t + [n-1];
  }
}

// Here is one way to fix the problem, by choosing a *stronger* postcondition,
// just like we often strengthen induction hypotheses (IHes) in math proofs.
method Upto_stronger(n: int) returns (r: seq<int>)
  requires n >= 0
  ensures |r| == n
  ensures forall i :: 0 <= i < n-1 ==> r[i] <= r[i+1]
  ensures forall v :: v in r ==> 0 <= v < n
{
  if n == 0 {
    return [];
  } else {
    var t := Upto_stronger(n-1);
    // assert n > 1 ==> t[n-2] in t;
    return t + [n-1];
  }
}

// Whoa, why did that one assertion make the difference?
// Let's show more of the proof-debugging action of repeatedly asserting facts
// to figure out what Dafny is having trouble proving.
// Read this one starting from the bottom.
method Upto_stronger_missing(n: int) returns (r: seq<int>)
  requires n >= 0
  ensures |r| == n
  ensures forall i :: 0 <= i < n-1 ==> r[i] <= r[i+1]
  ensures forall v :: v in r ==> 0 <= v < n
{
  if n == 0 {
    return [];
  } else {
    var t := Upto_stronger_missing(n-1);

    // Gosh, what could Dafny be missing to conclude that fact?
    // Oh, it apparently isn't using the last postcondition of the recursive call.
    // Let's remind Dafny that t[n-2] actually is part of the result of that call.
    assert n > 1 ==> t[n-2] in t;
    // That's what we were missing!  And now all the other extra lines we added
    // are redundant, or at least Dafny figures them out itself as needed.

    // OK, what else might Dafny not know?
    // Let's restate the obligation for this case, but using the equations we
    // just checked.
    assert n > 1 ==> t[n-2] <= n-1;
    // Oh, it doesn't know this one!

    // How about the other one?
    assert n > 1 ==> (t + [n-1])[n-1] == n-1;
    // This one is fine, too.

    // Hm, why can't Dafny prove that nice, simple, quantifier-free implication?
    // Let's check if it knows how to simplify each of the sequence accesses.
    assert n > 1 ==> (t + [n-1])[n-2] == t[n-2];
    // This one is fine.

    // Sure enough, Dafny says it can't prove this version.
    // (We did a little algebraic simplification, e.g. eliminating 'i'.)
    assert n > 1 ==> (t + [n-1])[n-2] <= (t + [n-1])[n-1];

    // OK, let's check if Dafny knows how to prove the first of the two cases.
    assert forall i :: 0 <= i < n-2 ==> (t + [n-1])[i] <= (t + [n-1])[i+1];
    // Yup, that goes through.  So we can focus on the other case.

    // Insight #1: let's split the premise of the implication into two cases:
    // (1) 'i' is in-bounds for 't'.
    // (2) 'i' is such that we compare the last element of 't' with the new element.
    //     Note this case is only possible if n > 1!  Otherwise, the final array has
    //     length 1, and we don't actually compare any numbers.
    assert forall i :: 0 <= i < n-2 || (n > 1 && i == n-2) ==> (t + [n-1])[i] <= (t + [n-1])[i+1];

    // When Dafny tells us it can't prove the postcondition, we can copy that
    // postcondition here, with current variable values substituted.
    assert forall i :: 0 <= i < n-1 ==> (t + [n-1])[i] <= (t + [n-1])[i+1];
    return t + [n-1];
  }
}

// Another approach is to give the most precise specification we can, which
// leaves us with less guesswork about which properties are worth remembering.
// This style is the ultimate in "IH strengthening"!
method Upto_determined_helper(n: int) returns (r: seq<int>)
  requires n >= 0
  ensures |r| == n
  ensures forall i :: 0 <= i < n ==> r[i] == i
{
  if n == 0 {
    return [];
  } else {
    var t := Upto_determined_helper(n-1);
    return t + [n-1];
  }
}

method Upto'(n: int) returns (r: seq<int>)
  requires n >= 0
  ensures forall i :: 0 <= i < |r|-1 ==> r[i] <= r[i+1]
{
  var t := Upto_determined_helper(n);
  return t;
}
// It's kind of a shame we had to introduce an extra call just to prove this
// implementation correct!  We can adopt a hybrid approach, though.

// We *can* still do it in one method.
method Upto_determined(n: int) returns (r: seq<int>)
  requires n >= 0
  ensures |r| == n
  ensures forall i :: 0 <= i < n ==> r[i] == i
  ensures forall i :: 0 <= i < |r|-1 ==> r[i] <= r[i+1]
{
  if n == 0 {
    return [];
  } else {
    var t := Upto_determined(n-1);
    return t + [n-1];
  }
}
// The specification is still "cluttered" with other facts that we might prefer
// to think of as internal proof details, so there is a definite trade-off to
// make between these two styles.

// As another example, let's write an efficient check for whether every element
// of one sequence is also in another one.  Just to mix things up, we'll work
// with strings instead of integers as data values.

// To start with, here is a version that is *not* efficient.  It takes time
// (at least) quadratic in the sequence lengths, because sequence 'in' takes
// linear time!
method SequenceInclusion_slow(a: seq<string>, b: seq<string>) returns (r: bool)
  ensures r == (forall x :: x in a ==> x in b)
{
  if |a| == 0 {
    return true;
  } else {
    var t := SequenceInclusion_slow(a[1..], b);
    return t && a[0] in b;
  }
}

// To do better, we'll write (and verify) a few helper methods.
// This first one may not quite take linear time due to our repeated
// creation of new sets, but a smart-enough compiler optimization could figure
// out the pattern and modify the set instead of making a new one, right? >:-)
method SetOfSeq(a: seq<string>) returns (r: set<string>)
  ensures forall x :: x in a <==> x in r
{
  if |a| == 0 {
    return {};
  } else {
    var t := SetOfSeq(a[1..]);
    return t + {a[0]};
  }
}

// Now we can implement the original functionality but where the second argument
// is a set.
method SequenceInclusionWithSet(a: seq<string>, b: set<string>) returns (r: bool)
  ensures r == (forall x :: x in a ==> x in b)
{
  if |a| == 0 {
    return true;
  } else {
    var t := SequenceInclusionWithSet(a[1..], b);
    return t && a[0] in b;
  }
}

// And we bring it all together here.
method SequenceInclusion(a: seq<string>, b: seq<string>) returns (r: bool)
  ensures r == (forall x :: x in a ==> x in b)
{
  var b' := SetOfSeq(b);
  var t := SequenceInclusionWithSet(a, b');
  return t;
}

// We can also get roughly the same performance with a version that leans more
// into set operations.
method SequenceInclusion_alt(a: seq<string>, b: seq<string>) returns (r: bool)
  ensures r == (forall x :: x in a ==> x in b)
{
  var a' := SetOfSeq(a);
  var b' := SetOfSeq(b);
  return a' <= b'; // That's a subset check!
}

// One last recursion example: counting how many times each element of a
// sequence occurs. Recall multisets from classes 1 & 2. They are like sets
// but with duplicates allowed -- perfect for expressing what it means to
// count occurrences.
method Counts(s: seq<string>) returns (r: map<string, int>)
  ensures forall k :: k in r.Keys <==> k in s
  ensures forall k :: k in r.Keys ==> r[k] == multiset(s)[k]
{
  if |s| == 0 {
    return map[];
  } else {
    var t := Counts(s[1..]);

    // Note the need for a case split to avoid accessing the map at a key that
    // it doesn't contain.
    if s[0] in t {
      assert s == [s[0]] + s[1..];
      // This extra assertion took a bit of experimentation to concoct!
      return t[s[0] := t[s[0]] + 1];
    } else {
      return t[s[0] := 1];
    }
  }
}