From fa9538ffa6cc9c206dd626589b18016a464e1bf4 Mon Sep 17 00:00:00 2001 From: Gabriel Ebner Date: Fri, 21 Oct 2022 10:36:40 -0700 Subject: [PATCH] perf: use old instantiateLevelParams in compiler --- src/Lean/Compiler/LCNF/BaseTypes.lean | 5 +---- src/Lean/Compiler/LCNF/Basic.lean | 6 +++--- src/Lean/Compiler/LCNF/Specialize.lean | 4 ++-- src/Lean/Util/InstantiateLevelParams.lean | 15 +++++++++++--- src/Lean/Util/ReplaceExpr.lean | 24 +++++++++++------------ 5 files changed, 30 insertions(+), 24 deletions(-) diff --git a/src/Lean/Compiler/LCNF/BaseTypes.lean b/src/Lean/Compiler/LCNF/BaseTypes.lean index b6fae06924..662b98853d 100644 --- a/src/Lean/Compiler/LCNF/BaseTypes.lean +++ b/src/Lean/Compiler/LCNF/BaseTypes.lean @@ -29,9 +29,6 @@ def getOtherDeclBaseType (declName : Name) (us : List Level) : CoreM Expr := do let type ← Meta.MetaM.run' <| toLCNFType info.type modifyEnv fun env => baseTypeExt.modifyState env fun s => { s with base := s.base.insert declName type } pure type - if us.isEmpty then - return type - else - return type.instantiateLevelParams info.levelParams us + return type.instantiateLevelParamsNoCache info.levelParams us end Lean.Compiler.LCNF diff --git a/src/Lean/Compiler/LCNF/Basic.lean b/src/Lean/Compiler/LCNF/Basic.lean index c78d585587..ffe9de5320 100644 --- a/src/Lean/Compiler/LCNF/Basic.lean +++ b/src/Lean/Compiler/LCNF/Basic.lean @@ -451,16 +451,16 @@ where | _ => none def Decl.instantiateTypeLevelParams (decl : Decl) (us : List Level) : Expr := - decl.type.instantiateLevelParams decl.levelParams us + decl.type.instantiateLevelParamsNoCache decl.levelParams us def Decl.instantiateParamsLevelParams (decl : Decl) (us : List Level) : Array Param := - decl.params.mapMono fun param => param.updateCore (param.type.instantiateLevelParams decl.levelParams us) + decl.params.mapMono fun param => param.updateCore (param.type.instantiateLevelParamsNoCache decl.levelParams us) partial def Decl.instantiateValueLevelParams (decl : Decl) (us : List Level) : Code := instCode decl.value where instExpr (e : Expr) := - e.instantiateLevelParams decl.levelParams us + e.instantiateLevelParamsNoCache decl.levelParams us instParams (ps : Array Param) := ps.mapMono fun p => p.updateCore (instExpr p.type) diff --git a/src/Lean/Compiler/LCNF/Specialize.lean b/src/Lean/Compiler/LCNF/Specialize.lean index 2a452a7ddc..23ade849a2 100644 --- a/src/Lean/Compiler/LCNF/Specialize.lean +++ b/src/Lean/Compiler/LCNF/Specialize.lean @@ -228,10 +228,10 @@ where modify fun s => s.insert param.fvarId arg else -- Keep the parameter - let param := { param with type := param.type.instantiateLevelParams decl.levelParams us } + let param := { param with type := param.type.instantiateLevelParamsNoCache decl.levelParams us } params := params.push (← internalizeParam param) for param in decl.params[argMask.size:] do - let param := { param with type := param.type.instantiateLevelParams decl.levelParams us } + let param := { param with type := param.type.instantiateLevelParamsNoCache decl.levelParams us } params := params.push (← internalizeParam param) let value := decl.instantiateValueLevelParams us let value ← internalizeCode value diff --git a/src/Lean/Util/InstantiateLevelParams.lean b/src/Lean/Util/InstantiateLevelParams.lean index 41809a96e5..642abd3535 100644 --- a/src/Lean/Util/InstantiateLevelParams.lean +++ b/src/Lean/Util/InstantiateLevelParams.lean @@ -12,7 +12,9 @@ namespace Lean.Expr Instantiate level parameters -/ @[specialize] def instantiateLevelParamsCore (s : Name → Option Level) (e : Expr) : Expr := - e.replace fun e => + e.replace replaceFn +where + @[specialize] replaceFn (e : Expr) : Option Expr := if !e.hasLevelParam then e else match e with | const _ us => e.updateConst! (us.map fun u => u.instantiateParams s) | sort u => e.updateSort! (u.instantiateParams s) @@ -30,6 +32,15 @@ def instantiateLevelParams (e : Expr) (paramNames : List Name) (lvls : List Leve if paramNames.isEmpty || lvls.isEmpty then e else instantiateLevelParamsCore (getParamSubst paramNames lvls) e +/-- +Instantiate univeres level parameters names `paramNames` with `lvls` in `e`. +If the two lists have different length, the smallest one is used. +(Does not preserve expression sharing.) +-/ +def instantiateLevelParamsNoCache (e : Expr) (paramNames : List Name) (lvls : List Level) : Expr := + if paramNames.isEmpty || lvls.isEmpty then e else + e.replaceNoCache (instantiateLevelParamsCore.replaceFn (getParamSubst paramNames lvls)) + private partial def getParamSubstArray (ps : Array Name) (us : Array Level) (p' : Name) (i : Nat) : Option Level := if h : i < ps.size then let p := ps.get ⟨i, h⟩ @@ -47,5 +58,3 @@ def instantiateLevelParamsArray (e : Expr) (paramNames : Array Name) (lvls : Arr if paramNames.isEmpty || lvls.isEmpty then e else e.instantiateLevelParamsCore fun p => getParamSubstArray paramNames lvls p 0 - -end Expr diff --git a/src/Lean/Util/ReplaceExpr.lean b/src/Lean/Util/ReplaceExpr.lean index 2c0d27742f..e1c59c87da 100644 --- a/src/Lean/Util/ReplaceExpr.lean +++ b/src/Lean/Util/ReplaceExpr.lean @@ -78,19 +78,19 @@ end ReplaceImpl /- TODO: use withPtrAddr, withPtrEq to avoid unsafe tricks above. We also need an invariant at `State` and proofs for the `uget` operations. -/ -@[implemented_by ReplaceImpl.replaceUnsafe] -partial def replace (f? : Expr → Option Expr) (e : Expr) : Expr := - /- This is a reference implementation for the unsafe one above -/ +@[specialize] +def replaceNoCache (f? : Expr → Option Expr) (e : Expr) : Expr := match f? e with | some eNew => eNew | none => match e with - | Expr.forallE _ d b _ => let d := replace f? d; let b := replace f? b; e.updateForallE! d b - | Expr.lam _ d b _ => let d := replace f? d; let b := replace f? b; e.updateLambdaE! d b - | Expr.mdata _ b => let b := replace f? b; e.updateMData! b - | Expr.letE _ t v b _ => let t := replace f? t; let v := replace f? v; let b := replace f? b; e.updateLet! t v b - | Expr.app f a => let f := replace f? f; let a := replace f? a; e.updateApp! f a - | Expr.proj _ _ b => let b := replace f? b; e.updateProj! b - | e => e -end Expr + | .forallE _ d b _ => let d := replaceNoCache f? d; let b := replaceNoCache f? b; e.updateForallE! d b + | .lam _ d b _ => let d := replaceNoCache f? d; let b := replaceNoCache f? b; e.updateLambdaE! d b + | .mdata _ b => let b := replaceNoCache f? b; e.updateMData! b + | .letE _ t v b _ => let t := replaceNoCache f? t; let v := replaceNoCache f? v; let b := replaceNoCache f? b; e.updateLet! t v b + | .app f a => let f := replaceNoCache f? f; let a := replaceNoCache f? a; e.updateApp! f a + | .proj _ _ b => let b := replaceNoCache f? b; e.updateProj! b + | e => e -end Lean +@[implemented_by ReplaceImpl.replaceUnsafe] +partial def replace (f? : Expr → Option Expr) (e : Expr) : Expr := + e.replaceNoCache f?