From 23bae264fd1fb3e2de50a1af2c65b473ac25b514 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 9 Jul 2022 20:09:15 -0700
Subject: [PATCH] perf: add cache for `check (e : Expr) : MetaM Unit`

Address one of the performance problems exposed by #1287
---
 src/Lean/Meta/Check.lean | 39 ++++++++++++++++++++++-----------------
 1 file changed, 22 insertions(+), 17 deletions(-)

diff --git a/src/Lean/Meta/Check.lean b/src/Lean/Meta/Check.lean
index 5e9d65b909..ae46897d14 100644
--- a/src/Lean/Meta/Check.lean
+++ b/src/Lean/Meta/Check.lean
@@ -143,40 +143,45 @@ def checkApp (f a : Expr) : MetaM Unit := do
       throwAppTypeMismatch f a
   | _ => throwFunctionExpected (mkApp f a)
 
-private partial def checkAux : Expr → MetaM Unit
-  | e@(Expr.forallE ..)  => checkForall e
-  | e@(Expr.lam ..)      => checkLambdaLet e
-  | e@(Expr.letE ..)     => checkLambdaLet e
-  | Expr.const c lvls _  => checkConstant c lvls
-  | Expr.app f a _       => do checkAux f; checkAux a; checkApp f a
-  | Expr.mdata _ e _     => checkAux e
-  | Expr.proj _ _ e _    => checkAux e
-  | _                    => pure ()
+private partial def checkAux (e : Expr) : MetaM Unit := do
+  check e |>.run
 where
-  checkLambdaLet (e : Expr) : MetaM Unit :=
+  check (e : Expr) : MonadCacheT ExprStructEq Unit MetaM Unit :=
+    checkCache { val := e : ExprStructEq } fun _ => do
+      match e with
+      | .forallE ..      => checkForall e
+      | .lam ..          => checkLambdaLet e
+      | .letE ..         => checkLambdaLet e
+      | .const c lvls _  => checkConstant c lvls
+      | .app f a _       => check f; check a; checkApp f a
+      | .mdata _ e _     => check e
+      | .proj _ _ e _    => check e
+      | _                => return ()
+
+  checkLambdaLet (e : Expr) : MonadCacheT ExprStructEq Unit MetaM Unit :=
     lambdaLetTelescope e fun xs b => do
       xs.forM fun x => do
         let xDecl ← getFVarLocalDecl x;
         match xDecl with
         | LocalDecl.cdecl (type := t) .. =>
           ensureType t
-          checkAux t
+          check t
         | LocalDecl.ldecl (type := t) (value := v) .. =>
           ensureType t
-          checkAux t
+          check t
           let vType ← inferType v
           unless (← isDefEq t vType) do throwLetTypeMismatchMessage x.fvarId!
-          checkAux v
-      checkAux b
+          check v
+      check b
 
-  checkForall (e : Expr) : MetaM Unit :=
+  checkForall (e : Expr) : MonadCacheT ExprStructEq Unit MetaM Unit :=
     forallTelescope e fun xs b => do
       xs.forM fun x => do
         let xDecl ← getFVarLocalDecl x
         ensureType xDecl.type
-        checkAux xDecl.type
+        check xDecl.type
       ensureType b
-      checkAux b
+      check b
 
 def check (e : Expr) : MetaM Unit :=
   traceCtx `Meta.check do