refactor: ensure that unfoldDefinitionAux is not specialized multiple times

library/Init/Lean/Meta/ExprDefEq.lean

@@ -610,7 +610,10 @@ private partial def processAssignmentFOApprox (mvar : Expr) (args : Array Expr)
   trace! `Meta.isDefEq.foApprox (mvar ++ " " ++ args ++ " := " ++ v);
   condM (try $ processAssignmentFOApproxAux mvar args v)
     (pure true)
-    (unfoldDefinitionAux v (pure false) processAssignmentFOApprox)
+    (do v? ← unfoldDefinition v;
+        match v? with
+        | none   => pure false
+        | some v => processAssignmentFOApprox v)
 
 private partial def simpAssignmentArgAux : Expr → MetaM Expr
 | Expr.mdata _ e _       => simpAssignmentArgAux e
@@ -733,7 +736,11 @@ traceCtx `Meta.isDefEq.delta $
     isListLevelDefEqAux tFn.constLevels! sFn.constLevels!
 
 /-- Auxiliary method for isDefEqDelta -/
-private abbrev unfold := @unfoldDefinitionAux
+private abbrev unfold {α} (e : Expr) (failK : MetaM α) (successK : Expr → MetaM α) : MetaM α :=
+do e? ← unfoldDefinition e;
+   match e? with
+   | some e => successK e
+   | none   => failK
 
 /-- Auxiliary method for isDefEqDelta -/
 private def unfoldBothDefEq (fn : Name) (t s : Expr) : MetaM LBool :=
@@ -891,9 +898,6 @@ do tType ← inferType t;
      (do sType ← inferType s; toLBoolM $ isExprDefEqAux tType sType)
      (pure LBool.undef)
 
-private def whnfCoreAux (e : Expr) : MetaM Expr :=
-Lean.whnfCore getConstNoEx isAuxDef? whnf inferType isExprDefEqAux getLocalDecl getExprMVarAssignment e
-
 @[inline] def tryL (x : MetaM LBool) (k : MetaM Bool) : MetaM Bool :=
 do status ← x;
    match status with
@@ -904,8 +908,8 @@ do status ← x;
 @[specialize] private partial def isDefEqWHNF
     (t s : Expr)
     (k : Expr → Expr → MetaM Bool) : MetaM Bool :=
-do t' ← whnfCoreAux t;
-   s' ← whnfCoreAux s;
+do t' ← whnfCore t;
+   s' ← whnfCore s;
    if t == t' && s == s' then
      k t' s'
    else
@@ -914,7 +918,7 @@ do t' ← whnfCoreAux t;
 @[specialize] private def unstuckMVar
     (e : Expr)
     (successK : Expr → MetaM Bool) (failK : MetaM Bool): MetaM Bool :=
-do s? ← getStuckMVar getConst whnf e;
+do s? ← WHNF.getStuckMVar getConst whnf e;
    match s? with
    | some s =>
      condM (synthPending s)

library/Init/Lean/Meta/WHNF.lean

@@ -15,16 +15,19 @@ namespace Meta
 def isAuxDef? (constName : Name) : MetaM Bool :=
 do env ← getEnv; pure (isAuxRecursor env constName || isNoConfusion env constName)
 
-@[specialize] def unfoldDefinitionAux {α}
-    (e : Expr) (failK : MetaM α) (successK : Expr → MetaM α) : MetaM α :=
-Lean.unfoldDefinitionAux getConstNoEx isAuxDef? whnf inferType isExprDefEq synthPending getLocalDecl
-  getExprMVarAssignment e (fun _ => failK) successK
+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 :=
+Lean.WHNF.whnfCore getConstNoEx isAuxDef? whnf inferType isExprDefEqAux getLocalDecl getExprMVarAssignment e
 
 partial def whnfImpl : Expr → MetaM Expr
-| e => whnfEasyCases getLocalDecl getExprMVarAssignment e $ fun e => do
-  e ← whnfCore getConstNoEx isAuxDef? whnfImpl inferType isExprDefEqAux getLocalDecl getExprMVarAssignment e;
-  Lean.unfoldDefinitionAux getConstNoEx isAuxDef? whnf inferType isExprDefEq synthPending getLocalDecl
-    getExprMVarAssignment e (fun _ => pure e) whnfImpl
+| e => Lean.WHNF.whnfEasyCases getLocalDecl getExprMVarAssignment e $ fun e => do
+  e ← whnfCore e;
+  e? ← unfoldDefinition e;
+  match e? with
+  | some e => whnfImpl e
+  | none   => pure e
 
 @[init] def setWHNFRef : IO Unit :=
 whnfRef.set whnfImpl

library/Init/Lean/WHNF.lean

@@ -8,6 +8,7 @@ import Init.Lean.Declaration
 import Init.Lean.LocalContext
 
 namespace Lean
+namespace WHNF
 /- ===========================
    Smart unfolding support
    =========================== -/
@@ -330,7 +331,7 @@ else
   if succeeded then whnfCoreUnstuck e else pure e
 
 /-- Unfold definition using "smart unfolding" if possible. -/
-@[specialize] def unfoldDefinitionAux {α} {m : Type → Type} [Monad m]
+@[specialize] def unfoldDefinitionAux {m : Type → Type} [Monad m]
     (getConst          : Name → m (Option ConstantInfo))
     (isAuxDef?         : Name → m Bool)
     (whnf              : Expr → m Expr)
@@ -339,27 +340,30 @@ else
     (synthesizePending : Expr → m Bool)
     (getLocalDecl      : Name → m LocalDecl)
     (getMVarAssignment : Name → m (Option Expr))
-    (e : Expr)
-    (failK : Unit → m α) (successK : Expr → m α) : m α :=
+    (e : Expr) : m (Option Expr) :=
 match e with
 | Expr.app f _ _ =>
-  matchConstAux getConst f.getAppFn failK $ fun fInfo fLvls =>
-    if fInfo.lparams.length != fLvls.length then failK ()
+  matchConstAux getConst f.getAppFn (fun _ => pure none) $ fun fInfo fLvls =>
+    if fInfo.lparams.length != fLvls.length then pure none
     else do
       fAuxInfo? ← getConst (mkSmartUnfoldingNameFor fInfo.name);
       match fAuxInfo? with
       | some $ fAuxInfo@(ConstantInfo.defnInfo _) =>
-        deltaBetaDefinition fAuxInfo fLvls e.getAppRevArgs failK $ fun e₁ => do
+        deltaBetaDefinition fAuxInfo fLvls e.getAppRevArgs (fun _ => pure none) $ fun e₁ => do
           e₂ ← whnfCoreUnstuck getConst isAuxDef? whnf inferType isDefEq synthesizePending getLocalDecl getMVarAssignment e₁;
           if isIdRhsApp e₂ then
-            successK $ extractIdRhs e₂
+            pure (some (extractIdRhs e₂))
           else
-            failK ()
-      | _ => if fInfo.hasValue then deltaBetaDefinition fInfo fLvls e.getAppRevArgs failK successK else failK ()
+            pure none
+      | _ =>
+        if fInfo.hasValue then
+          deltaBetaDefinition fInfo fLvls e.getAppRevArgs (fun _ => pure none) (fun e => pure (some e))
+        else
+          pure none
 | Expr.const name lvls _ => do
-  (some (cinfo@(ConstantInfo.defnInfo _))) ← getConst name | failK ();
-  deltaDefinition cinfo lvls failK successK
-| _ => failK ()
+  (some (cinfo@(ConstantInfo.defnInfo _))) ← getConst name | pure none;
+  deltaDefinition cinfo lvls (fun _ => pure none) (fun e => pure (some e))
+| _ => pure none
 
 /- Reference implementation for `whnf`. It does not cache any results.
 
@@ -381,6 +385,10 @@ match e with
     : Expr → m Expr
 | e => do
   e ← whnfCore getConst isAuxDef? whnfMain inferType isDefEq getLocalDecl getMVarAssignment e;
-  unfoldDefinitionAux getConst isAuxDef? whnfMain inferType isDefEq synthesizePending getLocalDecl getMVarAssignment e (fun _ => pure e) whnfMain
+  e? ← unfoldDefinitionAux getConst isAuxDef? whnfMain inferType isDefEq synthesizePending getLocalDecl getMVarAssignment e;
+  match e? with
+  | some e => whnfMain e
+  | none   => pure e
 
+end WHNF
 end Lean

tests/lean/run/meta2.lean

@@ -5,7 +5,7 @@ open Lean.Meta
 def dbgOpt : Options :=
 let opt : Options := {};
 let opt := opt.setBool `trace.Meta true;
--- let opt := opt.setBool `trace.Meta.check false;
+let opt := opt.setBool `trace.Meta.check false;
 opt
 
 def print (msg : MessageData) : MetaM Unit :=
@@ -24,9 +24,12 @@ do env ← importModules $ mods.map $ fun m => {module := m};
      s.traceState.traces.forM $ fun m => IO.println $ format m;
      throw (IO.userError (toString err))
 
+def nat  := mkConst `Nat
+def succ := mkConst `Nat.succ
+def add  := mkConst `Nat.add
+
 def tst1 : MetaM Unit :=
 do print "----- tst1 -----";
-   let nat := mkConst `Nat;
    mvar ← mkFreshExprMVar nat;
    check $ isExprDefEq mvar (mkNatLit 10);
    check $ isExprDefEq mvar (mkNatLit 10);
@@ -36,8 +39,6 @@ do print "----- tst1 -----";
 
 def tst2 : MetaM Unit :=
 do print "----- tst2 -----";
-   let nat  := mkConst `Nat;
-   let succ := mkConst `Nat.succ;
    mvar ← mkFreshExprMVar nat;
    check $ isExprDefEq (mkApp succ mvar) (mkApp succ (mkNatLit 10));
    check $ isExprDefEq mvar (mkNatLit 10);
@@ -47,8 +48,6 @@ do print "----- tst2 -----";
 
 def tst3 : MetaM Unit :=
 do print "----- tst3 -----";
-   let nat := mkConst `Nat;
-   let add := mkConst `Nat.add;
    let t   := mkLambda `x BinderInfo.default nat $ mkBVar 0;
    mvar ← mkFreshExprMVar (mkForall `x BinderInfo.default nat nat);
    lambdaTelescope t $ fun xs _ => do {
@@ -64,8 +63,6 @@ do print "----- tst3 -----";
 
 def tst4 : MetaM Unit :=
 do print "----- tst4 -----";
-   let nat := mkConst `Nat;
-   let add := mkConst `Nat.add;
    let t   := mkLambda `x BinderInfo.default nat $ mkBVar 0;
    lambdaTelescope t $ fun xs _ => do {
      let x := xs.get! 0;
@@ -80,3 +77,60 @@ do print "----- tst4 -----";
    pure ()
 
 #eval run [`Init.Data.Nat] tst4
+
+def mkProd (a b : Expr) : MetaM Expr :=
+do u ← getLevel a;
+   v ← getLevel b;
+   let r := mkAppN (mkConst `Prod [u.dec.getD u, v.dec.getD v]) #[a, b];
+   check r;
+   pure r
+
+def mkPair (a b : Expr) : MetaM Expr :=
+do aType ← inferType a;
+   bType ← inferType b;
+   u ← getLevel aType;
+   v ← getLevel bType;
+   let r := mkAppN (mkConst `Prod.mk [u.dec.getD u, v.dec.getD v]) #[aType, bType, a, b];
+   check r;
+   pure r
+
+def mkFst (s : Expr) : MetaM Expr :=
+do sType ← inferType s;
+   sType ← whnfUsingDefault sType;
+   unless (sType.isAppOfArity `Prod 2) $ throw $ Exception.other "product expected";
+   let lvls := sType.getAppFn.constLevels!;
+   let r := mkAppN (mkConst `Prod.fst lvls) #[sType.getArg! 0, sType.getArg! 1, s];
+   check r;
+   pure r
+
+def mkSnd (s : Expr) : MetaM Expr :=
+do sType ← inferType s;
+   sType ← whnfUsingDefault sType;
+   unless (sType.isAppOfArity `Prod 2) $ throw $ Exception.other "product expected";
+   let lvls := sType.getAppFn.constLevels!;
+   let r := mkAppN (mkConst `Prod.snd lvls) #[sType.getArg! 0, sType.getArg! 1, s];
+   check r;
+   pure r
+
+def mkId (a : Expr) : MetaM Expr :=
+do aType ← inferType a;
+   lvl   ← getLevel aType;
+   let r := mkAppN (mkConst `id [lvl]) #[aType, a];
+   check r;
+   pure r
+
+#print id
+
+def tst5 : MetaM Unit :=
+do print "----- tst5 -----";
+   p₁ ← mkPair (mkNatLit 1) (mkNatLit 2);
+   x  ← mkFst p₁;
+   x  ← mkId x;
+   print x;
+   prod ← mkProd nat nat;
+   m ← mkFreshExprMVar prod;
+   y ← mkFst m;
+   check $ isExprDefEq y x;
+   print y
+
+#eval run [`Init.Data.Nat] tst5