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

fix: consider params to be ground variables in specialization (#8008)

Browse source 
This PR changes specialization in the new code generator to consider
callee params to be ground variables, which improves the specialization
of polymorphic functions.
This commit is contained in:
Cameron Zwarich 2025-04-17 15:34:16 -07:00 • committed by GitHub
parent 7bbcfdf712
commit f0033cd15e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 7 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -261,6 +261,9 @@ 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) {}
mutual
/--
Try to specialize the function application in the given let-declaration.
@ -295,7 +298,8 @@ mutual
specDecl.saveBase
let specDecl ← specDecl.simp {}
let specDecl ← specDecl.simp { etaPoly := true, inlinePartial := true, implementedBy := true }
let value ← withReader (fun _ => { declName := specDecl.name }) do
let ground := paramsToVarSet specDecl.params
let value ← withReader (fun _ => { declName := specDecl.name, ground }) do
withParams specDecl.params <| specDecl.value.mapCodeM visitCode
let specDecl := { specDecl with value }
modify fun s => { s with decls := s.decls.push specDecl }
@ -337,7 +341,8 @@ def main (decl : Decl) : SpecializeM Decl := do
end Specialize
partial def Decl.specialize (decl : Decl) : CompilerM (Array Decl) := do
let (decl, s) ← Specialize.main decl |>.run { declName := decl.name } |>.run {}
let ground := Specialize.paramsToVarSet decl.params
let (decl, s) ← Specialize.main decl |>.run { declName := decl.name, ground } |>.run {}
return s.decls.push decl
def specialize : Pass where