From f0033cd15e07828bb773db4b2616cd6a308ba4bd Mon Sep 17 00:00:00 2001 From: Cameron Zwarich Date: Thu, 17 Apr 2025 15:34:16 -0700 Subject: [PATCH] fix: consider params to be ground variables in specialization (#8008) This PR changes specialization in the new code generator to consider callee params to be ground variables, which improves the specialization of polymorphic functions. --- src/Lean/Compiler/LCNF/Specialize.lean | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/Lean/Compiler/LCNF/Specialize.lean b/src/Lean/Compiler/LCNF/Specialize.lean index 00404b92c5..30ea5a4127 100644 --- a/src/Lean/Compiler/LCNF/Specialize.lean +++ b/src/Lean/Compiler/LCNF/Specialize.lean @@ -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