From afc50a801de4100e2dfa13e00be7b8b0f2f37134 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Fri, 13 Dec 2019 09:52:44 -0800
Subject: [PATCH] feat: add lazy unfolding to field name resolution

---
 src/Init/Lean/Elab/Term.lean      | 24 ++++++++++++++++++++----
 src/Init/Lean/Meta/ExprDefEq.lean |  4 ++--
 src/Init/Lean/Meta/WHNF.lean      |  4 ++--
 3 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/src/Init/Lean/Elab/Term.lean b/src/Init/Lean/Elab/Term.lean
index d05e8db860..31c939cf73 100644
--- a/src/Init/Lean/Elab/Term.lean
+++ b/src/Init/Lean/Elab/Term.lean
@@ -169,7 +169,7 @@ def inferType (ref : Syntax) (e : Expr) : TermElabM Expr := liftMetaM ref $ Meta
 def whnf (ref : Syntax) (e : Expr) : TermElabM Expr := liftMetaM ref $ Meta.whnf e
 def whnfForall (ref : Syntax) (e : Expr) : TermElabM Expr := liftMetaM ref $ Meta.whnfForall e
 def whnfCore (ref : Syntax) (e : Expr) : TermElabM Expr := liftMetaM ref $ Meta.whnfCore e
-def unfoldDefinition (ref : Syntax) (e : Expr) : TermElabM (Option Expr) := liftMetaM ref $ Meta.unfoldDefinition e
+def unfoldDefinition? (ref : Syntax) (e : Expr) : TermElabM (Option Expr) := liftMetaM ref $ Meta.unfoldDefinition? e
 def instantiateMVars (ref : Syntax) (e : Expr) : TermElabM Expr := liftMetaM ref $ Meta.instantiateMVars e
 def isClass (ref : Syntax) (t : Expr) : TermElabM (Option Name) := liftMetaM ref $ Meta.isClass t
 def mkFreshLevelMVar (ref : Syntax) : TermElabM Level := liftMetaM ref $ Meta.mkFreshLevelMVar
@@ -621,7 +621,7 @@ inductive FieldResolution
 private def throwFieldError {α} (ref : Syntax) (e : Expr) (eType : Expr) (msg : MessageData) : TermElabM α :=
 throwError ref $ msg ++ indentExpr e ++ Format.line ++ "has type" ++ indentExpr eType
 
-private def resolveField (ref : Syntax) (e : Expr) (eType : Expr) (field : Field) : TermElabM FieldResolution :=
+private def resolveFieldAux (ref : Syntax) (e : Expr) (eType : Expr) (field : Field) : TermElabM FieldResolution :=
 match eType.getAppFn, field with
 | Expr.const structName _ _, Field.num idx => do
   when (idx == 0) $
@@ -669,10 +669,26 @@ match eType.getAppFn, field with
     searchLCtx ()
 | _, _ => throwFieldError ref e eType "invalid field notation, type is not of the form (C ...) where C is a constant"
 
+private partial def resolveFieldLoop (ref : Syntax) (e : Expr) (field : Field) : Expr → Array Exception → TermElabM FieldResolution
+| eType, previousExceptions => do
+  eType ← whnfCore ref eType;
+  catch (resolveFieldAux ref e eType field)
+    (fun ex => do
+      eType? ← unfoldDefinition? ref eType;
+      match eType? with
+      | some eType => resolveFieldLoop eType (previousExceptions.push ex)
+      | none => do
+        previousExceptions.forM $ fun ex => logMessage ex;
+        throw ex)
+
+private def resolveField (ref : Syntax) (e : Expr) (field : Field) : TermElabM FieldResolution := do
+eType ← inferType ref e;
+resolveFieldLoop ref e field eType #[]
+
 private def elabAppFieldsAux (ref : Syntax) (namedArgs : Array NamedArg) (args : Array Syntax) (expectedType? : Option Expr) (explicit : Bool)
     : Expr → List Field → TermElabM Expr
-| f, []          => elabAppArgs ref f namedArgs args expectedType? explicit
-| f, proj::fields => do
+| f, []            => elabAppArgs ref f namedArgs args expectedType? explicit
+| f, field::fields => do
   fType ← inferType ref f;
   -- TODO
   elabAppArgs ref f namedArgs args expectedType? explicit
diff --git a/src/Init/Lean/Meta/ExprDefEq.lean b/src/Init/Lean/Meta/ExprDefEq.lean
index cfcee76857..6b2af88c1b 100644
--- a/src/Init/Lean/Meta/ExprDefEq.lean
+++ b/src/Init/Lean/Meta/ExprDefEq.lean
@@ -607,7 +607,7 @@ private partial def processAssignmentFOApprox (mvar : Expr) (args : Array Expr)
   trace! `Meta.isDefEq.foApprox (mvar ++ " " ++ args ++ " := " ++ v);
   condM (try $ processAssignmentFOApproxAux mvar args v)
     (pure true)
-    (do v? ← unfoldDefinition v;
+    (do v? ← unfoldDefinition? v;
         match v? with
         | none   => pure false
         | some v => processAssignmentFOApprox v)
@@ -733,7 +733,7 @@ traceCtx `Meta.isDefEq.delta $
 
 /-- Auxiliary method for isDefEqDelta -/
 private abbrev unfold {α} (e : Expr) (failK : MetaM α) (successK : Expr → MetaM α) : MetaM α := do
-e? ← unfoldDefinition e;
+e? ← unfoldDefinition? e;
 match e? with
 | some e => successK e
 | none   => failK
diff --git a/src/Init/Lean/Meta/WHNF.lean b/src/Init/Lean/Meta/WHNF.lean
index cef0c07f8c..9d9d943954 100644
--- a/src/Init/Lean/Meta/WHNF.lean
+++ b/src/Init/Lean/Meta/WHNF.lean
@@ -15,7 +15,7 @@ namespace Meta
 def isAuxDef? (constName : Name) : MetaM Bool := do
 env ← getEnv; pure (isAuxRecursor env constName || isNoConfusion env constName)
 
-def unfoldDefinition (e : Expr) : MetaM (Option Expr)  :=
+def unfoldDefinition? (e : Expr) : MetaM (Option Expr)  :=
 Lean.WHNF.unfoldDefinitionAux getConstNoEx isAuxDef? whnf inferType isExprDefEq synthPending getLocalDecl getExprMVarAssignment e
 
 def whnfCore (e : Expr) : MetaM Expr :=
@@ -24,7 +24,7 @@ Lean.WHNF.whnfCore getConstNoEx isAuxDef? whnf inferType isExprDefEqAux getLocal
 partial def whnfImpl : Expr → MetaM Expr
 | e => Lean.WHNF.whnfEasyCases getLocalDecl getExprMVarAssignment e $ fun e => do
   e ← whnfCore e;
-  e? ← unfoldDefinition e;
+  e? ← unfoldDefinition? e;
   match e? with
   | some e => whnfImpl e
   | none   => pure e