withCompilerModIdx ~> findExtEntry?

Author: Kha

src/Lean/Compiler/IR/CompilerM.lean

@@ -145,14 +145,7 @@ private def exportIREntries (env : Environment) : Array (Name × Array EnvExtens
     (Lean.regularInitAttr.ext.name, initDecls)]
 
 def findEnvDecl (env : Environment) (declName : Name) : Option Decl :=
-  Lean.Compiler.LCNF.withCompilerModIdx env declName
-    (fun modIdx =>
-      -- `meta import/import all`
-      guard (env.header.modules[modIdx]?.any (·.irPhases != .runtime)) *>
-      findAtSorted? (declMapExt.getModuleIREntries env modIdx) declName <|>
-      -- (closure of) `meta def`; will report `.extern`s for other `def`s so needs to come second
-      findAtSorted? (declMapExt.getModuleEntries env modIdx) declName)
-    (fun _ => declMapExt.getState env |>.find? declName)
+  Lean.Compiler.LCNF.findExtEntry? env declMapExt declName findAtSorted? (·.2.find?)
 
 @[export lean_ir_find_env_decl]
 private def findInterpDecl (env : Environment) (declName : Name) : Option Decl :=

src/Lean/Compiler/IR/ElimDeadBranches.lean

@@ -134,9 +134,7 @@ def addFunctionSummary (env : Environment) (fid : FunId) (v : Value) : Environme
   functionSummariesExt.addEntry env (fid, v)
 
 def getFunctionSummary? (env : Environment) (fid : FunId) : Option Value :=
-  Lean.Compiler.LCNF.withCompilerModIdx env fid
-    (fun modIdx => findAtSorted? (functionSummariesExt.getModuleEntries env modIdx) fid)
-    (fun _ => functionSummariesExt.getState env |>.find? fid)
+  Lean.Compiler.LCNF.findExtEntry? env functionSummariesExt fid findAtSorted? (·.2.find?)
 
 abbrev Assignment := Std.HashMap VarId Value
 

src/Lean/Compiler/LCNF/Basic.lean

@@ -757,10 +757,11 @@ def instantiateRevRangeArgs (e : Expr) (beginIdx endIdx : Nat) (args : Array Arg
   else
     e.instantiateRevRange beginIdx endIdx (args.map (·.toExpr))
 
-def withCompilerModIdx (env : Environment) (declName : Name)
-    (importedOrLocal : ModuleIdx → Option α) («local» : Unit → Option α) : Option α :=
-  match env.getModuleIdxFor? declName with
-  | some modIdx => importedOrLocal modIdx <|> «local» ()
-  | none        => «local» ()
+def findExtEntry? [Inhabited σ] (env : Environment) (ext : PersistentEnvExtension α β σ) (declName : Name)
+    (findAtSorted? : Array α → Name → Option α')
+    (findInState? : σ → Name → Option α') : Option α' :=
+  (env.getModuleIdxFor? declName).bind (fun modIdx =>
+    findAtSorted? (ext.getModuleIREntries env modIdx) declName <|> findAtSorted? (ext.getModuleEntries env modIdx) declName)
+  <|> findInState? (ext.getState env) declName
 
 end Lean.Compiler.LCNF

src/Lean/Compiler/LCNF/ElimDeadBranches.lean

@@ -265,9 +265,7 @@ def addFunctionSummary (env : Environment) (fid : Name) (v : Value) : Environmen
 Obtain the `Value` for a function name if possible.
 -/
 def getFunctionSummary? (env : Environment) (fid : Name) : Option Value :=
-  withCompilerModIdx env fid
-    (fun modIdx => findAtSorted? (functionSummariesExt.getModuleEntries env modIdx) fid)
-    (fun _ => functionSummariesExt.getState env |>.find? fid)
+  findExtEntry? env functionSummariesExt fid findAtSorted? (·.2.find?)
 
 /--
 A map from variable identifiers to the `Value` produced by the abstract

src/Lean/Compiler/LCNF/PhaseExt.lean

@@ -132,9 +132,7 @@ builtin_initialize baseExt : DeclExt ← mkDeclExt .base
 builtin_initialize monoExt : DeclExt ← mkDeclExt .mono
 
 def getDeclCore? (env : Environment) (ext : DeclExt) (declName : Name) : Option Decl :=
-  withCompilerModIdx env declName
-    (fun modIdx => findAtSorted? (ext.getModuleEntries env modIdx) declName)
-    (fun _ => ext.getState env |>.find? declName)
+  findExtEntry? env ext declName findAtSorted? (·.find?)
 
 def getBaseDecl? (declName : Name) : CoreM (Option Decl) := do
   return getDeclCore? (← getEnv) baseExt declName

src/Lean/Compiler/LCNF/SpecInfo.lean

@@ -206,9 +206,7 @@ def saveSpecParamInfo (decls : Array Decl) : CompilerM Unit := do
         modifyEnv fun env => specExtension.addEntry env { declName := decl.name, paramsInfo }
 
 def getSpecParamInfoCore? (env : Environment) (declName : Name) : Option (Array SpecParamInfo) :=
-  withCompilerModIdx env declName
-    (fun modIdx => findAtSorted? (specExtension.getModuleEntries env modIdx) declName |>.map (·.paramsInfo))
-    (fun _ => (specExtension.getState env).specInfo.find? declName)
+  findExtEntry? env specExtension declName (findAtSorted? · · |>.map (·.paramsInfo)) (·.2.specInfo.find?)
 
 def getSpecParamInfo? [Monad m] [MonadEnv m] (declName : Name) : m (Option (Array SpecParamInfo)) :=
   return getSpecParamInfoCore? (← getEnv) declName