// Expression evaluation and optimization

// We saw this version in the reading...
module ExprEvaluation {

  datatype Expr =
    | Const(value: int)
    | Var(name: string)
    | Plus(a: Expr, b: Expr)

  type State = map<string, int>

  function Eval(e: Expr, st: State): int
  {
    match e
    case Const(value) => value
    case Var(name) =>
      if name in st then st[name] else -42
    case Plus(a, b) =>
      Eval(a, st) + Eval(b, st)
  }

  lemma Example() {
    // x + 3 + y
    var e := Plus(Plus(Var("x"), Const(3)), Var("y"));

    var st := map["x" := 10, "y" := 20, "z" := 30];

    assert Eval(e, st) == 33;

    // Note, the testing technique in this lemma (also in Example2 and ExampleWellTypedExpr below) essentially
    // asks the verifier to compute values. If you think of it as really running the program, then you'd
    // think this technique could handle expressions of any complexity. But the verifier isn't really executing
    // the program. Instead, it tries to prove things about repeated unfolding of functions. The verifier
    // doesn't always succeed in doing this automatically. Therefore, if this technique fails to test an
    // expected property, then it might just be that the verifier didn't manage to do the whole "computation".
  }

  // We may want to avoid evaluating an expression unless all its variables
  // are defined in the state. The variables in an expression are called
  // "free variables" (to distinguish them from "bound variables" that are
  // declared inside the expression).
  // Here is a function that returns the set of (names of) free variables in
  // an expression.
  function FreeVars(e: Expr): set<string> {
    match e
    case Const(_) => {}
    case Var(name) => {name}
    case Plus(a, b) => FreeVars(a) + FreeVars(b)
  }

  // Using FreeVars, we write a precondition for Eval:
  function Eval2(e: Expr, st: State): int
    requires FreeVars(e) <= st.Keys
  {
    match e
    case Const(value) => value
    case Var(name) => st[name]
    case Plus(a, b) =>
      Eval2(a, st) + Eval2(b, st)
  }

  // Our Example still works:
  lemma Example2() {
    // x + 3 + y
    var e := Plus(Plus(Var("x"), Const(3)), Var("y"));

    var st := map["x" := 10, "y" := 20, "z" := 30];

    assert Eval2(e, st) == 33;
  }

} // module ExprEvaluation

// Let's try that again, now supporting two different data types in our interpreter.
module TypedExpressions {

  datatype Value = Bool(b: bool) | Int(x: int) | Error

  datatype Expr =
    | Const(value: Value)
    | Var(name: string)
    | Plus(a: Expr, b: Expr)
    | ITE(cond: Expr, thn: Expr, els: Expr) // `if .. then .. else ..` expressions

  type State = map<string, Value>

  function Eval(e: Expr, st: State): Value
  {
    match e
    case Const(value) => value
    case Var(name) =>
      if name in st then st[name] else Error
    case Plus(a, b) =>
      var aa, bb := Eval(a, st), Eval(b, st);
      if aa.Int? && bb.Int? then Int(aa.x + bb.x) else Error
    case ITE(cond, thn, els) =>
      match Eval(cond, st)
      case Bool(bvalue) => Eval(if bvalue then thn else els, st)
      case _ => Error
  }

  // We'd like to avoid expressions that evaluate to errors. To do that, we
  // need more information than that certain variables are defined in the state.
  // We will also need the type of the variable. To that end, we will define
  //   - types
  //   - typing environments
  //   - what it means for a state to respect a typing environment
  //   - what it means for an expression to be valid in a given typing environment

  // There are two types.
  datatype Type = TBool | TInt

  // A value is *good* for a given type if the value has that type.
  predicate GoodValue(v: Value, ty: Type) {
    || (v.Bool? && ty == TBool)
    || (v.Int? && ty == TInt)
  }

  // The Error value is not a good value for any type.
  lemma ErrorIsNotGood(ty: Type)
    ensures !GoodValue(Error, ty)
  {
  }

  // A typing environment is a map from variable names to types.
  type Environment = map<string, Type>

  // A state is *good* w.r.t. a typing environment when each for each typing "name: ty"
  // the state contains a variable "name" with a value that is suitable for type "ty".
  predicate GoodState(st: State, env: Environment) {
    forall name :: name in env ==> name in st && GoodValue(st[name], env[name])
  }

  // We consider an expression to have a given type only if all its variables are used
  // in ways consistent with their types.
  predicate HasType(e: Expr, ty: Type, env: Environment) {
    match e
    case Const(value) =>
      GoodValue(value, ty)
    case Var(name) =>
      name in env && env[name] == ty
    case Plus(a, b) =>
      HasType(a, TInt, env) && HasType(b, TInt, env) && ty == TInt
    case ITE(cond, thn, els) =>
      HasType(cond, TBool, env) && HasType(thn, ty, env) && HasType(els, ty, env)
  }

  // We can now prove that, in a good state, a well-typed expression evaluates to a value
  // that matches the expression's type. Consequently, a well-typed expression does not
  // evaluate to Error (see lemma ErrorIsNotGood).
  lemma WellTypedExpressionEvaluatesToGoodValue(e: Expr, ty: Type, st: State, env: Environment)
    requires HasType(e, ty, env) && GoodState(st, env)
    ensures GoodValue(Eval(e, st), ty)
  {
  }

  // All seems really good. Perhaps too good? What if we made some mistake in the definitions of
  // HasType or GoodState? For example, what if it's *never* possible to satisfy the precondition
  // of the WellTypedExpressionEvaluatesToGoodValue lemma? Then, the lemma would hold trivially.
  //
  // If we get worried in this way, then a quick first check is to try to prove `false` from the
  // precondition, e.g. by writing `assert false;` in the body of WellTypedExpressionEvaluatesToGoodValue.
  // If the verifier can prove `false`, then we've made a mistake somewhere. (Here, we do get an
  // error for such an `assert false;`, so at least it's not totally obvious that we've made a mistake.)
  //
  // A more detailed way to gain confidence in our specification of WellTypedExpressionEvaluatesToGoodValue
  // is to construct values for its parameters and call it.
  lemma ExampleWellTypedExpr() {
    var e := ITE(Var("b"), Const(Int(3)), Plus(Var("x"), Var("y")));
    var ty := TInt;
    var st := map["b" := Bool(false), "x" := Int(12), "y" := Int(11)];
    var env := map["b" := TBool, "x" := TInt, "y" := TInt];

    WellTypedExpressionEvaluatesToGoodValue(e, ty, st, env); // Great! The lemma's precondition check goes through.

    // So, we then know that evaluating the expression gives a value of the expected type.
    var r := Eval(e, st);
    assert GoodValue(r, ty);

    // Note, the previous assertion follows from the lemma's postcondition, which was proved above.
    // But we can also try to prove GoodValue(r, ty) directly from the definition. To do that, comment
    // out the lemma call above. Yay, that works, too.
  }


  // Let us consider some peephole optimizations on expressions. We'll start simple, by removing
  // any additions by zero.
  function NoZeroAdd(e: Expr): Expr {
    match e
    case Const(_) => e
    case Var(_) => e
    case Plus(a, b) =>
      var a', b' := NoZeroAdd(a), NoZeroAdd(b);
      var zero := Const(Int(0));
      if a' == zero then b' else if b' == zero then a' else Plus(a', b')
    case ITE(cond, thn, els) =>
      ITE(NoZeroAdd(cond), NoZeroAdd(thn), NoZeroAdd(els))
  }

  // Let's prove that NoZeroAdd does not alter the value that the expression denotes. Along the way,
  // we prove that the expression it returns is type-correct.
  lemma {:induction false} NoZeroAddCorrect(e: Expr, ty: Type, st: State, env: Environment)
    requires HasType(e, ty, env) && GoodState(st, env)
    ensures HasType(NoZeroAdd(e), ty, env)
    ensures Eval(e, st) == Eval(NoZeroAdd(e), st)
  {
    match e
    case Const(_) =>
    case Var(_) =>
    case Plus(a, b) =>
      var a', b' := NoZeroAdd(a), NoZeroAdd(b);
      assert HasType(a', TInt, env) && HasType(b', TInt, env) && Eval(e, st) == Eval(Plus(a', b'), st) by {
        NoZeroAddCorrect(a, TInt, st, env);
        NoZeroAddCorrect(b, TInt, st, env);
      }

      var aa, bb := Eval(a', st), Eval(b', st);
      assert aa.Int? && bb.Int? by {
        WellTypedExpressionEvaluatesToGoodValue(a', TInt, st, env);
        WellTypedExpressionEvaluatesToGoodValue(b', TInt, st, env);
      }
      var plus := Plus(a', b');
      assert Eval(plus, st) == Int(aa.x + bb.x);

      var zero := Const(Int(0));
      assert Eval(zero, st) == Int(0);
      if a' == zero {
        var r := b';
        assert Eval(plus, st) == Eval(b', st);
      } else if b' == zero {
        var r := a';
      } else {
        var r := Plus(a', b');
      }

    case ITE(cond, thn, els) =>
      NoZeroAddCorrect(cond, TBool, st, env);
      NoZeroAddCorrect(thn, ty, st, env);
      NoZeroAddCorrect(els, ty, st, env);

  }

  // Here is a more ambitious optimization, able to collapse arithmetic on constants.
  function PartiallyEvaluate(e: Expr): (e': Expr)
  {
    match e
    case Const(_) => e
    case Var(_) => e
    case Plus(a, b) =>
      var aa, bb := PartiallyEvaluate(a), PartiallyEvaluate(b);
      match (aa, bb) {
        case (Const(Int(avalue)), Const(Int(bvalue))) =>
          Const(Int(avalue + bvalue))
        case _ =>
          Plus(aa, bb)
      }
    case ITE(cond, thn, els) =>
      var tt, ee := PartiallyEvaluate(thn), PartiallyEvaluate(els);
      ITE(cond, tt, ee)
  }

  lemma PartiallyEvaluateCorrect(e: Expr, ty: Type, st: State, env: Environment)
    requires HasType(e, ty, env) && GoodState(st, env)
    ensures HasType(PartiallyEvaluate(e), ty, env)
    ensures Eval(e, st) == Eval(PartiallyEvaluate(e), st)
  {
    match e
    case Const(_) =>
    case Var(_) =>
    case Plus(a, b) =>
      PartiallyEvaluateCorrect(a, TInt, st, env);
      PartiallyEvaluateCorrect(b, TInt, st, env);
    case ITE(cond, thn, els) =>
  }

  // Here is a simple optimization on ITE expressions. It looks for Boolean literals being
  // used as ITE conditions.
  function OptimizeITE(e: Expr): Expr {
    match e
    case Const(_) => e
    case Var(_) => e
    case Plus(a, b) =>
      var aa, bb := OptimizeITE(a), OptimizeITE(b);
      Plus(aa, bb)
    case ITE(cond, thn, els) =>
      var cc := OptimizeITE(cond);
      match cc
      case Const(Bool(bvalue)) =>
        OptimizeITE(if bvalue then thn else els)
      case _ =>
        var tt := OptimizeITE(thn);
        var ee := OptimizeITE(els);
        ITE(cc, tt, ee)
  }

  lemma OptimizeITECorrect(e: Expr, st: State)
    ensures Eval(e, st) == Eval(OptimizeITE(e), st)
  {
    match e
    case Const(_) =>
    case Var(_) =>
    case Plus(_, _) =>
    case ITE(cond, thn, els) =>
      OptimizeITECorrect(cond, st);
  }

  // Here's a more advanced optimization for ITE expressions. It additionally
  // looks for an ITE condition being a variable. If it is, the optimization
  // uses `Bool(true)` and `Bool(false)` for that variable in the `thn` and
  // `els` branches, respectively.
  function OptimizeITE'(e: Expr, partialState: State): Expr {
    match e
    case Const(_) =>
      e
    case Var(name) =>
      if name in partialState then Const(partialState[name]) else e
    case Plus(a, b) =>
      var aa, bb := OptimizeITE'(a, partialState), OptimizeITE'(b, partialState);
      Plus(aa, bb)
    case ITE(cond, thn, els) =>
      var cc := OptimizeITE'(cond, partialState);
      match cc
      case Const(Bool(bvalue)) =>
        OptimizeITE'(if bvalue then thn else els, partialState)
      case Var(name) =>
        var tt := OptimizeITE'(thn, partialState[name := Bool(true)]);
        var ee := OptimizeITE'(els, partialState[name := Bool(false)]);
        ITE(cc, tt, ee)
      case _ =>
        var tt := OptimizeITE'(thn, partialState);
        var ee := OptimizeITE'(els, partialState);
        ITE(cc, tt, ee)
        // Q: Is the following a valid optimization?
        //     if tt == ee then tt else ITE(cc, tt, ee)
  }

  lemma OptimizeITECorrect'(e: Expr, partialState: State, st: State)
    requires st + partialState == st
    ensures Eval(e, st) == Eval(OptimizeITE'(e, partialState), st)
  {
    match e
    case Const(_) =>
    case Var(_) =>
    case Plus(_, _) =>
    case ITE(cond, thn, els) =>
      OptimizeITECorrect'(cond, partialState, st);
  }

} // module TypedExpressions

module DatabaseQueries {

  datatype List<T> = Nil | Cons(head: T, tail: List<T>)

  function Length<T>(xs: List<T>): nat {
    match xs
    case Nil => 0
    case Cons(_, tail) => 1 + Length(tail)
  }

  function Append<X>(xs: List<X>, ys: List<X>): List<X>
    ensures Length(Append(xs, ys)) == Length(xs) + Length(ys)
  {
    match xs
    case Nil => ys
    case Cons(x, tail) => Cons(x, Append(tail, ys))
  }

  function Map<X, Y>(xs: List<X>, f: X -> Y): List<Y> {
    match xs
    case Nil => Nil
    case Cons(x, tail) => Cons(f(x), Map(tail, f))
  }

  function Filter<X>(xs: List<X>, p: X -> bool): List<X> {
    match xs
    case Nil => Nil
    case Cons(x, tail) =>
      if p(x) then Cons(x, Filter(tail, p)) else Filter(tail, p)
  }

  type Name = string

  // A row in a database table is a labeled tuple. We can represent it as a map
  // from attribute names to values. For simplicity, we'll only consider integer
  // values.
  type Row = map<Name, int>

  // A table is a list of rows.
  type Table = List<Row>

  // Now, here are some operations for computing tables based on other tables.
  datatype Query =
    | Source(rows: Table)                 // A literal table of data
    | Union(q0: Query, q1: Query)         // Combining the rows of two tables
    | Project(names: set<Name>, q: Query) // Postprocessing a query by
                                          // *dropping* attributes not in the
                                          // given set
    | Select(cond: Pred, q: Query)        // Filtering query results according
                                          // to a predicate on rows.

  // For filtering queries, the condition is an expression we evaluate.
  // Let's keep them very simple for now:
  datatype Pred =
    | AttrEq(name: Name, value: int)
    | And(p0: Pred, p1: Pred)

  type State = map<Name, int>

  // How do we evaluate predicates?
  function PEval(pred: Pred, st: State): bool {
    match pred
    case AttrEq(name, value) => if name in st then st[name] == value else false
    case And(p0, p1) => PEval(p0, st) && PEval(p1, st)
  }

  // And how do we evaluate database queries?
  function QEval(q: Query): Table {
    match q
    case Source(rows) => rows
    case Union(q0, q1) => Append(QEval(q0), QEval(q1))
    case Project(names, q) => Map(QEval(q), ProjectNames(names))
    case Select(cond, q) => Filter(QEval(q), EvalCondition(cond))
  }
  // Oh look, each DB operator is defined by one of our classic list ops :)

  function ProjectNames(names: set<Name>): (Row -> Row) {
    (row: Row) => row - (row.Keys - names)
  }

  function EvalCondition(cond: Pred): (Row -> bool) {
    (row: Row) => PEval(cond, row)
  }

  // Idea: performing a Union followed by a Select is the same as
  // performing Selects (with the same condition) and then Union.

  lemma SelectUnion(p: Pred, q0: Query, q1: Query)
    ensures QEval(Select(p, Union(q0, q1))) == QEval(Union(Select(p, q0), Select(p, q1)))
  {
    // We can do some calculation...
    calc {
      QEval(Select(p, Union(q0, q1)));
      Filter(QEval(Union(q0, q1)), EvalCondition(p));
      Filter(Append(QEval(q0), QEval(q1)), EvalCondition(p));
      // ... and find the relevant underlying property to prove:
      { AppendFilters(QEval(q0), QEval(q1), EvalCondition(p)); }
      Append(Filter(QEval(q0), EvalCondition(p)), Filter(QEval(q1), EvalCondition(p)));
      Append(QEval(Select(p, q0)), QEval(Select(p, q1)));
      QEval(Union(Select(p, q0), Select(p, q1)));
    }
  }

  lemma AppendFilters<X>(l1: List<X>, l2: List<X>, p: X -> bool)
    ensures Append(Filter(l1, p), Filter(l2, p)) == Filter(Append(l1, l2), p)
  {
  }

  // We will revisit database queries in Problem Set 4.

} // module DatabaseQueries