datatype Expr = Const(nat) | Var(string) | Node(op: Op, args: List<Expr>)

datatype Op = Add | Mul

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

function Eval(e: Expr, env: map<string, nat>): nat 
    requires GoodEnv(e, env)
{
    match e
    case Const(c) => c
    case Var(s) => env[s]
    case Node(op, args) => EvalList(op, args, env)
}

/* E.4.9 => The default decreases clause use lexicographical ordering of the arguments.
This implies the following decreases clause:
    decreases op, args, env

But as "op" is either Add or Mul which and does not have fixed order in a list<Expr>
it jumps around.
for example 10 * 5 + 3 * 10. 
=> Node(Mul, Cons(Const(10), Cons()))

*/

function EvalList(op: Op, args: List<Expr>,
env: map<string, nat>): nat
decreases args
requires GoodEnvList(args, env)
{
match args
    case Nil =>
    (match op case Add => 0 case Mul => 1)
    case Cons(e, tail) =>
        var v0, v1 := Eval(e, env), EvalList(op, tail, env);
        match op
            case Add => v0 + v1
            case Mul => v0 * v1
}

predicate GoodEnv(e: Expr, env: map<string, nat>) {
    match e
    case Const(_) => true
    case Var(s) => s in env.Keys
    case Node(_, args) => GoodEnvList(args, env)
}

predicate GoodEnvList(args: List<Expr>, env: map<string, nat>)
{
    match args
    case Nil => true
    case Cons(e, tail) => GoodEnv(e, env) && GoodEnvList(tail, env)
}


ghost method TestGoodEnvExtensive() {

    // ---------- Simple constants ----------
    var c1 := Const(5);
    assert GoodEnv(c1, map[]);

    // ---------- Single variable ----------
    var v1 := Var("x");
    var env1 := map["x" := 10];
    assert GoodEnv(v1, env1);

    var env2 := map["y" := 10];
    assert !GoodEnv(v1, env2);

    // ---------- Simple Node ----------
    var e1 := Node(Add, Cons(Const(2), Cons(Const(3), Nil)));
    assert GoodEnv(e1, map[]);

    // ---------- Node with variable ----------
    var e2 := Node(Add, Cons(Var("x"), Cons(Const(4), Nil)));
    var env3 := map["x" := 7];
    assert GoodEnv(e2, env3);

    var env4 := map["y" := 7];
    assert !GoodEnv(e2, env4);

    // ---------- Multiple variables ----------
    var e3 :=
        Node(Mul,
            Cons(Var("x"),
            Cons(Var("y"),
            Cons(Const(5), Nil))));

    var env5 := map["x" := 1, "y" := 2];
    assert GoodEnv(e3, env5);

    var env6 := map["x" := 1];
    assert !GoodEnv(e3, env6);

    // ---------- Nested nodes ----------
    var nested :=
        Node(Add,
            Cons(
                Node(Mul,
                    Cons(Var("x"),
                    Cons(Const(3), Nil))),
            Cons(
                Node(Add,
                    Cons(Const(1),
                    Cons(Var("y"), Nil))),
            Nil)));

    var env7 := map["x" := 2, "y" := 4];
    assert GoodEnv(nested, env7);

    var env8 := map["x" := 2];
    assert !GoodEnv(nested, env8);

    // ---------- Deep tree ----------
    var deep :=
        Node(Add,
            Cons(
                Node(Add,
                    Cons(
                        Node(Mul,
                            Cons(Var("a"),
                            Cons(Const(2), Nil))),
                    Cons(Const(3), Nil))),
            Cons(
                Node(Mul,
                    Cons(Const(4),
                    Cons(Var("b"), Nil))),
            Nil)));

    var env9 := map["a" := 1, "b" := 5];
    assert GoodEnv(deep, env9);

    var env10 := map["a" := 1];
    assert !GoodEnv(deep, env10);

}