From 9151fef49d913677682edba2b79b6960fa250189 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Tue, 8 Sep 2020 12:28:04 -0700
Subject: [PATCH] feat: abstract proofs occurring in binders

---
 src/Lean/LocalContext.lean              | 19 +++++++++++++++++--
 src/Lean/Meta/AbstractNestedProofs.lean | 25 +++++++++++++++++++++----
 tests/lean/precissues.lean.expected.out |  2 +-
 tests/lean/run/newfrontend3.lean        | 10 ++++++++++
 4 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/src/Lean/LocalContext.lean b/src/Lean/LocalContext.lean
index 66ee4d5b79..83c6c76a54 100644
--- a/src/Lean/LocalContext.lean
+++ b/src/Lean/LocalContext.lean
@@ -51,6 +51,10 @@ def type : LocalDecl → Expr
 | cdecl _ _ _ t _   => t
 | ldecl _ _ _ t _ _ => t
 
+def setType : LocalDecl → Expr → LocalDecl
+| cdecl idx id n _ bi, t   => cdecl idx id n t bi
+| ldecl idx id n _ v nd, t => ldecl idx id n t v nd
+
 def binderInfo : LocalDecl → BinderInfo
 | cdecl _ _ _ _ bi  => bi
 | ldecl _ _ _ _ _ _ => BinderInfo.default
@@ -63,6 +67,10 @@ def value : LocalDecl → Expr
 | cdecl _ _ _ _ _   => panic! "let declaration expected"
 | ldecl _ _ _ _ v _ => v
 
+def setValue : LocalDecl → Expr → LocalDecl
+| ldecl idx id n t _ nd, v => ldecl idx id n t v nd
+| d, _                     => d
+
 def updateUserName : LocalDecl → Name → LocalDecl
 | cdecl index id _ type bi,     userName => cdecl index id userName type bi
 | ldecl index id _ type val nd, userName => ldecl index id userName type val nd
@@ -214,16 +222,23 @@ match lctx with
     { fvarIdToDecl := map.insert decl.fvarId decl,
       decls        := decls.set decl.index decl }
 
-def updateBinderInfo (lctx : LocalContext) (fvarId : FVarId) (bi : BinderInfo) : LocalContext :=
+/--
+  Low-level function for updating the local context.
+  Assumptions about `f`, the resulting nested expressions must be definitionally equal to their original values,
+  the `index` nor `fvarId` are modified.  -/
+@[inline] def modifyLocalDecl (lctx : LocalContext) (fvarId : FVarId) (f : LocalDecl → LocalDecl) : LocalContext :=
 match lctx with
 | { fvarIdToDecl := map, decls := decls } =>
   match lctx.find? fvarId with
   | none      => lctx
   | some decl =>
-    let decl := decl.updateBinderInfo bi;
+    let decl := f decl;
     { fvarIdToDecl := map.insert decl.fvarId decl,
       decls        := decls.set decl.index decl }
 
+def updateBinderInfo (lctx : LocalContext) (fvarId : FVarId) (bi : BinderInfo) : LocalContext :=
+modifyLocalDecl lctx fvarId fun decl => decl.updateBinderInfo bi
+
 @[export lean_local_ctx_num_indices]
 def numIndices (lctx : LocalContext) : Nat :=
 lctx.decls.size
diff --git a/src/Lean/Meta/AbstractNestedProofs.lean b/src/Lean/Meta/AbstractNestedProofs.lean
index f4db708fd3..e543bcbd6e 100644
--- a/src/Lean/Meta/AbstractNestedProofs.lean
+++ b/src/Lean/Meta/AbstractNestedProofs.lean
@@ -31,10 +31,27 @@ mkAuxDefinitionFor lemmaName e
 partial def visit : Expr → M Expr
 | e =>
   if e.isAtomic then pure e
-  else checkCache e fun e => condM (liftM $ isNonTrivialProof e) (mkAuxLemma e) $ match e with
-    | Expr.lam _ _ _ _     => lambdaLetTelescope e fun xs b => do b ← visit b; mkLambdaFVars xs b
-    | Expr.letE _ _ _ _ _  => lambdaLetTelescope e fun xs b => do b ← visit b; mkLambdaFVars xs b
-    | Expr.forallE _ _ _ _ => forallTelescope e fun xs b => do b ← visit b; mkForallFVars xs b
+  else do
+    let visitBinders (xs : Array Expr) (k : M Expr) : M Expr := do {
+      localInstances ← getLocalInstances;
+      lctx ← getLCtx;
+      lctx ← xs.foldlM
+        (fun (lctx : LocalContext) x => do
+          let xFVarId := x.fvarId!;
+          localDecl ← getLocalDecl xFVarId;
+          type      ← visit localDecl.type;
+          let localDecl := localDecl.setType type;
+          localDecl ← match localDecl.value? with
+            | some value => do value ← visit value; pure $ localDecl.setValue value
+            | none       => pure localDecl;
+          pure $ lctx.modifyLocalDecl xFVarId fun _ => localDecl)
+        lctx;
+      withLCtx lctx localInstances k
+    };
+    checkCache e fun e => condM (liftM $ isNonTrivialProof e) (mkAuxLemma e) $ match e with
+    | Expr.lam _ _ _ _     => lambdaLetTelescope e fun xs b => visitBinders xs do b ← visit b; mkLambdaFVars xs b
+    | Expr.letE _ _ _ _ _  => lambdaLetTelescope e fun xs b => visitBinders xs do b ← visit b; mkLambdaFVars xs b
+    | Expr.forallE _ _ _ _ => forallTelescope e fun xs b => visitBinders xs do b ← visit b; mkForallFVars xs b
     | Expr.mdata _ b _     => do b ← visit b; pure $ e.updateMData! b
     | Expr.proj _ _ b _    => do b ← visit b; pure $ e.updateProj! b
     | Expr.app _ _ _       => e.withApp fun f args => do args ← args.mapM visit; pure $ mkAppN f args
diff --git a/tests/lean/precissues.lean.expected.out b/tests/lean/precissues.lean.expected.out
index 8d95783c43..d59b634240 100644
--- a/tests/lean/precissues.lean.expected.out
+++ b/tests/lean/precissues.lean.expected.out
@@ -3,7 +3,7 @@ id (fun (x : ?m.4) => x) : ?m.4 → ?m.4
 f 1 (fun (x : Nat) => x) : Nat
 0 : Nat
 f 1 (fun (x : Nat) => x) : Nat
-id : ?m.90 → ?m.90
+id : ?m.91 → ?m.91
 precissues.lean:15:10: error: expected command
 1 : Nat
 id ((fun (this : True) => this) True.intro) : True
diff --git a/tests/lean/run/newfrontend3.lean b/tests/lean/run/newfrontend3.lean
index c3390c3d67..11a86ff3bd 100644
--- a/tests/lean/run/newfrontend3.lean
+++ b/tests/lean/run/newfrontend3.lean
@@ -14,3 +14,13 @@ by {
   intro;
   assumption
 }
+
+def g (i j k : Nat) (a : Array Nat) (h₁ : i < k) (h₂ : k < j) (h₃ : j < a.size) : Nat :=
+let vj := a.get ⟨j, h₃⟩;
+let vi := a.get ⟨i, Nat.ltTrans h₁ (Nat.ltTrans h₂ h₃)⟩;
+vi + vj
+
+set_option pp.all true in
+#print g
+
+#check g.proof_1