diff --git a/src/Lean/Compiler/CSE.lean b/src/Lean/Compiler/CSE.lean index 785fa5fe30..8c2e429b4f 100644 --- a/src/Lean/Compiler/CSE.lean +++ b/src/Lean/Compiler/CSE.lean @@ -27,13 +27,13 @@ partial def visitCases (casesInfo : CasesInfo) (cases : Expr) : M Expr := do partial def visitLambda (e : Expr) : M Expr := withNewScope do - let (as, e) ← Compiler.visitLambda e - let e ← mkLetUsingScope (← visitLet e) + let (as, e) ← Compiler.visitLambdaCore e + let e ← mkLetUsingScope (← visitLet e as) mkLambda as e -partial def visitLet (e : Expr) : M Expr := do +partial def visitLet (e : Expr) (xs : Array Expr): M Expr := do let saved ← get - try go e #[] finally set saved + try go e xs finally set saved where go (e : Expr) (xs : Array Expr) : M Expr := do match e with diff --git a/src/Lean/Compiler/CompilerM.lean b/src/Lean/Compiler/CompilerM.lean index 7fb8c1d22e..3e317cac41 100644 --- a/src/Lean/Compiler/CompilerM.lean +++ b/src/Lean/Compiler/CompilerM.lean @@ -120,13 +120,16 @@ def ensureUniqueLetVarNames (e : Expr) : CompilerM Expr := do /-- Move through all consecutive lambda abstractions at the top level of `e`. -Returning the body of the last one we find with all bound variables -instantiated as new free variables. +Returning the body of the last one we find with **loose bound variables**. Returns a tuple consisting of: -1. An `Array` of all the newly created free variables -2. The (fully instantiated) body of the last lambda binder that was visited +1. An `Array` of all the newly created free variables. +2. The body of the last lambda binder that was visited. + The caller is responsible for replacing the loose bound variables + with the newly created free variables. + +See `visitLambda`. -/ -def visitLambda (e : Expr) : CompilerM (Array Expr × Expr) := +def visitLambdaCore (e : Expr) : CompilerM (Array Expr × Expr) := go e #[] where go (e : Expr) (fvars : Array Expr) := do @@ -135,7 +138,19 @@ where let fvar ← mkLocalDecl binderName type binderInfo go body (fvars.push fvar) else - return (fvars, e.instantiateRev fvars) + return (fvars, e) + +/-- +Move through all consecutive lambda abstractions at the top level of `e`. +Returning the body of the last one we find with all bound variables +instantiated as new free variables. +Returns a tuple consisting of: +1. An `Array` of all the newly created free variables +2. The (fully instantiated) body of the last lambda binder that was visited +-/ +def visitLambda (e : Expr) : CompilerM (Array Expr × Expr) := do + let (fvars, e) ← visitLambdaCore e + return (fvars, e.instantiateRev fvars) /-- Given an expression representing a `match` return a tuple consisting of: @@ -184,10 +199,10 @@ class VisitLet (m : Type → Type) where export VisitLet (visitLet) -def visitLetImp (e : Expr) (f : Name → Expr → CompilerM Expr) : CompilerM Expr := +@[inline] def visitLetImp (e : Expr) (f : Name → Expr → CompilerM Expr) : CompilerM Expr := go e #[] where - go (e : Expr) (fvars : Array Expr) : CompilerM Expr := do + @[specialize] go (e : Expr) (fvars : Array Expr) : CompilerM Expr := do if let .letE binderName type value body nonDep := e then let type := type.instantiateRev fvars let value := value.instantiateRev fvars @@ -293,8 +308,8 @@ partial def attachJp (e : Expr) (jp : Expr) : CompilerM Expr := do where visitLambda (e : Expr) : ReaderT FVarIdSet CompilerM Expr := do withNewScope do - let (as, e) ← Compiler.visitLambda e - let e ← mkLetUsingScope (← visitLet e #[]) + let (as, e) ← Compiler.visitLambdaCore e + let e ← mkLetUsingScope (← visitLet e as) mkLambda as e visitCases (casesInfo : CasesInfo) (cases : Expr) : ReaderT FVarIdSet CompilerM Expr := do diff --git a/src/Lean/Compiler/PullLocalDecls.lean b/src/Lean/Compiler/PullLocalDecls.lean index 7e85d3334e..a9137e3267 100644 --- a/src/Lean/Compiler/PullLocalDecls.lean +++ b/src/Lean/Compiler/PullLocalDecls.lean @@ -42,9 +42,9 @@ def mkLambda (xs : Array Expr) (letFVarsSavedSize : Nat) (e : Expr) : PullM Expr mutual partial def visitLambda (e : Expr) (topLevel := false) : PullM Expr := do - let (as, e) ← Compiler.visitLambda e + let (as, e) ← Compiler.visitLambdaCore e let letFVarsSavedSize := (← getThe CompilerM.State).letFVars.size - let e ← visitLet e + let e ← visitLet e as if topLevel then Compiler.mkLambda as (← mkLetUsingScope e) else diff --git a/src/Lean/Compiler/Simp.lean b/src/Lean/Compiler/Simp.lean index b8f8f31951..4d602be188 100644 --- a/src/Lean/Compiler/Simp.lean +++ b/src/Lean/Compiler/Simp.lean @@ -84,8 +84,8 @@ partial def internalize (e : Expr) : SimpM Expr := do where visitLambda (e : Expr) : SimpM Expr := do withNewScope do - let (as, e) ← Compiler.visitLambda e - let e ← mkLetUsingScope (← visitLet e #[]) + let (as, e) ← Compiler.visitLambdaCore e + let e ← mkLetUsingScope (← visitLet e as) mkLambda as e visitCases (casesInfo : CasesInfo) (cases : Expr) : SimpM Expr := do @@ -389,14 +389,15 @@ If `checkEmptyTypes := true`, then return `fun a_i : t_i => lcUnreachable` if -/ partial def visitLambda (e : Expr) (checkEmptyTypes := false): SimpM Expr := withNewScope do - let (as, e) ← Compiler.visitLambda e + let (as, e) ← Compiler.visitLambdaCore e if checkEmptyTypes then for a in as do if (← isEmptyType (← inferType a)) then + let e := e.instantiateRev as let unreach ← mkLcUnreachable (← inferType e) let r ← mkLambda as unreach return r - let e ← mkLetUsingScope (← visitLet e) + let e ← mkLetUsingScope (← visitLet e as) mkLambda as e partial def visitCases (casesInfo : CasesInfo) (e : Expr) : SimpM Expr := do diff --git a/src/Lean/Compiler/TerminalCases.lean b/src/Lean/Compiler/TerminalCases.lean index 539067e697..90859b6cfb 100644 --- a/src/Lean/Compiler/TerminalCases.lean +++ b/src/Lean/Compiler/TerminalCases.lean @@ -21,8 +21,8 @@ private partial def visitCases (casesInfo : CasesInfo) (cases : Expr) : Compiler private partial def visitLambda (e : Expr) : CompilerM Expr := withNewScope do - let (as, e) ← Compiler.visitLambda e - let e ← mkLetUsingScope (← visitLet e #[]) + let (as, e) ← Compiler.visitLambdaCore e + let e ← mkLetUsingScope (← visitLet e as) mkLambda as e private partial def visitLet (e : Expr) (fvars : Array Expr) : CompilerM Expr := do