Removing thunks from the typechecker

Yatima/Typechecker/Datatypes.lean

@@ -90,11 +90,20 @@ def TypedConst.type : TypedConst → TypedExpr
   | recursor    type ..
   | quotient    type .. => type
 
-structure Env' (SusValue : Type) where
-  exprs : List SusValue
+structure Env' (Value : Type) where
+  exprs : List Value
   univs : List Univ
   deriving Inhabited
 
+inductive ValueException where
+  | custom : String → ValueException
+  | noVar : ValueException
+deriving Inhabited
+
+instance : ToString ValueException where toString
+  | .custom e => s!"{e}"
+  | .noVar => s!"Free variable found. This means the encoding of the expression is incorrect"
+
 mutual
   /--
   Values are the final result of the evaluation of well-typed expressions under a well-typed
@@ -109,16 +118,17 @@ mutual
     We also keep the `TypeInfo` of each subapplication (`neu a_1 a_2 ... a_i`), for
     i = 0, .. , n-1; this preserves information necessary to implement the quoting
     (i.e. read-back) functionality that is used in lambda inference -/
-    | app : Neutral → List (AddInfo (Thunk Value)) → List TypeInfo → Value
+    | app : Neutral → List (AddInfo Value) → List TypeInfo → Value
     /-- Lambdas are unevaluated expressions with environments for their free
     variables apart from their argument variables -/
-    | lam : AddInfo (Thunk Value) → TypedExpr → Env' (AddInfo (Thunk Value)) → Value
-    /-- Pi types will have thunks for their domains and unevaluated expressions
+    | lam : AddInfo Value → TypedExpr → Env' (AddInfo Value) → Value
+    /-- Pi types will have values for their domains and unevaluated expressions
     analogous to lambda bodies for their codomains -/
-    | pi : AddInfo (Thunk Value) → TypedExpr → Env' (AddInfo (Thunk Value)) → Value
+    | pi  : AddInfo Value → TypedExpr → Env' (AddInfo Value) → Value
     | lit : Literal → Value
-    -- An exception constructor is used to catch bugs in the evaluator/typechecker
-    | exception : String → Value
+    -- An exception constructor is added so that we don't have to wrap `Exception` around `Value`.
+    -- This is both for convenience and efficiency.
+    | exception : ValueException → Value
     deriving Inhabited
   /--
   A neutral term is either a variable or a constant with not enough arguments to
@@ -134,14 +144,6 @@ end
 
 abbrev TypedValue := AddInfo Value
 
-/--
-Suspended values are thunks that return a value. For optimization purposes, the value's
-`TypeInfo`, which by type preservation comes from the underlying expression that gave
-rise to this value by means of evaluation, is saved outside the thunk, instead of in
-the values themselves. This allows us to extract it without needing to force the thunk.
--/
-abbrev SusValue := AddInfo (Thunk Value)
-
 /--
 The environment will bind free variables to different things, depending on
 the evaluation strategy:
@@ -150,30 +152,15 @@ the evaluation strategy:
 2) Non-strict evaluation: binds free variables to unevaluated expressions
 3) Lazy evaluation (i.e. non-strict without duplication of work): binds free variables to thunks
 
-Here we chose lazy evaluation since it is more efficient for typechecking.
+Here we will chose strict evaluation
 
 Since we also have universes with free variables, we need to add a environment
 for universe variables as well
 -/
-abbrev Env := Env' SusValue
+abbrev Env := Env' TypedValue
 
 /-- The arguments of a stuck sequence of applications `(h a1 ... an)` -/
-abbrev Args := List SusValue
-
-instance : Inhabited SusValue where
-  default := .mk default {fn := default}
-
--- Auxiliary functions
-namespace AddInfo
-
-def expr (t : TypedExpr) : Expr := t.body
-def thunk (sus : SusValue) : Thunk Value := sus.body
-def get (sus : SusValue) : Value := sus.body.get
-def getTyped (sus : SusValue) : TypedValue := ⟨sus.info, sus.body.get⟩
-def value (val : TypedValue) : Value := val.body
-def sus (val : TypedValue) : SusValue := ⟨val.info, val.body⟩
-
-end AddInfo
+abbrev Args := List TypedValue
 
 def Value.neu (neu : Neutral) : Value := .app neu [] []
 
@@ -192,11 +179,11 @@ def Neutral.ctorName : Neutral → String
 
 namespace Env'
 /-- Stacks a new expression in the environment -/
-def extendWith (env : Env) (thunk : SusValue) : Env :=
-  .mk (thunk :: env.exprs) env.univs
+def extendWith (env : Env) (val : TypedValue) : Env :=
+  .mk (val :: env.exprs) env.univs
 
 /-- Sets a list of expressions to a environment -/
-def withExprs (env : Env) (exprs : List SusValue) : Env :=
+def withExprs (env : Env) (exprs : List TypedValue) : Env :=
   .mk exprs env.univs
 
 end Env'
@@ -205,9 +192,9 @@ end Env'
 def mkConst (f : F) (univs : List Univ) : Value :=
   .neu (.const f univs)
 
-/-- Creates a new variable as a thunk -/
-def mkSusVar (info : TypeInfo) (idx : Nat) : SusValue :=
-  .mk info (.mk fun _ => .neu (.fvar idx))
+/-- Creates a new variable as a typed value -/
+def mkVar (info : TypeInfo) (idx : Nat) : TypedValue :=
+  .mk info $ .neu (.fvar idx)
 
 inductive PrimConstOp
   | natAdd | natMul | natPow | natBeq | natBle | natBlt  | natSucc

Yatima/Typechecker/Equal.lean

@@ -26,16 +26,16 @@ Note: Generally the values are assumed to already have the same type in the func
 namespace Yatima.Typechecker
 
 /-- Reduces the application of a `pi` type to its arguments -/
-def applyType : Value → List SusValue → TypecheckM Value
+def applyType : Value → List TypedValue → TypecheckM Value
   | .pi _ img imgCtx, arg :: args => do
     let res ← withEnv (imgCtx.extendWith arg) (eval img)
     applyType res args
   | type, [] => pure type
   | _, _ => throw "Invalid case for applyType"
 
 mutual
-  partial def tryEtaStruct (lvl : Nat) (term term' : SusValue) : TypecheckM Bool := do
-    match term'.get with
+  partial def tryEtaStruct (lvl : Nat) (term term' : TypedValue) : TypecheckM Bool := do
+    match term'.body with
     | .app (.const k _) args _ =>
       match ← derefTypedConst k with
       | .constructor type .. =>
@@ -45,9 +45,9 @@ mutual
           | .inductive _ struct .. =>
             if struct then
               args.enum.foldlM (init := true) fun acc (i, arg) => do
-                match arg.get with
+                match arg.body with
                 | .app (.proj _ idx val) _ _ =>
-                  pure $ acc && i == idx && (← equal lvl term val.sus)
+                  pure $ acc && i == idx && (← equal lvl term val)
                 | _ => pure false
             else
               pure false
@@ -57,78 +57,75 @@ mutual
     | _ => pure false
 
   /--
-  Checks if two suspended values `term term' : SusValue` at level `lvl : Nat` are equal.
+  Checks if two values `term term' : TypedValue` at level `lvl : Nat` are equal.
 
   It is assumed here that the values are typechecked, have both the same type and their
   original unevaluated terms both lived in the same context.
   -/
-  partial def equal (lvl : Nat) (term term' : SusValue) : TypecheckM Bool :=
+  partial def equal (lvl : Nat) (term term' : TypedValue) : TypecheckM Bool :=
     match term.info, term'.info with
     | .unit, .unit => pure true
     | .proof, .proof => pure true
     | _, _ => do
-      let term! := term.get
-      let term'! := term'.get
-      match term!, term'! with
+      match term.body, term'.body with
       | .lit lit, .lit lit' => pure $ lit == lit'
       | .sort u, .sort u' => pure $ u.equalUniv u'
       | .pi dom img env, .pi dom' img' env' => do
         let res  ← equal lvl dom dom'
-        let img  := suspend img  { ← read with env := env.extendWith  (mkSusVar dom.info  lvl) } (← get)
-        let img' := suspend img' { ← read with env := env'.extendWith (mkSusVar dom'.info lvl) } (← get)
+        let img  ← withNewExtendedEnv env  (mkVar dom.info  lvl) $ evalTyped img
+        let img' ← withNewExtendedEnv env' (mkVar dom'.info lvl) $ evalTyped img'
         let res' ← equal (lvl + 1) img img'
         pure $ res && res'
       | .lam dom bod env, .lam dom' bod' env' => do
         let res  ← equal lvl dom dom'
-        let bod  := suspend bod  { ← read with env := env.extendWith  (mkSusVar dom.info  lvl) } (← get)
-        let bod' := suspend bod' { ← read with env := env'.extendWith (mkSusVar dom'.info lvl) } (← get)
+        let bod  ← withNewExtendedEnv env  (mkVar dom.info  lvl) $ evalTyped bod
+        let bod' ← withNewExtendedEnv env' (mkVar dom'.info lvl) $ evalTyped bod'
         let res' ← equal (lvl + 1) bod bod'
         pure $ res && res'
       | .lam dom bod env, .app neu' args' infos' => do
-        let var := mkSusVar dom.info lvl
-        let bod  := suspend bod { ← read with env := env.extendWith var } (← get)
+        let var := mkVar dom.info lvl
+        let bod  ← withNewExtendedEnv env var $ evalTyped bod
         let app := Value.app neu' (var :: args') (term'.info :: infos')
         equal (lvl + 1) bod (.mk bod.info app)
       | .app neu args infos, .lam dom bod env => do
-        let var := mkSusVar dom.info lvl
-        let bod  := suspend bod { ← read with env := env.extendWith var } (← get)
+        let var := mkVar dom.info lvl
+        let bod  ← withNewExtendedEnv env var $ evalTyped bod
         let app := Value.app neu (var :: args) (term.info :: infos)
         equal (lvl + 1) (.mk bod.info app) bod
       | .app (.fvar idx) args _, .app (.fvar idx') args' _ =>
         if idx == idx' then
           -- If our assumption is correct, i.e., that these values come from terms
           -- in the same environment then their types are equal when their indices
           -- are equal
-          equalThunks lvl args args'
+          equalValues lvl args args'
         else pure false
       | .app (.const k us) args _, .app (.const k' us') args' _ =>
         if k == k' && IR.Univ.equalUnivs us us' then
-          equalThunks lvl args args'
+          equalValues lvl args args'
         else pure false
       | _, .app (.const _ _) _ _ =>
         tryEtaStruct lvl term term'
       | .app (.const _ _) _ _, _ =>
         tryEtaStruct lvl term' term
       | .app (.proj ind idx val) args _, .app (.proj ind' idx' val') args' _ =>
           if ind == ind' && idx == idx' then do
-            let eqVal ← equal lvl val.sus val'.sus
-            let eqThunks ← equalThunks lvl args args'
-            pure (eqVal && eqThunks)
+            let eqVal ← equal lvl val val'
+            let eqVals ← equalValues lvl args args'
+            pure (eqVal && eqVals)
           else pure false
       | .exception e, _ | _, .exception e =>
         throw s!"exception in equal: {e}"
       | _, _ =>
         pure false
 
 /--
-Checks if two list of thunks `vals vals' : List SusValue` are equal by evaluating the thunks
-and checking the evaluated images are equal.
+Checks if two list of values `vals vals' : List TypedValue` are equal
 -/
-  partial def equalThunks (lvl : Nat) (vals vals' : List SusValue) : TypecheckM Bool :=
+  partial def equalValues (lvl : Nat) (vals vals' : List TypedValue) : TypecheckM Bool :=
     match vals, vals' with
     | val::vals, val'::vals' => do
       let eq ← equal lvl val val'
-      let eq' ← equalThunks lvl vals vals'
+      let eq' ← equalValues lvl vals vals'
       pure $ eq && eq'
     | [], [] => pure true
     | _, _ => pure false