chore: port InlineCandidate.lean

This commit is contained in:
Leonardo de Moura 2022-11-01 17:11:56 -07:00
parent 0c58913cf1
commit 123aed11ca
2 changed files with 27 additions and 21 deletions

View file

@ -92,6 +92,14 @@ def findLetDecl? (fvarId : FVarId) : CompilerM (Option LetDecl) :=
def findFunDecl? (fvarId : FVarId) : CompilerM (Option FunDecl) :=
return (← get).lctx.funDecls.find? fvarId
def isConstructorApp (fvarId : FVarId) : CompilerM Bool := do
let some { value := .const declName _ _, .. } ← findLetDecl? fvarId | return false
return (← getConstInfo declName) matches .ctorInfo ..
def Arg.isConstructorApp (arg : Arg) : CompilerM Bool := do
let .fvar fvarId := arg | return false
LCNF.isConstructorApp fvarId
def getParam (fvarId : FVarId) : CompilerM Param := do
let some param ← findParam? fvarId | throwError "unknown parameter {fvarId.name}"
return param

View file

@ -5,9 +5,6 @@ Authors: Leonardo de Moura
-/
import Lean.Compiler.LCNF.Simp.SimpM
set_option warningAsError false
#exit
namespace Lean.Compiler.LCNF
namespace Simp
@ -20,8 +17,8 @@ structure InlineCandidateInfo where
params : Array Param
/-- Value (lambda expression) of the function to be inlined. -/
value : Code
f : Expr
args : Array Expr
fType : Expr
args : Array Arg
/-- `ifReduce = true` if the declaration being inlined was tagged with `inline_if_reduce`. -/
ifReduce : Bool
/-- `recursive = true` if the declaration being inline is in a mutually recursive block. -/
@ -34,15 +31,14 @@ def InlineCandidateInfo.arity : InlineCandidateInfo → Nat
/--
Return `some info` if `e` should be inlined.
-/
def inlineCandidate? (e : Expr) : SimpM (Option InlineCandidateInfo) := do
def inlineCandidate? (e : LetExpr) : SimpM (Option InlineCandidateInfo) := do
let mut e := e
let mut mustInline := false
if e.isAppOfArity ``inline 2 then
e ← findExpr e.appArg!
if let .const ``inline _ #[_, .fvar argFVarId] := e then
let some decl ← findLetDecl? argFVarId | return none
e := decl.value
mustInline := true
let numArgs := e.getAppNumArgs
let f := e.getAppFn
if let .const declName us ← findExpr f then
if let .const declName us args := e then
unless (← read).config.inlineDefs do
return none
let some decl ← getDecl? declName | return none
@ -73,33 +69,35 @@ def inlineCandidate? (e : Expr) : SimpM (Option InlineCandidateInfo) := do
/- check arity -/
let arity := decl.getArity
let inlinePartial := (← read).config.inlinePartial
if !mustInline && !inlinePartial && numArgs < arity then return none
if !mustInline && !inlinePartial && args.size < arity then return none
if decl.inlineIfReduceAttr then
let some paramIdx := decl.isCasesOnParam? | return none
unless paramIdx < numArgs do return none
let arg ← findExpr (e.getArg! paramIdx)
unless arg.isConstructorApp (← getEnv) do return none
unless paramIdx < args.size do return none
let arg := args[paramIdx]!
unless (← arg.isConstructorApp) do return none
let params := decl.instantiateParamsLevelParams us
let value := decl.instantiateValueLevelParams us
let type := decl.instantiateTypeLevelParams us
incInline
return some {
isLocal := false
f := e.getAppFn
args := e.getAppArgs
fType := type
args := args
ifReduce := decl.inlineIfReduceAttr
recursive := decl.recursive
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
else if let .fvar f args := e then
let some decl ← findFunDecl? f | return none
unless args.size > 0 do return none -- It is not worth to inline a local function that does not take any arguments
unless mustInline || (← shouldInlineLocal decl) do return none
-- Remark: we inline local function declarations even if they are partial applied
incInlineLocal
modify fun s => { s with inlineLocal := s.inlineLocal + 1 }
return some {
isLocal := true
f := e.getAppFn
args := e.getAppArgs
fType := (← getType f)
args := args
params := decl.params
value := decl.value
ifReduce := false