Encapsulate partially applied builtin in a separate struct

6 files changed, 120 insertions, 74 deletions

diff --git a/dhall/src/semantics/builtins.rs b/dhall/src/semantics/builtins.rs
index 0b105c7..a23591e 100644
--- a/dhall/src/semantics/builtins.rs
+++ b/dhall/src/semantics/builtins.rs
@@ -1,5 +1,7 @@
 use crate::semantics::phase::Normalized;
-use crate::semantics::{typecheck, NzEnv, Value, ValueKind};
+use crate::semantics::{
+    typecheck, NzEnv, TyExpr, TyExprKind, Value, ValueKind, VarEnv,
+};
 use crate::syntax::map::DupTreeMap;
 use crate::syntax::Const::Type;
 use crate::syntax::{
@@ -9,6 +11,66 @@ use crate::syntax::{
 use std::collections::HashMap;
 use std::convert::TryInto;
 
+/// A partially applied builtin.
+/// Invariant: TODO
+#[derive(Debug, Clone)]
+pub(crate) struct BuiltinClosure<Value> {
+    pub env: NzEnv,
+    pub b: Builtin,
+    /// Arguments applied to the closure so far.
+    pub args: Vec<Value>,
+    /// Keeps the types of the partial applications around to be able to convert back to TyExpr.
+    /// If the args so far are `x_1`, ..., `x_n`, this contains the types of `b`, `b x1`, ...,
+    /// `b x_1 x_2 ... x_(n-1)`.
+    pub types: Vec<Value>,
+}
+
+impl BuiltinClosure<Value> {
+    pub fn new(b: Builtin, env: NzEnv) -> Self {
+        BuiltinClosure {
+            env,
+            b,
+            args: Vec::new(),
+            types: Vec::new(),
+        }
+    }
+
+    pub fn apply(
+        &self,
+        a: Value,
+        f_ty: Value,
+        ret_ty: &Value,
+    ) -> ValueKind<Value> {
+        use std::iter::once;
+        let args = self.args.iter().cloned().chain(once(a.clone())).collect();
+        let types = self.types.iter().cloned().chain(once(f_ty)).collect();
+        apply_builtin(self.b, args, ret_ty, types, self.env.clone())
+    }
+    pub fn ensure_whnf(self, ret_ty: &Value) -> ValueKind<Value> {
+        apply_builtin(self.b, self.args, ret_ty, self.types, self.env)
+    }
+    pub fn normalize_mut(&mut self) {
+        for x in self.args.iter_mut() {
+            x.normalize_mut();
+        }
+    }
+    pub fn to_tyexprkind(&self, venv: VarEnv) -> TyExprKind {
+        TyExprKind::Expr(self.args.iter().zip(self.types.iter()).fold(
+            ExprKind::Builtin(self.b),
+            |acc, (v, ty)| {
+                ExprKind::App(
+                    TyExpr::new(
+                        TyExprKind::Expr(acc),
+                        Some(ty.clone()),
+                        Span::Artificial,
+                    ),
+                    v.to_tyexpr(venv),
+                )
+            },
+        ))
+    }
+}
+
 pub(crate) fn rc<E>(x: UnspannedExpr<E>) -> Expr<E> {
     Expr::new(x, Span::Artificial)
 }
@@ -208,7 +270,7 @@ macro_rules! make_closure {
 }
 
 #[allow(clippy::cognitive_complexity)]
-pub(crate) fn apply_builtin(
+fn apply_builtin(
     b: Builtin,
     args: Vec<Value>,
     ty: &Value,
@@ -488,6 +550,18 @@ pub(crate) fn apply_builtin(
     match ret {
         Ret::ValueKind(v) => v,
         Ret::Value(v) => v.to_whnf_check_type(ty),
-        Ret::DoneAsIs => AppliedBuiltin(b, args, types, env),
+        Ret::DoneAsIs => AppliedBuiltin(BuiltinClosure {
+            b,
+            args,
+            types,
+            env,
+        }),
+    }
+}
+
+impl<Value: std::cmp::PartialEq> std::cmp::PartialEq for BuiltinClosure<Value> {
+    fn eq(&self, other: &Self) -> bool {
+        self.b == other.b && self.args == other.args
     }
 }
+impl<Value: std::cmp::Eq> std::cmp::Eq for BuiltinClosure<Value> {}
diff --git a/dhall/src/semantics/core/value.rs b/dhall/src/semantics/core/value.rs
index 7ecb2bf..0fda870 100644
--- a/dhall/src/semantics/core/value.rs
+++ b/dhall/src/semantics/core/value.rs
@@ -7,7 +7,7 @@ use crate::semantics::core::var::Binder;
 use crate::semantics::phase::normalize::{apply_any, normalize_whnf};
 use crate::semantics::phase::{Normalized, NormalizedExpr, ToExprOptions};
 use crate::semantics::{type_of_builtin, TyExpr, TyExprKind};
-use crate::semantics::{NzEnv, NzVar, VarEnv};
+use crate::semantics::{BuiltinClosure, NzEnv, NzVar, VarEnv};
 use crate::syntax::{
     BinOp, Builtin, Const, ExprKind, Integer, InterpolatedTextContents, Label,
     NaiveDouble, Natural, Span,
@@ -76,9 +76,7 @@ pub(crate) enum ValueKind<Value> {
         closure: Closure,
     },
     // Invariant: in whnf, the evaluation must not be able to progress further.
-    // Keep types around to be able to recover the types of partial applications.
-    // Keep env around to construct Foo/build closures; hopefully temporary.
-    AppliedBuiltin(Builtin, Vec<Value>, Vec<Value>, NzEnv),
+    AppliedBuiltin(BuiltinClosure<Value>),
 
     Var(NzVar),
     Const(Const),
@@ -292,21 +290,7 @@ impl Value {
                     .apply_var(NzVar::new(venv.size()))
                     .to_tyexpr(venv.insert()),
             )),
-            ValueKind::AppliedBuiltin(b, args, types, _) => {
-                TyExprKind::Expr(args.into_iter().zip(types.into_iter()).fold(
-                    ExprKind::Builtin(*b),
-                    |acc, (v, ty)| {
-                        ExprKind::App(
-                            TyExpr::new(
-                                TyExprKind::Expr(acc),
-                                Some(ty.clone()),
-                                Span::Artificial,
-                            ),
-                            v.to_tyexpr(venv),
-                        )
-                    },
-                ))
-            }
+            ValueKind::AppliedBuiltin(closure) => closure.to_tyexprkind(venv),
             ValueKind::UnionConstructor(l, kts, t) => {
                 TyExprKind::Expr(ExprKind::Field(
                     TyExpr::new(
@@ -504,11 +488,7 @@ impl ValueKind<Value> {
             | ValueKind::PiClosure { annot, .. } => {
                 annot.normalize_mut();
             }
-            ValueKind::AppliedBuiltin(_, args, _, _) => {
-                for x in args.iter_mut() {
-                    x.normalize_mut();
-                }
-            }
+            ValueKind::AppliedBuiltin(closure) => closure.normalize_mut(),
             ValueKind::NEListLit(elts) => {
                 for x in elts.iter_mut() {
                     x.normalize_mut();
@@ -555,7 +535,7 @@ impl ValueKind<Value> {
         ValueKind::from_builtin_env(b, NzEnv::new())
     }
     pub(crate) fn from_builtin_env(b: Builtin, env: NzEnv) -> ValueKind<Value> {
-        ValueKind::AppliedBuiltin(b, vec![], vec![], env)
+        ValueKind::AppliedBuiltin(BuiltinClosure::new(b, env))
     }
 }
 
diff --git a/dhall/src/semantics/core/visitor.rs b/dhall/src/semantics/core/visitor.rs
index 2191ce3..a78c219 100644
--- a/dhall/src/semantics/core/visitor.rs
+++ b/dhall/src/semantics/core/visitor.rs
@@ -1,6 +1,6 @@
 use std::iter::FromIterator;
 
-use crate::semantics::{Binder, ValueKind};
+use crate::semantics::{Binder, BuiltinClosure, ValueKind};
 use crate::syntax::Label;
 
 /// A visitor trait that can be used to traverse `ValueKind`s. We need this pattern so that Rust lets
@@ -90,12 +90,17 @@ where
             annot: v.visit_val(annot)?,
             closure: closure.clone(),
         },
-        AppliedBuiltin(b, xs, types, env) => AppliedBuiltin(
-            *b,
-            v.visit_vec(xs)?,
-            v.visit_vec(types)?,
-            env.clone(),
-        ),
+        AppliedBuiltin(BuiltinClosure {
+            b,
+            args,
+            types,
+            env,
+        }) => AppliedBuiltin(BuiltinClosure {
+            b: *b,
+            args: v.visit_vec(args)?,
+            types: v.visit_vec(types)?,
+            env: env.clone(),
+        }),
         Var(v) => Var(*v),
         Const(k) => Const(*k),
         BoolLit(b) => BoolLit(*b),
diff --git a/dhall/src/semantics/nze/env.rs b/dhall/src/semantics/nze/env.rs
index dee597a..a083076 100644
--- a/dhall/src/semantics/nze/env.rs
+++ b/dhall/src/semantics/nze/env.rs
@@ -70,11 +70,3 @@ impl NzEnv {
         }
     }
 }
-
-/// Ignore NzEnv when comparing; useful because we store them in `AppliedBuiltin`.
-impl std::cmp::PartialEq for NzEnv {
-    fn eq(&self, _other: &Self) -> bool {
-        true
-    }
-}
-impl std::cmp::Eq for NzEnv {}
diff --git a/dhall/src/semantics/phase/normalize.rs b/dhall/src/semantics/phase/normalize.rs
index 2d4b4b3..d7720c7 100644
--- a/dhall/src/semantics/phase/normalize.rs
+++ b/dhall/src/semantics/phase/normalize.rs
@@ -3,10 +3,11 @@ use std::collections::HashMap;
 use crate::semantics::phase::Normalized;
 use crate::semantics::NzEnv;
 use crate::semantics::{
-    apply_builtin, Binder, Closure, TyExpr, TyExprKind, Value, ValueKind,
+    Binder, BuiltinClosure, Closure, TyExpr, TyExprKind, Value, ValueKind,
+};
+use crate::syntax::{
+    BinOp, Builtin, Const, ExprKind, InterpolatedTextContents,
 };
-use crate::syntax;
-use crate::syntax::{BinOp, Const, ExprKind, InterpolatedTextContents};
 
 pub(crate) fn apply_any(f: Value, a: Value, ty: &Value) -> ValueKind<Value> {
     let f_borrow = f.as_whnf();
@@ -14,15 +15,8 @@ pub(crate) fn apply_any(f: Value, a: Value, ty: &Value) -> ValueKind<Value> {
         ValueKind::LamClosure { closure, .. } => {
             closure.apply(a).to_whnf_check_type(ty)
         }
-        ValueKind::AppliedBuiltin(b, args, types, env) => {
-            use std::iter::once;
-            let args = args.iter().cloned().chain(once(a.clone())).collect();
-            let types = types
-                .iter()
-                .cloned()
-                .chain(once(f.get_type().unwrap()))
-                .collect();
-            apply_builtin(*b, args, ty, types, env.clone())
+        ValueKind::AppliedBuiltin(closure) => {
+            closure.apply(a, f.get_type().unwrap(), ty)
         }
         ValueKind::UnionConstructor(l, kts, uniont) => ValueKind::UnionLit(
             l.clone(),
@@ -293,12 +287,11 @@ pub(crate) fn normalize_one_layer(
         ExprKind::SomeLit(e) => Ret::ValueKind(NEOptionalLit(e)),
         ExprKind::EmptyListLit(t) => {
             let arg = match &*t.as_whnf() {
-                ValueKind::AppliedBuiltin(
-                    syntax::Builtin::List,
+                ValueKind::AppliedBuiltin(BuiltinClosure {
+                    b: Builtin::List,
                     args,
-                    _,
-                    _,
-                ) if args.len() == 1 => args[0].clone(),
+                    ..
+                }) if args.len() == 1 => args[0].clone(),
                 _ => panic!("internal type error"),
             };
             Ret::ValueKind(ValueKind::EmptyListLit(arg))
@@ -462,9 +455,7 @@ pub(crate) fn normalize_whnf(
     ty: &Value,
 ) -> ValueKind<Value> {
     match v {
-        ValueKind::AppliedBuiltin(b, args, types, env) => {
-            apply_builtin(b, args, ty, types, env)
-        }
+        ValueKind::AppliedBuiltin(closure) => closure.ensure_whnf(ty),
         ValueKind::PartialExpr(e) => normalize_one_layer(e, ty, &NzEnv::new()),
         ValueKind::TextLit(elts) => {
             ValueKind::TextLit(squash_textlit(elts.into_iter()))
diff --git a/dhall/src/semantics/tck/typecheck.rs b/dhall/src/semantics/tck/typecheck.rs
index 5422e6d..e1a9c11 100644
--- a/dhall/src/semantics/tck/typecheck.rs
+++ b/dhall/src/semantics/tck/typecheck.rs
@@ -6,10 +6,9 @@ use crate::error::{TypeError, TypeMessage};
 use crate::semantics::phase::normalize::merge_maps;
 use crate::semantics::phase::Normalized;
 use crate::semantics::{
-    type_of_builtin, Binder, Closure, NzVar, TyEnv, TyExpr, TyExprKind, Type,
-    Value, ValueKind,
+    type_of_builtin, Binder, BuiltinClosure, Closure, NzVar, TyEnv, TyExpr,
+    TyExprKind, Type, Value, ValueKind,
 };
-use crate::syntax;
 use crate::syntax::{
     BinOp, Builtin, Const, Expr, ExprKind, InterpolatedTextContents, Span,
 };
@@ -96,8 +95,11 @@ fn type_one_layer(
         ExprKind::EmptyListLit(t) => {
             let t = t.normalize_nf(env.as_nzenv());
             match &*t.as_whnf() {
-                ValueKind::AppliedBuiltin(Builtin::List, args, _, _)
-                    if args.len() == 1 => {}
+                ValueKind::AppliedBuiltin(BuiltinClosure {
+                    b: Builtin::List,
+                    args,
+                    ..
+                }) if args.len() == 1 => {}
                 _ => return mkerr("InvalidListType"),
             };
             t
@@ -365,7 +367,10 @@ fn type_one_layer(
         ExprKind::BinOp(BinOp::ListAppend, l, r) => {
             let l_ty = l.get_type()?;
             match &*l_ty.as_whnf() {
-                ValueKind::AppliedBuiltin(Builtin::List, _, _, _) => {}
+                ValueKind::AppliedBuiltin(BuiltinClosure {
+                    b: Builtin::List,
+                    ..
+                }) => {}
                 _ => return mkerr("BinOpTypeMismatch"),
             }
 
@@ -423,12 +428,11 @@ fn type_one_layer(
             let union_borrow = union_type.as_whnf();
             let variants = match &*union_borrow {
                 ValueKind::UnionType(kts) => Cow::Borrowed(kts),
-                ValueKind::AppliedBuiltin(
-                    syntax::Builtin::Optional,
+                ValueKind::AppliedBuiltin(BuiltinClosure {
+                    b: Builtin::Optional,
                     args,
-                    _,
-                    _,
-                ) if args.len() == 1 => {
+                    ..
+                }) if args.len() == 1 => {
                     let ty = &args[0];
                     let mut kts = HashMap::new();
                     kts.insert("None".into(), None);