Share type alongside the value in a thunk

2 files changed, 110 insertions, 40 deletions

diff --git a/dhall/src/core/thunk.rs b/dhall/src/core/thunk.rs
index b6358ef..325a410 100644
--- a/dhall/src/core/thunk.rs
+++ b/dhall/src/core/thunk.rs
@@ -1,5 +1,5 @@
 use std::borrow::Cow;
-use std::cell::{Ref, RefCell};
+use std::cell::{Ref, RefCell, RefMut};
 use std::rc::Rc;
 
 use dhall_syntax::{Const, ExprF};
@@ -21,7 +21,7 @@ enum Marker {
 }
 use Marker::{NF, WHNF};
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 enum ThunkInternal {
     /// Partially normalized value whose subexpressions have been thunked (this is returned from
     /// typechecking). Note that this is different from `ValueF::PartialExpr` because there is no
@@ -32,19 +32,29 @@ enum ThunkInternal {
     ValueF(Marker, ValueF),
 }
 
+#[derive(Debug)]
+struct TypedThunkInternal {
+    internal: ThunkInternal,
+    typ: Option<Type>,
+}
+
 /// Stores a possibly unevaluated value. Gets (partially) normalized on-demand,
 /// sharing computation automatically. Uses a RefCell to share computation.
 /// Can optionally store a Type from the typechecking phase to preserve type information through
 /// the normalization phase.
 #[derive(Clone)]
-pub struct Thunk(Rc<RefCell<ThunkInternal>>, Option<Box<Type>>);
+pub struct Thunk(Rc<RefCell<TypedThunkInternal>>);
 
 #[derive(Debug, Clone)]
 pub struct TypedThunk(Thunk);
 
 impl ThunkInternal {
-    fn into_thunk(self, t: Option<Box<Type>>) -> Thunk {
-        Thunk(Rc::new(RefCell::new(self)), t)
+    fn into_thunk(self, t: Option<Type>) -> Thunk {
+        TypedThunkInternal {
+            internal: self,
+            typ: t,
+        }
+        .into_thunk()
     }
 
     fn normalize_whnf(&mut self) {
@@ -92,37 +102,78 @@ impl ThunkInternal {
     }
 }
 
+impl TypedThunkInternal {
+    fn into_thunk(self) -> Thunk {
+        Thunk(Rc::new(RefCell::new(self)))
+    }
+    fn as_internal(&self) -> &ThunkInternal {
+        &self.internal
+    }
+    fn as_internal_mut(&mut self) -> &mut ThunkInternal {
+        &mut self.internal
+    }
+
+    fn get_type(&self) -> Result<Type, TypeError> {
+        match &self.typ {
+            Some(t) => Ok(t.clone()),
+            None => Err(TypeError::new(
+                &TypecheckContext::new(),
+                TypeMessage::Untyped,
+            )),
+        }
+    }
+}
+
 impl Thunk {
     pub(crate) fn from_valuef(v: ValueF) -> Thunk {
         ThunkInternal::ValueF(WHNF, v).into_thunk(None)
     }
     pub(crate) fn from_valuef_and_type(v: ValueF, t: Type) -> Thunk {
-        ThunkInternal::ValueF(WHNF, v).into_thunk(Some(Box::new(t)))
+        ThunkInternal::ValueF(WHNF, v).into_thunk(Some(t))
     }
     pub(crate) fn from_partial_expr(e: ExprF<Thunk, Normalized>) -> Thunk {
         ThunkInternal::PartialExpr(e).into_thunk(None)
     }
     pub(crate) fn with_type(self, t: Type) -> Thunk {
-        Thunk(self.0, Some(Box::new(t)))
+        self.as_internal().clone().into_thunk(Some(t))
     }
 
-    // Normalizes contents to normal form; faster than `normalize_nf` if
-    // no one else shares this thunk
-    pub(crate) fn normalize_mut(&mut self) {
+    /// Mutates the contents. If no one else shares this thunk,
+    /// mutates directly, thus avoiding a RefCell lock.
+    fn mutate_internal(&mut self, f: impl FnOnce(&mut TypedThunkInternal)) {
         match Rc::get_mut(&mut self.0) {
             // Mutate directly if sole owner
-            Some(refcell) => RefCell::get_mut(refcell).normalize_nf(),
+            Some(refcell) => f(RefCell::get_mut(refcell)),
             // Otherwise mutate through the refcell
-            None => self.0.borrow_mut().normalize_nf(),
+            None => f(&mut self.as_tinternal_mut()),
         }
     }
 
+    /// Normalizes contents to normal form; faster than `normalize_nf` if
+    /// no one else shares this thunk
+    pub(crate) fn normalize_mut(&mut self) {
+        self.mutate_internal(|i| i.as_internal_mut().normalize_nf())
+    }
+
+    fn as_tinternal(&self) -> Ref<TypedThunkInternal> {
+        self.0.borrow()
+    }
+    fn as_tinternal_mut(&mut self) -> RefMut<TypedThunkInternal> {
+        self.0.borrow_mut()
+    }
+    fn as_internal(&self) -> Ref<ThunkInternal> {
+        Ref::map(self.as_tinternal(), TypedThunkInternal::as_internal)
+    }
+    fn as_internal_mut(&self) -> RefMut<ThunkInternal> {
+        RefMut::map(self.0.borrow_mut(), TypedThunkInternal::as_internal_mut)
+    }
+
     fn do_normalize_whnf(&self) {
-        let borrow = self.0.borrow();
+        let borrow = self.as_internal();
         match &*borrow {
             ThunkInternal::PartialExpr(_) => {
                 drop(borrow);
-                self.0.borrow_mut().normalize_whnf();
+                self.as_internal_mut().normalize_whnf();
             }
             // Already at least in WHNF
             ThunkInternal::ValueF(_, _) => {}
@@ -130,11 +181,11 @@ impl Thunk {
     }
 
     fn do_normalize_nf(&self) {
-        let borrow = self.0.borrow();
+        let borrow = self.as_internal();
         match &*borrow {
             ThunkInternal::PartialExpr(_) | ThunkInternal::ValueF(WHNF, _) => {
                 drop(borrow);
-                self.0.borrow_mut().normalize_nf();
+                self.as_internal_mut().normalize_nf();
             }
             // Already in NF
             ThunkInternal::ValueF(NF, _) => {}
@@ -145,14 +196,14 @@ impl Thunk {
     // or you could run into BorrowMut panics
     pub(crate) fn normalize_nf(&self) -> Ref<ValueF> {
         self.do_normalize_nf();
-        Ref::map(self.0.borrow(), ThunkInternal::as_nf)
+        Ref::map(self.as_internal(), ThunkInternal::as_nf)
     }
 
     // WARNING: avoid normalizing any thunk while holding on to this ref
     // or you could run into BorrowMut panics
     pub(crate) fn as_valuef(&self) -> Ref<ValueF> {
         self.do_normalize_whnf();
-        Ref::map(self.0.borrow(), ThunkInternal::as_whnf)
+        Ref::map(self.as_internal(), ThunkInternal::as_whnf)
     }
 
     pub(crate) fn to_valuef(&self) -> ValueF {
@@ -173,6 +224,10 @@ impl Thunk {
     pub(crate) fn app_thunk(&self, th: Thunk) -> ValueF {
         apply_any(self.clone(), th)
     }
+
+    pub(crate) fn get_type(&self) -> Result<Cow<'_, Type>, TypeError> {
+        Ok(Cow::Owned(self.as_tinternal().get_type()?))
+    }
 }
 
 impl TypedThunk {
@@ -185,9 +240,7 @@ impl TypedThunk {
     }
 
     pub(crate) fn normalize_nf(&self) -> ValueF {
-        match self {
-            TypedThunk(thunk) => thunk.normalize_nf().clone(),
-        }
+        self.0.normalize_nf().clone()
     }
 
     pub(crate) fn to_typed(&self) -> Typed {
@@ -217,10 +270,7 @@ impl TypedThunk {
         }
     }
     pub(crate) fn from_valuef_and_type(v: ValueF, t: Type) -> Self {
-        TypedThunk::from_thunk_and_type(
-            Thunk::from_valuef_and_type(v, t.clone()),
-            t,
-        )
+        TypedThunk(Thunk::from_valuef_and_type(v, t))
     }
 
     pub(crate) fn to_valuef(&self) -> ValueF {
@@ -254,19 +304,13 @@ impl TypedThunk {
     }
 
     pub(crate) fn get_type(&self) -> Result<Cow<'_, Type>, TypeError> {
-        match &(self.0).1 {
-            Some(t) => Ok(Cow::Borrowed(&*t)),
-            None => Err(TypeError::new(
-                &TypecheckContext::new(),
-                TypeMessage::Untyped,
-            )),
-        }
+        self.0.get_type()
     }
 }
 
 impl Shift for Thunk {
     fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
-        Some(Thunk(self.0.shift(delta, var)?, self.1.shift(delta, var)?))
+        Some(Thunk(self.0.shift(delta, var)?))
     }
 }
 
@@ -285,15 +329,22 @@ impl Shift for ThunkInternal {
 
 impl Shift for TypedThunk {
     fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
-        Some(match self {
-            TypedThunk(th) => TypedThunk(th.shift(delta, var)?),
+        Some(TypedThunk(self.0.shift(delta, var)?))
+    }
+}
+
+impl Shift for TypedThunkInternal {
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some(TypedThunkInternal {
+            internal: self.internal.shift(delta, var)?,
+            typ: self.typ.shift(delta, var)?,
         })
     }
 }
 
 impl Subst<Typed> for Thunk {
     fn subst_shift(&self, var: &AlphaVar, val: &Typed) -> Self {
-        Thunk(self.0.subst_shift(var, val), self.1.subst_shift(var, val))
+        Thunk(self.0.subst_shift(var, val))
     }
 }
 
@@ -311,14 +362,21 @@ impl Subst<Typed> for ThunkInternal {
     }
 }
 
-impl Subst<Typed> for TypedThunk {
+impl Subst<Typed> for TypedThunkInternal {
     fn subst_shift(&self, var: &AlphaVar, val: &Typed) -> Self {
-        match self {
-            TypedThunk(th) => TypedThunk(th.subst_shift(var, val)),
+        TypedThunkInternal {
+            internal: self.internal.subst_shift(var, val),
+            typ: self.typ.subst_shift(var, val),
         }
     }
 }
 
+impl Subst<Typed> for TypedThunk {
+    fn subst_shift(&self, var: &AlphaVar, val: &Typed) -> Self {
+        TypedThunk(self.0.subst_shift(var, val))
+    }
+}
+
 impl std::cmp::PartialEq for Thunk {
     fn eq(&self, other: &Self) -> bool {
         *self.as_valuef() == *other.as_valuef()
@@ -335,7 +393,7 @@ impl std::cmp::Eq for TypedThunk {}
 
 impl std::fmt::Debug for Thunk {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let b: &ThunkInternal = &self.0.borrow();
+        let b: &ThunkInternal = &self.as_internal();
         match b {
             ThunkInternal::ValueF(m, v) => {
                 f.debug_tuple(&format!("Thunk@{:?}", m)).field(v).finish()
diff --git a/dhall/src/core/var.rs b/dhall/src/core/var.rs
index 8d02171..ce4d137 100644
--- a/dhall/src/core/var.rs
+++ b/dhall/src/core/var.rs
@@ -157,6 +157,12 @@ impl Shift for () {
     }
 }
 
+impl<A: Shift, B: Shift> Shift for (A, B) {
+    fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
+        Some((self.0.shift(delta, var)?, self.1.shift(delta, var)?))
+    }
+}
+
 impl<T: Shift> Shift for Option<T> {
     fn shift(&self, delta: isize, var: &AlphaVar) -> Option<Self> {
         Some(match self {
@@ -230,6 +236,12 @@ impl<S> Subst<S> for () {
     fn subst_shift(&self, _var: &AlphaVar, _val: &S) -> Self {}
 }
 
+impl<S, A: Subst<S>, B: Subst<S>> Subst<S> for (A, B) {
+    fn subst_shift(&self, var: &AlphaVar, val: &S) -> Self {
+        (self.0.subst_shift(var, val), self.1.subst_shift(var, val))
+    }
+}
+
 impl<S, T: Subst<S>> Subst<S> for Option<T> {
     fn subst_shift(&self, var: &AlphaVar, val: &S) -> Self {
         self.as_ref().map(|x| x.subst_shift(var, val))