Normalize lazily

Closes #7

1 files changed, 212 insertions, 166 deletions

diff --git a/dhall/src/normalize.rs b/dhall/src/normalize.rs
index c23f887..c8ec310 100644
--- a/dhall/src/normalize.rs
+++ b/dhall/src/normalize.rs
@@ -3,194 +3,240 @@ use dhall_core::core::*;
 use dhall_generator::dhall_expr;
 use std::fmt;
 
-/// Reduce an expression to its normal form, performing beta reduction
-///
-/// `normalize` does not type-check the expression.  You may want to type-check
-/// expressions before normalizing them since normalization can convert an
-/// ill-typed expression into a well-typed expression.
-///
-/// However, `normalize` will not fail if the expression is ill-typed and will
-/// leave ill-typed sub-expressions unevaluated.
-///
-pub fn normalize<S, T, A>(e: &Expr<S, A>) -> Expr<T, A>
+/// Reduce an expression to its weak head normal form, i.e. normalize
+/// just enough to get the first constructor of the final expression
+/// Is identical to normalize on primitive types, but not on more complex
+/// types like functions and records.
+/// This allows normalization to be lazy.
+pub fn normalize_whnf<S, A>(e: &Expr<S, A>) -> Expr<S, A>
 where
     S: Clone + fmt::Debug,
-    T: Clone + fmt::Debug,
     A: Clone + fmt::Debug,
 {
     use dhall_core::BinOp::*;
     use dhall_core::Builtin::*;
     use dhall_core::Expr::*;
     match e {
-        // Matches that don't normalize everything right away
         Let(f, _, r, b) => {
-            let r2 = shift::<_, S, _>(1, &V(f.clone(), 0), r);
-            let b2 = subst::<_, S, _>(&V(f.clone(), 0), &r2, b);
-            let b3 = shift::<_, T, _>(-1, &V(f.clone(), 0), &b2);
-            normalize(&b3)
+            let vf0 = &V(f.clone(), 0);
+            let r2 = shift::<_, S, _>(1, vf0, r);
+            let b2 = subst::<_, S, _>(vf0, &r2, b);
+            let b3 = shift::<_, S, _>(-1, vf0, &b2);
+            normalize_whnf(&b3)
         }
-        BoolIf(b, t, f) => match normalize(b) {
-            BoolLit(true) => normalize(t),
-            BoolLit(false) => normalize(f),
-            b2 => BoolIf(bx(b2), bx(normalize(t)), bx(normalize(f))),
-        },
-        Annot(x, _) => normalize(x),
-        Note(_, e) => normalize(e),
-        App(f, a) => match normalize::<S, T, A>(f) {
-            Lam(x, _A, b) => {
+        Annot(x, _) => normalize_whnf(x),
+        Note(_, e) => normalize_whnf(e),
+        App(f, a) => match (normalize_whnf(f), a) {
+            (Lam(x, _, b), a) => {
                 // Beta reduce
                 let vx0 = &V(x, 0);
                 let a2 = shift::<S, S, A>(1, vx0, a);
-                let b2 = subst::<S, T, A>(vx0, &a2, &b);
-                let b3 = shift::<S, T, A>(-1, vx0, &b2);
-                normalize(&b3)
+                let b2 = subst::<S, S, A>(vx0, &a2, &b);
+                let b3 = shift::<S, S, A>(-1, vx0, &b2);
+                normalize_whnf(&b3)
             }
-            f2 => match (f2, normalize::<S, T, A>(a)) {
-                // fold/build fusion for `List`
-                (App(box Builtin(ListBuild), _), App(box App(box Builtin(ListFold), _), box e2)) |
-                (App(box Builtin(ListFold), _), App(box App(box Builtin(ListBuild), _), box e2)) |
+            (Builtin(b), a) => match (b, normalize_whnf(a)) {
+                (NaturalIsZero, NaturalLit(n)) => BoolLit(n == 0),
+                (NaturalEven, NaturalLit(n)) => BoolLit(n % 2 == 0),
+                (NaturalOdd, NaturalLit(n)) => BoolLit(n % 2 != 0),
+                (NaturalToInteger, NaturalLit(n)) => IntegerLit(n as isize),
+                (NaturalShow, NaturalLit(n)) => TextLit(n.to_string()),
 
-                // fold/build fusion for `Optional`
-                (App(box Builtin(OptionalBuild), _), App(box App(box Builtin(OptionalFold), _), box e2)) |
-                (App(box Builtin(OptionalFold), _), App(box App(box Builtin(OptionalBuild), _), box e2)) |
+                (b, App(f, x)) => match (b, normalize_whnf(&f), x) {
+                    // fold/build fusion for `Natural`
+                    (NaturalBuild, Builtin(NaturalFold), x) => {
+                        normalize_whnf(&x)
+                    }
+                    (NaturalFold, Builtin(NaturalBuild), x) => {
+                        normalize_whnf(&x)
+                    }
+                    (b, f, x) => app(Builtin(b), app(f, *x)),
+                },
+                (b, a) => app(Builtin(b), a),
+            },
+            (App(f, x), y) => match (normalize_whnf(&f), x, y) {
+                // TODO: use whnf
+                (Builtin(b), x, y) => match (b, x, normalize::<S, S, A>(&y)) {
+                    (ListLength, _, ListLit(_, ys)) => NaturalLit(ys.len()),
+                    (ListHead, _, ListLit(t, ys)) => {
+                        OptionalLit(t, ys.into_iter().take(1).collect())
+                    }
+                    (ListLast, _, ListLit(t, ys)) => OptionalLit(
+                        t,
+                        ys.into_iter().last().into_iter().collect(),
+                    ),
+                    (ListReverse, _, ListLit(t, ys)) => {
+                        let mut xs = ys;
+                        xs.reverse();
+                        ListLit(t, xs)
+                    }
 
-                // fold/build fusion for `Natural`
-                (Builtin(NaturalBuild), App(box Builtin(NaturalFold), box e2)) |
-                (Builtin(NaturalFold), App(box Builtin(NaturalBuild), box e2)) => normalize(&e2),
+                    // fold/build fusion for `List`
+                    (
+                        ListBuild,
+                        _,
+                        App(box App(box Builtin(ListFold), _), box e2),
+                    ) => normalize_whnf(&e2),
+                    (
+                        ListFold,
+                        _,
+                        App(box App(box Builtin(ListBuild), _), box e2),
+                    ) => normalize_whnf(&e2),
 
-            /*
-                App (App (App (App NaturalFold (NaturalLit n0)) _) succ') zero ->
-                    normalize (go n0)
-                  where
-                    go !0 = zero
-                    go !n = App succ' (go (n - 1))
-                App NaturalBuild k
-                    | check     -> NaturalLit n
-                    | otherwise -> App f' a'
-                  where
-                    labeled =
-                        normalize (App (App (App k Natural) "Succ") "Zero")
+                    // fold/build fusion for `Optional`
+                    (
+                        OptionalBuild,
+                        _,
+                        App(box App(box Builtin(OptionalFold), _), box e2),
+                    ) => normalize_whnf(&e2),
+                    (
+                        OptionalFold,
+                        _,
+                        App(box App(box Builtin(OptionalBuild), _), box e2),
+                    ) => normalize_whnf(&e2),
 
-                    n = go 0 labeled
-                      where
-                        go !m (App (Var "Succ") e') = go (m + 1) e'
-                        go !m (Var "Zero")          = m
-                        go !_  _                    = internalError text
-                    check = go labeled
-                      where
-                        go (App (Var "Succ") e') = go e'
-                        go (Var "Zero")          = True
-                        go  _                    = False
-                        */
-                (Builtin(NaturalIsZero), NaturalLit(n)) => BoolLit(n == 0),
-                (Builtin(NaturalEven), NaturalLit(n)) => BoolLit(n % 2 == 0),
-                (Builtin(NaturalOdd), NaturalLit(n)) => BoolLit(n % 2 != 0),
-                (Builtin(NaturalToInteger), NaturalLit(n)) => IntegerLit(n as isize),
-                (Builtin(NaturalShow), NaturalLit(n)) => TextLit(n.to_string()),
-                (App(box Builtin(ListBuild), a0), k) => {
-                    let k = bx(k);
-                    let a1 = bx(shift(1, &V("a".into(), 0), &a0));
-                    normalize(&dhall_expr!(k (List a0) (λ(a : a0) -> λ(as : List a1) -> [ a ] # as) ([] : List a0)))
-                }
-                (App(box App(box App(box App(box Builtin(ListFold), _), box ListLit(_, xs)), _), cons), nil) => {
-                    let e2: Expr<_, _> = xs.into_iter().rev().fold(nil, |y, ys| {
-                        let y = bx(y);
-                        let ys = bx(ys);
-                        dhall_expr!(cons y ys)
-                    });
-                    normalize(&e2)
-                }
-                (App(f, x_), ListLit(t, ys)) => match *f {
-                    Builtin(ListLength) =>
-                        NaturalLit(ys.len()),
-                    Builtin(ListHead) =>
-                        normalize(&OptionalLit(t, ys.into_iter().take(1).collect())),
-                    Builtin(ListLast) =>
-                        normalize(&OptionalLit(t, ys.into_iter().last().into_iter().collect())),
-                    Builtin(ListReverse) => {
-                        let mut xs = ys;
-                        xs.reverse();
-                        normalize(&ListLit(t, xs))
+                    (ListBuild, a0, k) => {
+                        let k = bx(k);
+                        let a1 = bx(shift(1, &V("a".into(), 0), &a0));
+                        normalize_whnf(
+                            &dhall_expr!(k (List a0) (λ(a : a0) -> λ(as : List a1) -> [ a ] # as) ([] : List a0)),
+                        )
+                    }
+                    (OptionalBuild, a0, g) => {
+                        let g = bx(g);
+                        normalize_whnf(
+                            &dhall_expr!((g (Optional a0)) (λ(x: a0) -> [x] :  Optional a0) ([] :  Optional a0)),
+                        )
+                    }
+                    (b, x, y) => app(App(bx(Builtin(b)), x), y),
+                },
+                (App(f, x), y, z) => match (normalize_whnf(&f), x, y, z) {
+                    (App(f, x), y, z, w) => {
+                        match (normalize_whnf(&f), x, y, z, w) {
+                            (App(f, x), y, z, w, t) => match (
+                                normalize_whnf(&f),
+                                x,
+                                normalize_whnf(&y),
+                                z,
+                                w,
+                                t,
+                            ) {
+                                (
+                                    Builtin(ListFold),
+                                    _,
+                                    ListLit(_, xs),
+                                    _,
+                                    cons,
+                                    nil,
+                                ) => {
+                                    let e2: Expr<_, _> = xs
+                                        .into_iter()
+                                        .rev()
+                                        .fold((**nil).clone(), |y, ys| {
+                                            let y = bx(y);
+                                            let ys = bx(ys);
+                                            dhall_expr!(cons y ys)
+                                        });
+                                    normalize_whnf(&e2)
+                                }
+                                (
+                                    Builtin(OptionalFold),
+                                    _,
+                                    OptionalLit(_, xs),
+                                    _,
+                                    just,
+                                    nothing,
+                                ) => {
+                                    let e2: Expr<_, _> = xs.into_iter().fold(
+                                        (**nothing).clone(),
+                                        |y, _| {
+                                            let y = bx(y);
+                                            dhall_expr!(just y)
+                                        },
+                                    );
+                                    normalize_whnf(&e2)
+                                }
+                                (f, x, y, z, w, t) => app(
+                                    App(
+                                        bx(App(
+                                            bx(App(bx(App(bx(f), x)), bx(y))),
+                                            z,
+                                        )),
+                                        w,
+                                    ),
+                                    (**t).clone(),
+                                ),
+                            },
+                            (f, x, y, z, w) => app(
+                                App(bx(App(bx(App(bx(f), x)), y)), z),
+                                (**w).clone(),
+                            ),
+                        }
+                    }
+                    (f, x, y, z) => {
+                        app(App(bx(App(bx(f), x)), y), (**z).clone())
                     }
-                    _ => app(App(f, x_), ListLit(t, ys)),
                 },
-                /*
-                App (App ListIndexed _) (ListLit t xs) ->
-                    normalize (ListLit t' (fmap adapt (Data.Vector.indexed xs)))
-                  where
-                    t' = Record (Data.Map.fromList kts)
-                      where
-                        kts = [ ("index", Natural)
-                              , ("value", t)
-                              ]
-                    adapt (n, x) = RecordLit (Data.Map.fromList kvs)
-                      where
-                        kvs = [ ("index", NaturalLit (fromIntegral n))
-                              , ("value", x)
-                              ]
-            */
-                (App(box App(box App(box App(box Builtin(OptionalFold), _), box OptionalLit(_, xs)), _), just), nothing) => {
-                    let e2: Expr<_, _> = xs.into_iter().fold(nothing, |y, _| {
-                        let y = bx(y);
-                        dhall_expr!(just y)
-                    });
-                    normalize(&e2)
-                }
-                (App(box Builtin(OptionalBuild), a0), g) => {
-                    let g = bx(g);
-                    normalize(&dhall_expr!((g (Optional a0)) (λ(x: a0) -> [x] :  Optional a0) ([] :  Optional a0)))
-                }
-                (f2, a2) => app(f2, a2),
+                (f, x, y) => app(App(bx(f), x), (**y).clone()),
             },
+            (f, a) => app(f, (**a).clone()),
         },
-
-        // Normalize everything else before matching
-        e => {
-            match e.map_shallow(
-                normalize,
-                |_| unreachable!(),
-                |x| x.clone(),
-                |x| x.clone(),
-            ) {
-                BinOp(BoolAnd, box BoolLit(x), box BoolLit(y)) => {
-                    BoolLit(x && y)
-                }
-                BinOp(BoolOr, box BoolLit(x), box BoolLit(y)) => {
-                    BoolLit(x || y)
-                }
-                BinOp(BoolEQ, box BoolLit(x), box BoolLit(y)) => {
-                    BoolLit(x == y)
-                }
-                BinOp(BoolNE, box BoolLit(x), box BoolLit(y)) => {
-                    BoolLit(x != y)
-                }
-                BinOp(NaturalPlus, box NaturalLit(x), box NaturalLit(y)) => {
-                    NaturalLit(x + y)
-                }
-                BinOp(NaturalTimes, box NaturalLit(x), box NaturalLit(y)) => {
-                    NaturalLit(x * y)
-                }
-                BinOp(TextAppend, box TextLit(x), box TextLit(y)) => {
-                    TextLit(x + &y)
-                }
-                BinOp(ListAppend, box ListLit(t1, xs), box ListLit(t2, ys)) => {
-                    // Drop type annotation if the result is nonempty
-                    let t = if xs.is_empty() && ys.is_empty() {
-                        t1.or(t2)
-                    } else {
-                        None
-                    };
-                    let xs = xs.into_iter();
-                    let ys = ys.into_iter();
-                    ListLit(t, xs.chain(ys).collect())
-                }
-                // Merge(_x, _y, _t) => unimplemented!(),
-                Field(box RecordLit(kvs), x) => match kvs.get(&x) {
-                    Some(r) => r.clone(),
-                    None => Field(bx(RecordLit(kvs)), x),
-                },
-                e => e,
+        BoolIf(b, t, f) => match normalize_whnf(b) {
+            BoolLit(true) => normalize_whnf(t),
+            BoolLit(false) => normalize_whnf(f),
+            b2 => BoolIf(bx(b2), t.clone(), f.clone()),
+        },
+        BinOp(o, x, y) => match (o, normalize_whnf(&x), normalize_whnf(&y)) {
+            (BoolAnd, BoolLit(x), BoolLit(y)) => BoolLit(x && y),
+            (BoolOr, BoolLit(x), BoolLit(y)) => BoolLit(x || y),
+            (BoolEQ, BoolLit(x), BoolLit(y)) => BoolLit(x == y),
+            (BoolNE, BoolLit(x), BoolLit(y)) => BoolLit(x != y),
+            (NaturalPlus, NaturalLit(x), NaturalLit(y)) => NaturalLit(x + y),
+            (NaturalTimes, NaturalLit(x), NaturalLit(y)) => NaturalLit(x * y),
+            (TextAppend, TextLit(x), TextLit(y)) => TextLit(x + &y),
+            (ListAppend, ListLit(t1, xs), ListLit(t2, ys)) => {
+                // Drop type annotation if the result is nonempty
+                let t = if xs.is_empty() && ys.is_empty() {
+                    t1.or(t2)
+                } else {
+                    None
+                };
+                let xs = xs.into_iter();
+                let ys = ys.into_iter();
+                ListLit(t, xs.chain(ys).collect())
             }
-        }
+            (o, x, y) => BinOp(*o, bx(x), bx(y)),
+        },
+        Field(e, x) => match (normalize_whnf(&e), x) {
+            (RecordLit(kvs), x) => match kvs.get(&x) {
+                Some(r) => normalize_whnf(r),
+                None => Field(bx(RecordLit(kvs)), x.clone()),
+            },
+            (e, x) => Field(bx(e), x.clone()),
+        },
+        e => e.clone(),
     }
 }
+
+/// Reduce an expression to its normal form, performing beta reduction
+///
+/// `normalize` does not type-check the expression.  You may want to type-check
+/// expressions before normalizing them since normalization can convert an
+/// ill-typed expression into a well-typed expression.
+///
+/// However, `normalize` will not fail if the expression is ill-typed and will
+/// leave ill-typed sub-expressions unevaluated.
+///
+pub fn normalize<S, T, A>(e: &Expr<S, A>) -> Expr<T, A>
+where
+    S: Clone + fmt::Debug,
+    T: Clone + fmt::Debug,
+    A: Clone + fmt::Debug,
+{
+    normalize_whnf(e).map_shallow(
+        normalize,
+        |_| unreachable!(),
+        |x| x.clone(),
+        |x| x.clone(),
+    )
+}