diff --git a/src/Lean/Compiler/LCNF/Simp/Basic.lean b/src/Lean/Compiler/LCNF/Simp/Basic.lean index e29c71657f..55a6ef200f 100644 --- a/src/Lean/Compiler/LCNF/Simp/Basic.lean +++ b/src/Lean/Compiler/LCNF/Simp/Basic.lean @@ -11,6 +11,23 @@ import Lean.Compiler.LCNF.CompilerM namespace Lean.Compiler.LCNF namespace Simp +/-- +Similar to `findFunDecl?`, but follows aliases (i.e., `let _x.i := _x.j`). +Consider the following example +``` +fun _f.1 ... := ... +let _x.2 := _f.1 +``` +`findFunDecl? _x.2` returns `none`, but `findFunDecl'? _x.2` returns the declaration for `_f.1`. +-/ +partial def findFunDecl'? (fvarId : FVarId) : CompilerM (Option FunDecl) := do + if let some decl ← findFunDecl? fvarId then + return decl + else if let some (.fvar fvarId' #[]) ← findLetExpr? fvarId then + findFunDecl'? fvarId' + else + return none + /- -- TODO: cleanup partial def findExpr (e : LetExpr) : CompilerM LetExpr := do @@ -22,18 +39,6 @@ partial def findExpr (e : LetExpr) : CompilerM LetExpr := do else return e | _ => return e - -partial def findFunDecl? (e : LetExpr) : CompilerM (Option FunDecl) := do - match e with - | .fvar fvarId args => - - if let some decl ← LCNF.findFunDecl? fvarId then - return some decl - else if let some decl ← findLetDecl? fvarId then - findFunDecl? decl.value - else - return none - | _ => return none -/ end Simp end Lean.Compiler.LCNF diff --git a/src/Lean/Compiler/LCNF/Simp/FunDeclInfo.lean b/src/Lean/Compiler/LCNF/Simp/FunDeclInfo.lean index f363105577..ac4fec8e8e 100644 --- a/src/Lean/Compiler/LCNF/Simp/FunDeclInfo.lean +++ b/src/Lean/Compiler/LCNF/Simp/FunDeclInfo.lean @@ -96,7 +96,7 @@ where addArgOcc (arg : Arg) : StateRefT FunDeclInfoMap CompilerM Unit := do match arg with | .fvar fvarId => - let some funDecl ← findFunDecl? fvarId | return () + let some funDecl ← findFunDecl'? fvarId | return () modify fun s => s.addHo funDecl.fvarId | .erased .. | .type .. => return () @@ -105,7 +105,7 @@ where | .erased | .value .. | .proj .. => return () | .const _ _ args => args.forM addArgOcc | .fvar fvarId args => - let some funDecl ← findFunDecl? fvarId | return () + let some funDecl ← findFunDecl'? fvarId | return () modify fun s => s.add funDecl.fvarId args.forM addArgOcc diff --git a/src/Lean/Compiler/LCNF/Simp/InlineCandidate.lean b/src/Lean/Compiler/LCNF/Simp/InlineCandidate.lean index 69441ab6cb..2b98328fd1 100644 --- a/src/Lean/Compiler/LCNF/Simp/InlineCandidate.lean +++ b/src/Lean/Compiler/LCNF/Simp/InlineCandidate.lean @@ -88,7 +88,7 @@ def inlineCandidate? (e : LetExpr) : SimpM (Option InlineCandidateInfo) := do params, value } else if let .fvar f args := e then - let some decl ← findFunDecl? f | return none + 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 diff --git a/src/Lean/Compiler/LCNF/Simp/Main.lean b/src/Lean/Compiler/LCNF/Simp/Main.lean index 5d255ccfa5..dd99685079 100644 --- a/src/Lean/Compiler/LCNF/Simp/Main.lean +++ b/src/Lean/Compiler/LCNF/Simp/Main.lean @@ -60,7 +60,8 @@ def specializePartialApp (info : InlineCandidateInfo) : SimpM FunDecl := do Try to inline a join point. -/ partial def inlineJp? (fvarId : FVarId) (args : Array Arg) : SimpM (Option Code) := do - let some decl ← LCNF.findFunDecl? fvarId | return none + /- Remark: we don't need to use `findFunDecl'?` here. -/ + let some decl ← findFunDecl? fvarId | return none unless (← shouldInlineLocal decl) do return none markSimplified betaReduce decl.params decl.value args