diff --git a/src/Lean/Compiler/CompilerM.lean b/src/Lean/Compiler/CompilerM.lean index 411477189f..af5f4e6098 100644 --- a/src/Lean/Compiler/CompilerM.lean +++ b/src/Lean/Compiler/CompilerM.lean @@ -221,6 +221,17 @@ Shorthand for `LocalContext.mkLambda` with the `LocalContext` of `CompilerM`. def mkLambda (xs : Array Expr) (e : Expr) : CompilerM Expr := return (← get).lctx.mkLambda xs e +/-- +Given a join point `jp` of the form `fun y => body`, if `jp` is simple (see `isSimpleLCNF`), just return it +Otherwise, create `let jp := fun y => body` declaration and return `jp`. +-/ +def mkJpDeclIfNotSimple (jp : Expr) : CompilerM Expr := do + if (← isSimpleLCNF jp.bindingBody!) then + -- Join point is too simple, we eagerly inline it. + return jp + else + mkJpDecl jp + /-- Create "jump" to join point `jp` with value `e`. Remarks: @@ -244,4 +255,8 @@ def mkJump (jp : Expr) (e : Expr) : CompilerM Expr := do let x ← mkAuxLetDecl (← mkLcCast x d) return mkJpApp x +def mkOptJump (jp? : Option Expr) (e : Expr) : CompilerM Expr := do + let some jp := jp? | return e + mkJump jp e + end Lean.Compiler diff --git a/src/Lean/Compiler/Simp.lean b/src/Lean/Compiler/Simp.lean index ec9cd9586f..0ee815e25b 100644 --- a/src/Lean/Compiler/Simp.lean +++ b/src/Lean/Compiler/Simp.lean @@ -17,12 +17,19 @@ structure Config where structure Context where config : Config := {} + /-- + Current continuation. It is a join point or lambda abstraction. + -/ jp? : Option Expr := none -structure State where - unit : Unit := () +abbrev SimpM := ReaderT Context CompilerM -abbrev SimpM := ReaderT Context $ StateRefT State CompilerM +def withJp (jp : Expr) (x : SimpM α) : SimpM α := do + let jp ← mkJpDeclIfNotSimple jp + withReader (fun ctx => { ctx with jp? := some jp }) x + +def withoutJp (x : SimpM α) : SimpM α := + withReader (fun ctx => { ctx with jp? := none }) x def inlineCandidate? (e : Expr) : SimpM (Option Nat) := do let .const declName _ := e.getAppFn | return none @@ -86,63 +93,80 @@ partial def inlineApp (e : Expr) : SimpM Expr := do assert! !value.isLambda visitLet value -partial def visitLet (e : Expr) : SimpM Expr := do - go e #[] -where - go (e : Expr) (xs : Array Expr) : SimpM Expr := do - let rec inlineApp? (e : Expr) (k? : Option Expr) : SimpM (Option Expr) := do - let some numExtraArgs ← inlineCandidate? e | return none - let args := e.getAppArgs - if k?.isNone && numExtraArgs == 0 then - inlineApp e - else - let toInline := mkAppN e.getAppFn args[:args.size - numExtraArgs] - let jpDomain ← inferType toInline - let binderName ← mkFreshUserName `_y - let bodyAbst ← withNewScope do - let y ← mkLocalDecl binderName jpDomain - let body ← if numExtraArgs == 0 then - go k?.get! (xs.push y) - else if let some k := k? then - let x ← mkAuxLetDecl (mkAppN y args[args.size - numExtraArgs:]) - go k (xs.push x) - else - pure <| mkAppN y args[args.size - numExtraArgs:] - let body ← mkLetUsingScope body - return body.abstract #[y] - let jp ← if (← isSimpleLCNF bodyAbst) then - -- Join point is too simple, we eagerly inline it. - pure <| .lam binderName jpDomain bodyAbst .default - else - mkJpDecl (.lam binderName jpDomain bodyAbst .default) - withReader (fun _ => { jp? := some jp }) do - inlineApp toInline - match e with - | .letE binderName type value body nonDep => - let mut value := value.instantiateRev xs - if value.isLambda then - value ← withReader (fun _ => {}) <| visitLambda value - if value.isFVar then - /- Eliminate `let _x_i := _x_j;` -/ - go body (xs.push value) - else if let some e ← inlineApp? value body then - return e - else - let type := type.instantiateRev xs - let x ← mkLetDecl binderName type value nonDep - go body (xs.push x) - | _ => - let e := e.instantiateRev xs - if let some casesInfo ← isCasesApp? e then - visitCases casesInfo e - else if let some e ← inlineApp? e none then - return e - else - let e ← expandTrivialExpr e - match (← read).jp? with - | none => return e - | some jp => mkJump jp e +/-- +If `e` is an application that can be inlined, inline it. +`k?` is the optional "continuation" for `e`, and it may contain loose bound variables +that need to instantiated with `xs`. That is, if `k? = some k`, then `k.instantiateRev xs` +is an expression without loose bound variables. +-/ +partial def inlineApp? (e : Expr) (xs : Array Expr) (k? : Option Expr) : SimpM (Option Expr) := do + let some numExtraArgs ← inlineCandidate? e | return none + let args := e.getAppArgs + if k?.isNone && numExtraArgs == 0 then + -- Easy case, there is not continuation and `e` is not over applied + inlineApp e + else + /- + There is a continuation `k` or `e` is over applied. + If `e` is over applied, the extra arguments act as continuation. + -/ + let toInline := mkAppN e.getAppFn args[:args.size - numExtraArgs] + /- + `toInline` is the application that is going to be inline + We create a new join point + ``` + let jp := fun y => + let x := y -- if `e` is over applied + k + ``` + Recall that `visitLet` incorporates the current continuation + to the new join point `jp`. + -/ + let jpDomain ← inferType toInline + let binderName ← mkFreshUserName `_y + let jp ← withNewScope do + let y ← mkLocalDecl binderName jpDomain + let body ← if numExtraArgs == 0 then + visitLet k?.get! (xs.push y) + else + let x ← mkAuxLetDecl (mkAppN y args[args.size - numExtraArgs:]) + if let some k := k? then + visitLet k (xs.push x) + else + visitLet x (xs.push x) + let body ← mkLetUsingScope body + mkLambda #[y] body + /- Inline `toInline` and "go-to" `jp` with the result. -/ + withJp jp do inlineApp toInline + +/-- +Let-declaration basic block visitor. `e` may contain loose bound variables that +still have to be instantiated with `xs`. +-/ +partial def visitLet (e : Expr) (xs : Array Expr := #[]): SimpM Expr := do + match e with + | .letE binderName type value body nonDep => + let mut value := value.instantiateRev xs + if value.isLambda then + value ← withoutJp <| visitLambda value + if value.isFVar then + /- Eliminate `let _x_i := _x_j;` -/ + visitLet body (xs.push value) + else if let some e ← inlineApp? value xs body then + return e + else + let type := type.instantiateRev xs + let x ← mkLetDecl binderName type value nonDep + visitLet body (xs.push x) + | _ => + let e := e.instantiateRev xs + if let some casesInfo ← isCasesApp? e then + visitCases casesInfo e + else if let some e ← inlineApp? e #[] none then + return e + else + mkOptJump (← read).jp? (← expandTrivialExpr e) end end Simp diff --git a/src/Lean/Compiler/TerminalCases.lean b/src/Lean/Compiler/TerminalCases.lean index bbe2d1e66b..dc1a426bed 100644 --- a/src/Lean/Compiler/TerminalCases.lean +++ b/src/Lean/Compiler/TerminalCases.lean @@ -17,6 +17,13 @@ structure Context where abbrev M := ReaderT Context CompilerM +def withJp (jp : Expr) (x : M α) : M α := do + let jp ← mkJpDeclIfNotSimple jp + withReader (fun ctx => { ctx with jp? := some jp }) x + +def withoutJp (x : M α) : M α := + withReader (fun ctx => { ctx with jp? := none }) x + mutual private partial def visitCases (casesInfo : CasesInfo) (cases : Expr) : M Expr := do @@ -40,30 +47,23 @@ private partial def visitLet (e : Expr) (fvars : Array Expr) : M Expr := do let type := type.instantiateRev fvars let mut value := value.instantiateRev fvars if let some casesInfo ← isCasesApp? value then - let bodyAbst ← withNewScope do + let jp ← withNewScope do let x ← mkLocalDecl binderName type let body ← visitLet body (fvars.push x) let body ← mkLetUsingScope body - return body.abstract #[x] - let jp ← if (← isSimpleLCNF bodyAbst) then - -- Join point is too simple, we eagerly inline it. - pure <| .lam binderName type bodyAbst .default - else - mkJpDecl (.lam binderName type bodyAbst .default) - withReader (fun _ => { jp? := some jp }) do - visitCases casesInfo value + mkLambda #[x] body + withJp jp do visitCases casesInfo value else if value.isLambda then - value ← withReader (fun _ => {}) <| visitLambda value + value ← withoutJp <| visitLambda value let fvar ← mkLetDecl binderName type value nonDep visitLet body (fvars.push fvar) | e => let e := e.instantiateRev fvars if let some casesInfo ← isCasesApp? e then visitCases casesInfo e - else match (← read).jp? with - | none => return e - | some jp => mkJump jp e + else + mkOptJump (← read).jp? e end