feat: shake: track [csimp] uses (#12351)

This PR extends the `@[csimp]` attribute to be correctly tracked by
`lake shake`

src/Lean/Compiler/CSimpAttr.lean

@@ -8,6 +8,7 @@ module
 prelude
 public import Lean.ScopedEnvExtension
 public import Lean.Util.Recognizers
+import Lean.ExtraModUses
 
 public section
 
@@ -21,7 +22,8 @@ structure Entry where
   deriving Inhabited
 
 structure State where
-  map : SMap Name Name := {}
+  /-- Map from `e.fromDeclName` to `e` -/
+  map : SMap Name Entry := {}
   thmNames : SSet Name := {}
   deriving Inhabited
 
@@ -31,7 +33,7 @@ def State.switch : State → State
 builtin_initialize ext : SimpleScopedEnvExtension Entry State ←
   registerSimpleScopedEnvExtension {
     initial        := {}
-    addEntry       := fun { map, thmNames } { fromDeclName, toDeclName, thmName } => { map := map.insert fromDeclName toDeclName, thmNames := thmNames.insert thmName }
+    addEntry       := fun { map, thmNames } e => { map := map.insert e.fromDeclName e, thmNames := thmNames.insert e.thmName }
     finalizeImport := fun s => s.switch
   }
 
@@ -77,15 +79,21 @@ private def initFn :=
       discard <| add declName attrKind
   }
 
-def replaceConstants (env : Environment) (e : Expr) : Expr :=
+/-- If `e` (as a whole) matches a `[csimp]` theorem, returns the replacement expression. -/
+def replaceConstant? (env : Environment) (e : Expr) : CoreM (Option Expr) := do
   let s := ext.getState env
-  e.replace fun e =>
-    if e.isConst then
-      match s.map.find? e.constName! with
-      | some declNameNew => some (mkConst declNameNew e.constLevels!)
-      | none => none
-    else
-      none
+  let .const declName _ := e | return none
+  if let some ent := s.map.find? declName then
+    withoutExporting <| recordExtraModUseFromDecl (isMeta := false) ent.thmName
+    return some (mkConst ent.toDeclName e.constLevels!)
+  else
+    return none
+
+/--
+If `e` (as a whole) matches a `[csimp]` theorem, returns the replacement expression, or else `e`.
+-/
+def replaceConstant (env : Environment) (e : Expr) : CoreM Expr :=
+  return (← replaceConstant? env e).getD e
 
 end CSimp
 

src/Lean/Compiler/LCNF/ToLCNF.lean

@@ -536,7 +536,7 @@ where
   /-- Giving `f` a constant `.const declName us`, convert `args` into `args'`, and return `.const declName us args'` -/
   visitAppDefaultConst (f : Expr) (args : Array Expr) : M (Arg .pure) := do
     let env ← getEnv
-    let .const declName us := CSimp.replaceConstants env f | unreachable!
+    let .const declName us ← CSimp.replaceConstant env f | unreachable!
     let args ← args.mapM visitAppArg
     if hasNeverExtractAttribute env declName then
       modify fun s => {s with shouldCache := false }
@@ -667,7 +667,7 @@ where
       let f := e.getAppFn
       let args := e.getAppArgs
       let env ← getEnv
-      let .const declName us := CSimp.replaceConstants env f | unreachable!
+      let .const declName us ← CSimp.replaceConstant env f | unreachable!
       let ctorInfo? ← isCtor? declName
       let args ← args.mapIdxM fun idx arg =>
         -- We can rely on `toMono` erasing ctor params eventually; we do not do so here so that type
@@ -793,7 +793,7 @@ where
       visit (f.beta e.getAppArgs)
 
   visitApp (e : Expr) : M (Arg .pure) := do
-    if let .const declName us := CSimp.replaceConstants (← getEnv) e.getAppFn then
+    if let .const declName us ← CSimp.replaceConstant (← getEnv) e.getAppFn then
       checkComputable declName
       if declName == ``Quot.lift then
         visitQuotLift e

tests/lake/tests/shake/.gitignore

@@ -3,6 +3,6 @@ lake-manifest.json
 lake-packages/
 produced.out
 # Working files copied from input/
-lakefile.toml
-Main.lean
-Lib/
+/lakefile.toml
+/Main.lean
+/Lib/

tests/lake/tests/shake/expected/Main.lean

@@ -1,6 +1,7 @@
 module
 
 import Lib.A
+import Lib.CSimp
 
--- Does not use Lib.B, uses Lib.A privately only
 def myValue : Nat := valueA
+def myCSimpedValue : Nat := valueC

tests/lake/tests/shake/input/Lib/C.lean

@@ -0,0 +1,3 @@
+module
+
+@[expose] public def valueC : Nat := 34

tests/lake/tests/shake/input/Lib/CSimp.lean

@@ -0,0 +1,8 @@
+module
+public import Lib.C
+
+public def valueC' : Nat := 34
+
+@[csimp]
+public theorem valueC_eq_valueC' : valueC = valueC' := by
+  rfl

tests/lake/tests/shake/input/Main.lean

@@ -1,7 +1,8 @@
 module
 
-public import Lib.A
-public import Lib.B
+public import Lib.A  -- used privately
+public import Lib.B  -- unused
+public import Lib.CSimp  -- used privately via `[csimp]` replacement
 
--- Does not use Lib.B, uses Lib.A privately only
 def myValue : Nat := valueA
+def myCSimpedValue : Nat := valueC

tests/lake/tests/shake/test.sh

@@ -22,5 +22,5 @@ cp -r input/* .
 test_run build
 test_run shake --fix Main
 
-# Verify Main.lean matches expected (Lib.B import removed)
+# Verify Main.lean matches expected
 check_diff expected/Main.lean Main.lean

tests/lean/csimpCore.lean

@@ -56,6 +56,6 @@ info: (Acc.rec, Acc.recC)
 #guard_msgs in
 run_elab do
   let env ← getEnv
-  for (l, r) in CSimp.ext.getState env |>.map.toList.mergeSort (le := (·.1.lt ·.1)) do
+  for (l, e) in CSimp.ext.getState env |>.map.toList.mergeSort (le := (·.1.lt ·.1)) do
     if !isNoncomputable env l then
-      IO.println (l, r)
+      IO.println (l, e.toDeclName)