max/lean4-htt
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

fix: only treat type/instance params as ground vars in specialization (#8565)

Browse source 
This PR makes the LCNF specialization pass only treat type/instance
params as ground vars. The current policy was too liberal and would
result on computations being floated into specialized loops.
This commit is contained in:
Cameron Zwarich 2025-05-31 14:18:24 -07:00 • committed by GitHub
parent 2c8ee4f29c
commit 0c43efc2c9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 8 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
src/Lean/Compiler/LCNF/Specialize.lean
View file

@ -268,8 +268,12 @@ def getRemainingArgs (paramsInfo : Array SpecParamInfo) (args : Array Arg) : Arr
result := result.push arg
return result ++ args[paramsInfo.size:]
def paramsToVarSet (params : Array Param) : FVarIdSet :=
params.foldl (fun r p => r.insert p.fvarId) {}
def paramsToGroundVars (params : Array Param) : CompilerM FVarIdSet :=
params.foldlM (init := {}) fun r p => do
if isTypeFormerType p.type || (← isArrowClass? p.type).isSome then
return r.insert p.fvarId
else
return r
mutual
/--
@ -305,7 +309,7 @@ mutual
specDecl.saveBase
let specDecl ← specDecl.simp {}
let specDecl ← specDecl.simp { etaPoly := true, inlinePartial := true, implementedBy := true }
let ground := paramsToVarSet specDecl.params
let ground ← paramsToGroundVars specDecl.params
let value ← withReader (fun _ => { declName := specDecl.name, ground }) do
withParams specDecl.params <| specDecl.value.mapCodeM visitCode
let specDecl := { specDecl with value }
@ -352,7 +356,7 @@ def main (decl : Decl) : SpecializeM Decl := do
end Specialize
partial def Decl.specialize (decl : Decl) : CompilerM (Array Decl) := do
let ground := Specialize.paramsToVarSet decl.params
let ground ← Specialize.paramsToGroundVars decl.params
let (decl, s) ← Specialize.main decl |>.run { declName := decl.name, ground } |>.run {}
return s.decls.push decl