From 66de09bc9c705257130faac0ca507426d6711047 Mon Sep 17 00:00:00 2001
From: Cameron Zwarich <cameron@lean-fro.org>
Date: Wed, 9 Jul 2025 17:01:56 -0700
Subject: [PATCH] refactor: use getEnv/modifyEnv more in IR (#9290)

---
 src/Lean/Compiler/IR/CompilerM.lean        | 10 +++++-----
 src/Lean/Compiler/IR/ElimDeadBranches.lean |  8 +++-----
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/src/Lean/Compiler/IR/CompilerM.lean b/src/Lean/Compiler/IR/CompilerM.lean
index 32650013dc..ef1326b510 100644
--- a/src/Lean/Compiler/IR/CompilerM.lean
+++ b/src/Lean/Compiler/IR/CompilerM.lean
@@ -71,6 +71,9 @@ private def logMessageIfAux {α : Type} [ToFormat α] (optName : Name) (a : α)
 @[inline] def modifyEnv (f : Environment → Environment) : CompilerM Unit :=
   modify fun s => { s with env := f s.env }
 
+def getEnv : CompilerM Environment := do
+  let s ← get; pure s.env
+
 abbrev DeclMap := PHashMap Name Decl
 
 private abbrev declLt (a b : Decl) :=
@@ -185,7 +188,7 @@ private def findInterpreterDecl (env : Environment) (declName : Name) : Option D
   | none => declMapExt.getState env |>.find? declName
 
 def findDecl (n : Name) : CompilerM (Option Decl) :=
-  return findEnvDecl (← get).env n
+  return findEnvDecl (← getEnv) n
 
 def containsDecl (n : Name) : CompilerM Bool :=
   return (← findDecl n).isSome
@@ -200,9 +203,6 @@ def addDeclAux (env : Environment) (decl : Decl) : Environment :=
 def getDecls (env : Environment) : List Decl :=
   declMapExt.getEntries env
 
-def getEnv : CompilerM Environment := do
-  let s ← get; pure s.env
-
 def addDecl (decl : Decl) : CompilerM Unit :=
   modifyEnv fun env => declMapExt.addEntry (env.addExtraName decl.name) decl
 
@@ -215,7 +215,7 @@ def findEnvDecl' (env : Environment) (n : Name) (decls : Array Decl) : Option De
   | none      => findEnvDecl env n
 
 def findDecl' (n : Name) (decls : Array Decl) : CompilerM (Option Decl) :=
-  return findEnvDecl' (← get).env n decls
+  return findEnvDecl' (← getEnv) n decls
 
 def containsDecl' (n : Name) (decls : Array Decl) : CompilerM Bool := do
   if decls.any fun decl => decl.name == n then
diff --git a/src/Lean/Compiler/IR/ElimDeadBranches.lean b/src/Lean/Compiler/IR/ElimDeadBranches.lean
index 88a2851b15..6d21c67434 100644
--- a/src/Lean/Compiler/IR/ElimDeadBranches.lean
+++ b/src/Lean/Compiler/IR/ElimDeadBranches.lean
@@ -320,8 +320,7 @@ end UnreachableBranches
 open UnreachableBranches
 
 def elimDeadBranches (decls : Array Decl) : CompilerM (Array Decl) := do
-  let s ← get
-  let env := s.env
+  let env ← getEnv
   let assignments : Array Assignment := decls.map fun _ => {}
   let funVals := mkPArray decls.size Value.bot
   let visitedJps := decls.map fun _ => {}
@@ -330,10 +329,9 @@ def elimDeadBranches (decls : Array Decl) : CompilerM (Array Decl) := do
   let (_, s) := (inferMain ctx).run s
   let funVals := s.funVals
   let assignments := s.assignments
-  modify fun s =>
-    let env := decls.size.fold (init := s.env) fun i _ env =>
+  modifyEnv fun env =>
+    decls.size.fold (init := env) fun i _ env =>
       addFunctionSummary env decls[i].name funVals[i]!
-    { s with env := env }
   return decls.mapIdx fun i decl => elimDead assignments[i]! decl
 
 builtin_initialize registerTraceClass `compiler.ir.elim_dead_branches (inherited := true)