From 24cbd4efbea2e8d4b398b6362941bb87d8cce138 Mon Sep 17 00:00:00 2001
From: Cameron Zwarich <cameron@lean-fro.org>
Date: Mon, 23 Jun 2025 16:03:10 -0700
Subject: [PATCH] fix: correctly handle never_extract attribute in LCNF CSE
 (#8952)

This PR fixes the handling of the `never_extract` attribute in the
compiler's CSE pass. There is an interesting debate to be had about
exactly how hard the compiler should try to avoid duplicating anything
that transitively uses `never_extract`, but this is the simplest form
and roughly matches the check in the old compiler (although due to
different handling of local function decls in the two compilers, the
consequences might be slightly different).

This gets half of the way to #8944.
---
 src/Lean/Compiler/LCNF/CSE.lean               | 27 +++++++++++++------
 tests/compiler/never_extract.lean             | 13 +++++++++
 .../compiler/never_extract.lean.expected.out  |  3 +++
 3 files changed, 35 insertions(+), 8 deletions(-)
 create mode 100644 tests/compiler/never_extract.lean
 create mode 100644 tests/compiler/never_extract.lean.expected.out

diff --git a/src/Lean/Compiler/LCNF/CSE.lean b/src/Lean/Compiler/LCNF/CSE.lean
index b49d260460..bf0cae5557 100644
--- a/src/Lean/Compiler/LCNF/CSE.lean
+++ b/src/Lean/Compiler/LCNF/CSE.lean
@@ -7,6 +7,7 @@ prelude
 import Lean.Compiler.LCNF.CompilerM
 import Lean.Compiler.LCNF.ToExpr
 import Lean.Compiler.LCNF.PassManager
+import Lean.Compiler.NeverExtractAttr
 
 namespace Lean.Compiler.LCNF
 
@@ -44,6 +45,13 @@ def replaceFun (decl : FunDecl) (fvarId : FVarId) : M Unit := do
   eraseFunDecl decl
   addFVarSubst decl.fvarId fvarId
 
+def hasNeverExtract (v : LetValue) : CompilerM Bool :=
+  match v with
+  | .const declName .. =>
+    return hasNeverExtractAttribute (← getEnv) declName
+  | .lit _ | .erased | .proj .. | .fvar .. =>
+    return false
+
 partial def _root_.Lean.Compiler.LCNF.Code.cse (shouldElimFunDecls : Bool) (code : Code) : CompilerM Code :=
   go code |>.run' {}
 where
@@ -57,15 +65,18 @@ where
     match code with
     | .let decl k =>
       let decl ← normLetDecl decl
-      -- We only apply CSE to pure code
-      let key := decl.value.toExpr
-      match (← get).map.find? key with
-      | some fvarId =>
-        replaceLet decl fvarId
-        go k
-      | none =>
-        addEntry key decl.fvarId
+      if (← hasNeverExtract decl.value) then
         return code.updateLet! decl (← go k)
+      else
+        -- We only apply CSE to pure code
+        let key := decl.value.toExpr
+        match (← get).map.find? key with
+        | some fvarId =>
+          replaceLet decl fvarId
+          go k
+        | none =>
+          addEntry key decl.fvarId
+          return code.updateLet! decl (← go k)
     | .fun decl k =>
       let decl ← goFunDecl decl
       if shouldElimFunDecls then
diff --git a/tests/compiler/never_extract.lean b/tests/compiler/never_extract.lean
new file mode 100644
index 0000000000..de16a9291d
--- /dev/null
+++ b/tests/compiler/never_extract.lean
@@ -0,0 +1,13 @@
+def test1 (a : Nat) : Nat :=
+  let f a :=
+    dbg_trace s!"{a}"
+    a
+  let g a :=
+    dbg_trace s!"{a + 0}"
+    a
+  (f a) + (g a)
+
+def main : IO Unit :=
+  -- Use `eprintln` because that is what `dbg_trace` uses.
+  IO.eprintln f!"{test1 1}"
+
diff --git a/tests/compiler/never_extract.lean.expected.out b/tests/compiler/never_extract.lean.expected.out
new file mode 100644
index 0000000000..33280629d4
--- /dev/null
+++ b/tests/compiler/never_extract.lean.expected.out
@@ -0,0 +1,3 @@
+1
+1
+2