From b6fbdd8679487f986754a553477e8c4130dbfd85 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leonardo@microsoft.com>
Date: Sat, 18 Dec 2021 08:25:56 -0800
Subject: [PATCH] feat: add `Meta.Context.canUnfold?`

---
 src/Lean/Meta/Basic.lean    |  4 ++++
 src/Lean/Meta/GetConst.lean | 22 ++++++++++++++--------
 src/Lean/Meta/WHNF.lean     |  2 +-
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/Lean/Meta/Basic.lean b/src/Lean/Meta/Basic.lean
index 2f45148f05..52f0dd4526 100644
--- a/src/Lean/Meta/Basic.lean
+++ b/src/Lean/Meta/Basic.lean
@@ -184,6 +184,10 @@ structure Context where
     Remark: in the current implementation, `synthPending` fails if `synthPendingDepth > 0`.
     We will add a configuration option if necessary. -/
   synthPendingDepth : Nat                  := 0
+  /--
+    A predicate to control whether a constant can be unfolded or not at `whnf`.
+    Note that we do not cache results at `whnf` when `canUnfold?` is not `none`. -/
+  canUnfold?        : Option (Config → ConstantInfo → CoreM Bool) := none
 
 abbrev MetaM  := ReaderT Context $ StateRefT State CoreM
 
diff --git a/src/Lean/Meta/GetConst.lean b/src/Lean/Meta/GetConst.lean
index a2c07f6600..a863325a2a 100644
--- a/src/Lean/Meta/GetConst.lean
+++ b/src/Lean/Meta/GetConst.lean
@@ -7,23 +7,29 @@ import Lean.Meta.GlobalInstances
 
 namespace Lean.Meta
 
-private def getDefInfo (info : ConstantInfo) : MetaM (Option ConstantInfo) := do
+private def canUnfoldDefault (info : ConstantInfo) : MetaM Bool := do
   match (← read).config.transparency with
-  | TransparencyMode.all => return some info
-  | TransparencyMode.default => return some info
+  | TransparencyMode.all => return true
+  | TransparencyMode.default => return true
   | m =>
     if (← isReducible info.name) then
-      return some info
+      return true
     else if m == TransparencyMode.instances && isGlobalInstance (← getEnv) info.name then
-      return some info
+      return true
     else
-      return none
+      return false
+
+def canUnfold (info : ConstantInfo) : MetaM Bool := do
+  if let some f := (← read).canUnfold? then
+    f (← read).config info
+  else
+    canUnfoldDefault info
 
 def getConst? (constName : Name) : MetaM (Option ConstantInfo) := do
   let env ← getEnv
   match env.find? constName with
   | some (info@(ConstantInfo.thmInfo _))  => getTheoremInfo info
-  | some (info@(ConstantInfo.defnInfo _)) => getDefInfo info
+  | some (info@(ConstantInfo.defnInfo _)) => if (← canUnfold info) then return info else return none
   | some info                             => pure (some info)
   | none                                  => throwUnknownConstant constName
 
@@ -31,7 +37,7 @@ def getConstNoEx? (constName : Name) : MetaM (Option ConstantInfo) := do
   let env ← getEnv
   match env.find? constName with
   | some (info@(ConstantInfo.thmInfo _))  => getTheoremInfo info
-  | some (info@(ConstantInfo.defnInfo _)) => getDefInfo info
+  | some (info@(ConstantInfo.defnInfo _)) => if (← canUnfold info) then return info else return none
   | some info                             => pure (some info)
   | none                                  => pure none
 
diff --git a/src/Lean/Meta/WHNF.lean b/src/Lean/Meta/WHNF.lean
index 2484fc2157..ee0c05eaef 100644
--- a/src/Lean/Meta/WHNF.lean
+++ b/src/Lean/Meta/WHNF.lean
@@ -646,7 +646,7 @@ def reduceNat? (e : Expr) : MetaM (Option Expr) :=
 @[inline] private def useWHNFCache (e : Expr) : MetaM Bool := do
   -- We cache only closed terms without expr metavars.
   -- Potential refinement: cache if `e` is not stuck at a metavariable
-  if e.hasFVar || e.hasExprMVar then
+  if e.hasFVar || e.hasExprMVar || (← read).canUnfold?.isSome then
     return false
   else
     match (← getConfig).transparency with