From ae79d14d2718cf66261c589808618c84474aa6d8 Mon Sep 17 00:00:00 2001 From: Sebastian Ullrich Date: Sat, 7 Feb 2026 15:27:00 +0100 Subject: [PATCH] feat: shake: track `[csimp]` uses (#12351) This PR extends the `@[csimp]` attribute to be correctly tracked by `lake shake` --- src/Lean/Compiler/CSimpAttr.lean | 28 +++++++++++++-------- src/Lean/Compiler/LCNF/ToLCNF.lean | 6 ++--- tests/lake/tests/shake/.gitignore | 6 ++--- tests/lake/tests/shake/expected/Main.lean | 3 ++- tests/lake/tests/shake/input/Lib/C.lean | 3 +++ tests/lake/tests/shake/input/Lib/CSimp.lean | 8 ++++++ tests/lake/tests/shake/input/Main.lean | 7 +++--- tests/lake/tests/shake/test.sh | 2 +- tests/lean/csimpCore.lean | 4 +-- 9 files changed, 44 insertions(+), 23 deletions(-) create mode 100644 tests/lake/tests/shake/input/Lib/C.lean create mode 100644 tests/lake/tests/shake/input/Lib/CSimp.lean diff --git a/src/Lean/Compiler/CSimpAttr.lean b/src/Lean/Compiler/CSimpAttr.lean index 25890ec45f..9d60dfaaec 100644 --- a/src/Lean/Compiler/CSimpAttr.lean +++ b/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 diff --git a/src/Lean/Compiler/LCNF/ToLCNF.lean b/src/Lean/Compiler/LCNF/ToLCNF.lean index d86c10da15..8171dcf45b 100644 --- a/src/Lean/Compiler/LCNF/ToLCNF.lean +++ b/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 diff --git a/tests/lake/tests/shake/.gitignore b/tests/lake/tests/shake/.gitignore index 193ac124d3..dc33c5a067 100644 --- a/tests/lake/tests/shake/.gitignore +++ b/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/ diff --git a/tests/lake/tests/shake/expected/Main.lean b/tests/lake/tests/shake/expected/Main.lean index cf0c385818..7435f3dba6 100644 --- a/tests/lake/tests/shake/expected/Main.lean +++ b/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 diff --git a/tests/lake/tests/shake/input/Lib/C.lean b/tests/lake/tests/shake/input/Lib/C.lean new file mode 100644 index 0000000000..776d7b299d --- /dev/null +++ b/tests/lake/tests/shake/input/Lib/C.lean @@ -0,0 +1,3 @@ +module + +@[expose] public def valueC : Nat := 34 diff --git a/tests/lake/tests/shake/input/Lib/CSimp.lean b/tests/lake/tests/shake/input/Lib/CSimp.lean new file mode 100644 index 0000000000..5728abe938 --- /dev/null +++ b/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 diff --git a/tests/lake/tests/shake/input/Main.lean b/tests/lake/tests/shake/input/Main.lean index 57f35e8b7c..f5460d242e 100644 --- a/tests/lake/tests/shake/input/Main.lean +++ b/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 diff --git a/tests/lake/tests/shake/test.sh b/tests/lake/tests/shake/test.sh index 18d6b02ddd..570e782a4b 100755 --- a/tests/lake/tests/shake/test.sh +++ b/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 diff --git a/tests/lean/csimpCore.lean b/tests/lean/csimpCore.lean index ff88b69d64..3cecc06017 100644 --- a/tests/lean/csimpCore.lean +++ b/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)