diff --git a/src/Lean/Compiler/LCNF/Passes.lean b/src/Lean/Compiler/LCNF/Passes.lean index e2e94111a7..ec685a4aca 100644 --- a/src/Lean/Compiler/LCNF/Passes.lean +++ b/src/Lean/Compiler/LCNF/Passes.lean @@ -13,6 +13,6 @@ import Lean.Compiler.LCNF.ReduceJpArity namespace Lean.Compiler.LCNF @[cpass] def builtin : PassInstaller := - .append #[pullInstances, cse, simp, pullFunDecls, reduceJpArity, simp { etaPoly := true }] + .append #[pullInstances, cse, simp, pullFunDecls, reduceJpArity, simp { etaPoly := true, inlinePartial := true }] end Lean.Compiler.LCNF diff --git a/src/Lean/Compiler/LCNF/Simp.lean b/src/Lean/Compiler/LCNF/Simp.lean index 69a8a8ce1b..59aa0b93f7 100644 --- a/src/Lean/Compiler/LCNF/Simp.lean +++ b/src/Lean/Compiler/LCNF/Simp.lean @@ -6,6 +6,7 @@ Authors: Leonardo de Moura import Lean.Util.Recognizers import Lean.Meta.Instances import Lean.Compiler.InlineAttrs +import Lean.Compiler.Specialize import Lean.Compiler.LCNF.CompilerM import Lean.Compiler.LCNF.ElimDead import Lean.Compiler.LCNF.Bind @@ -15,6 +16,28 @@ import Lean.Compiler.LCNF.PassManager import Lean.Compiler.LCNF.AlphaEqv namespace Lean.Compiler.LCNF + +/-- +Return `true` if the arrow type contains an instance implicit argument. +-/ +def hasLocalInst (type : Expr) : Bool := + match type with + | .forallE _ _ b bi => bi.isInstImplicit || hasLocalInst b + | _ => false + +/-- +Return `true` if `decl` is supposed to be inlined/specialized. +-/ +def isTemplateLike (decl : Decl) : CoreM Bool := do + if hasLocalInst decl.type then + return true -- `decl` applications will be specialized + else if (← Meta.isInstance decl.name) then + return true -- `decl` is "fuel" for code specialization + else if hasInlineAttribute (← getEnv) decl.name || hasSpecializeAttribute (← getEnv) decl.name then + return true -- `decl` is going to be inlined or specialized + else + return false + namespace Simp /-- @@ -115,6 +138,12 @@ structure Config where for this kind of application, and we want all of them to specialized or inlined. -/ etaPoly : Bool := false + /-- + If `inlinePartial` is `true`, we inline partial function applications tagged + with `[inline]`. Note that this option is automatically disabled when processing + declarations tagged with `[inline]`, marked to be specialized, or instances. + -/ + inlinePartial := false deriving Inhabited structure Context where @@ -283,15 +312,16 @@ def InlineCandidateInfo.arity : InlineCandidateInfo → Nat Return `some info` if `e` should be inlined. -/ def inlineCandidate? (e : Expr) : SimpM (Option InlineCandidateInfo) := do + let numArgs := e.getAppNumArgs let f := e.getAppFn if let .const declName us ← findExpr f then unless hasInlineAttribute (← getEnv) declName do return none -- TODO: check whether function is recursive or not. -- We can skip the test and store function inline so far. let some decl ← getStage1Decl? declName | return none - let numArgs := e.getAppNumArgs let arity := decl.getArity - if numArgs < arity then return none + let inlinePartial := (← read).config.inlinePartial + if !inlinePartial && numArgs < arity then return none let params := decl.instantiateParamsLevelParams us let value := decl.instantiateValueLevelParams us incInline @@ -300,6 +330,7 @@ def inlineCandidate? (e : Expr) : SimpM (Option InlineCandidateInfo) := do params, value } else if let some decl ← findFunDecl? f then + unless numArgs > 0 do return none -- It is not worth to inline a local function that does not take any arguments unless (← shouldInlineLocal decl) do return none -- Remark: we inline local function declarations even if they are partial applied incInlineLocal @@ -750,14 +781,6 @@ def eraseLocalDecl (fvarId : FVarId) : SimpM Unit := do eraseFVar fvarId markSimplified -/-- -Return `true` if the arrow type contains an instance implicit argument. --/ -def hasLocalInst (type : Expr) : Bool := - match type with - | .forallE _ _ b bi => bi.isInstImplicit || hasLocalInst b - | _ => false - /-- When the configuration flag `etaPoly = true`, we eta-expand partial applications of functions that take local instances as arguments. @@ -935,14 +958,14 @@ def Decl.simp? (decl : Decl) : SimpM (Option Decl) := do partial def Decl.simp (decl : Decl) (config : Config) : CompilerM Decl := do let mut config := config - if (← pure (Simp.hasLocalInst decl.type) <||> Meta.isInstance decl.name) then + if (← isTemplateLike decl) then /- - We do not eta-expand partial applications in instances or when the declaration type - has local instances. Recall we do not generate code for them. + We do not eta-expand or inline partial applications in template like code. + Recall we don't want to generate code for them. Remark: by eta-expanding partial applications in instaces, we also make the simplifier work harder when inlining instance projections. -/ - config := { config with etaPoly := false } + config := { config with etaPoly := false, inlinePartial := false } go decl config where go (decl : Decl) (config : Config) : CompilerM Decl := do