feat: cache specialization results

src/Lean/Compiler/LCNF/Level.lean

@@ -28,9 +28,11 @@ when comparing keys.
 /-- State for the universe level normalizer monad. -/
 structure State where
   /-- Counter for generating new (normalized) universe parameter names. -/
-  nextIdx : Nat := 1
+  nextIdx    : Nat := 1
   /-- Mapping from existing universe parameter names to the new ones. -/
-  map     : HashMap Name Level := {}
+  map        : HashMap Name Level := {}
+  /-- Parameters that have been normalized. -/
+  paramNames : Array Name := #[]
 
 /-- Monad for the universe leve normalizer -/
 abbrev M := StateM State
@@ -51,7 +53,8 @@ partial def normLevel (u : Level) : M Level := do
       | some u => return u
       | none   =>
         let u := Level.param <| (`u).appendIndexAfter (← get).nextIdx
-        modify fun s => { s with nextIdx := s.nextIdx + 1, map := s.map.insert n u }
+        modify fun { nextIdx, map, paramNames } =>
+          { nextIdx := nextIdx + 1, map := map.insert n u, paramNames := paramNames.push n }
         return u
 
 /--
@@ -76,9 +79,11 @@ end NormLevelParam
 
 /--
 Normalize universe level parameter names in the given expression.
+The function also returns the list of universe level parameter names that have been normalized.
 -/
-def normLevelParams (e : Expr) : Expr :=
-  NormLevelParam.normExpr e |>.run' {}
+def normLevelParams (e : Expr) : Expr × List Name :=
+  let (e, s) := NormLevelParam.normExpr e |>.run {}
+  (e, s.paramNames.toList)
 
 namespace CollectLevelParams
 /-!

src/Lean/Compiler/LCNF/Passes.lean

@@ -23,7 +23,8 @@ namespace Lean.Compiler.LCNF
     findJoinPoints,
     reduceJpArity,
     simp { etaPoly := true, inlinePartial := true, implementedBy := true } (occurrence := 1),
-    specialize
+    specialize,
+    simp (occurrence := 2)
   ]
 
 end Lean.Compiler.LCNF

src/Lean/Compiler/LCNF/Specialize.lean

@@ -14,6 +14,17 @@ import Lean.Compiler.LCNF.Level
 namespace Lean.Compiler.LCNF
 namespace Specialize
 
+abbrev Cache := Std.PHashMap Expr Name
+
+builtin_initialize specCacheExt : EnvExtension Cache ←
+  registerEnvExtension (pure {})
+
+def cacheSpec (key : Expr) (declName : Name) : CoreM Unit :=
+  modifyEnv fun env => specCacheExt.modifyState env (·.insert key declName)
+
+def findSpecCache? (key : Expr) : CoreM (Option Name) :=
+  return specCacheExt.getState (← getEnv) |>.find? key
+
 structure Context where
   /--
   Set of free variables in scope. The "collector" uses this information when collecting
@@ -209,19 +220,20 @@ end
 /--
 Given the specialization mask `paramsInfo` and the arguments `args`,
 collect their dependencies, and return an array `mask` of size `paramsInfo.size` s.t.
-- `mask[i] = args[i]` if `paramsInfo[i] != .other`
-- `mask[i] = lcErased`, otherwise.
+- `mask[i] = some args[i]` if `paramsInfo[i] != .other`
+- `mask[i] = none`, otherwise.
 That is, `mask` contains only the arguments that are contributing to the code specialization.
-We use this information to compute a "key" to uniquely identify the code specialization.
+We use this information to compute a "key" to uniquely identify the code specialization, and
+creating the specialized code.
 -/
-def collect (paramsInfo : Array SpecParamInfo) (args : Array Expr) : CollectorM (Array Expr) := do
+def collect (paramsInfo : Array SpecParamInfo) (args : Array Expr) : CollectorM (Array (Option Expr)) := do
   let mut argMask := #[]
   for paramInfo in paramsInfo, arg in args do
     match paramInfo with
     | .other =>
-      argMask := argMask.push (mkConst ``lcErased)
+      argMask := argMask.push none
     | .fixedNeutral | .user | .fixedInst | .fixedHO =>
-      argMask := argMask.push arg
+      argMask := argMask.push (some arg)
       collectExpr arg
   return argMask
 
@@ -261,8 +273,10 @@ termination_by go => values.size - i
 /--
 Create the "key" that uniquely identifies a code specialization.
 `params` and `decls` are the declarations collected by the `collect` function above.
+The result contains the list of universe level parameter names the key that `params`, `decls`, and `body` depends on.
+We use this information to create the new auxiliary declaration and resulting application.
 -/
-def mkKey (params : Array Param) (decls : Array CodeDecl) (body : Expr) : CompilerM Expr := do
+def mkKey (params : Array Param) (decls : Array CodeDecl) (body : Expr) : CompilerM (Expr × List Name) := do
   let body ← expandCodeDecls decls body
   let key := ToExpr.run do
     ToExpr.withParams params do
@@ -270,58 +284,83 @@ def mkKey (params : Array Param) (decls : Array CodeDecl) (body : Expr) : Compil
   return normLevelParams key
 
 open Internalize in
-def mkSpecDecl (decl : Decl) (us : List Level) (paramsInfo : Array SpecParamInfo) (params : Array Param) (decls : Array CodeDecl) (args : Array Expr) : SpecializeM Decl := do
+/--
+Specialize `decl` using
+- `us`: the universe level used to instantiate `decl.name`
+- `argMask`: arguments that are being used to specialize the declaration.
+- `params`: new parameters that arguments in `argMask` depend on.
+- `decls`: local declarations that arguments in `argMask` depend on.
+- `levelParamsNew`: the universe level parameters for the new declaration.
+-/
+def mkSpecDecl (decl : Decl) (us : List Level) (argMask : Array (Option Expr)) (params : Array Param) (decls : Array CodeDecl) (levelParamsNew : List Name) : SpecializeM Decl := do
   let nameNew := decl.name ++ `_at_ ++ (← read).declName ++ (`spec).appendIndexAfter (← get).decls.size
   go nameNew |>.run' {}
 where
   go (nameNew : Name) : InternalizeM Decl := do
     let mut params ← params.mapM internalizeParam
     let decls ← decls.mapM internalizeCodeDecl
-    for param in decl.params, info in paramsInfo, arg in args do
-      if info matches .other then
+    for param in decl.params, arg in argMask do
+      if let some arg := arg then
+        let arg ← normExpr arg
+        modify fun s => s.insert param.fvarId arg
+      else
         -- Keep the parameter
         let param := { param with type := param.type.instantiateLevelParams decl.levelParams us }
         params := params.push (← internalizeParam param)
-      else
-        let arg ← normExpr arg
-        modify fun s => s.insert param.fvarId arg
     let value := decl.instantiateValueLevelParams us
     let value ← internalizeCode value
     let value := attachCodeDecls decls value
     let type ← value.inferType
     let type ← mkForallParams params type
-    let decl := { name := nameNew, levelParams := [], params, type, value : Decl }
+    let decl := { name := nameNew, levelParams := levelParamsNew, params, type, value : Decl }
     return decl.setLevelParams
 
+/--
+Given the specialization mask `paramsInfo` and the arguments `args`,
+return the arguments that have not been considered for specialization.
+-/
+def getRemainingArgs (paramsInfo : Array SpecParamInfo) (args : Array Expr) : Array Expr := Id.run do
+  let mut result := #[]
+  for info in paramsInfo, arg in args do
+    if info matches .other then
+      result := result.push arg
+  return result ++ args[paramsInfo.size:]
+
 mutual
   /--
   Try to specialize the function application in the given let-declaration.
   `k` is the continuation for the let-declaration.
   -/
-  partial def specializeApp? (letDecl : LetDecl) (_k : Code) : SpecializeM (Option Code) := do
-    unless letDecl.value.isApp do return none
-    let f := letDecl.value.getAppFn
+  partial def specializeApp? (e : Expr) : SpecializeM (Option Expr) := do
+    unless e.isApp do return none
+    let f := e.getAppFn
     let .const declName us := f | return none
     if (← Meta.isInstance declName) then return none
     let some paramsInfo ← getSpecParamInfo? declName | return none
-    let args := letDecl.value.getAppArgs
+    let args := e.getAppArgs
     unless (← shouldSpecialize paramsInfo args) do return none
     let some decl ← getStage1Decl? declName | return none
-    trace[Compiler.specialize.candidate] "{letDecl.value}, {paramsInfo}"
+    trace[Compiler.specialize.candidate] "{e}, {paramsInfo}"
     let (argMask, { params, decls, .. }) ← Collector.collect paramsInfo args |>.run {}
-    let keyBody := mkAppN f argMask
-    let key ← mkKey params decls keyBody
+    let keyBody := mkAppN f (argMask.filterMap id)
+    let (key, levelParamsNew) ← mkKey params decls keyBody
     trace[Compiler.specialize.candidate] "key: {key}"
     assert! !key.hasLooseBVars
     assert! !key.hasFVar
-    trace[Compiler.specialize.candidate] "decls:  {decls.map fun | .let decl | .jp decl | .fun decl => decl.binderName}"
-    trace[Compiler.specialize.candidate] "params: {params.map fun p => p.binderName}"
-    -- TODO: create cache and check cached results, apply `visitCode` to `specDecl`, return new application
-    let specDecl ← mkSpecDecl decl us paramsInfo params decls args
-    let specDecl ← specDecl.etaExpand
-    let specDecl ← specDecl.simp {} -- TODO: `simp` config
-    modify fun s => { s with decls := s.decls.push specDecl }
-    return none
+    let usNew := levelParamsNew.map mkLevelParam
+    let argsNew := params.map (mkFVar ·.fvarId) ++ getRemainingArgs paramsInfo args
+    if let some declName ← findSpecCache? key then
+      trace[Compiler.specialize.step] "cached: {declName}"
+      return mkAppN (.const declName usNew) argsNew
+    else
+      let specDecl ← mkSpecDecl decl us argMask params decls levelParamsNew
+      trace[Compiler.specialize.step] "new: {specDecl.name}"
+      cacheSpec key specDecl.name
+      let specDecl ← specDecl.etaExpand
+      saveLCNFType specDecl.name specDecl.type
+      let specDecl ← specDecl.simp {} -- TODO: `simp` config
+      modify fun s => { s with decls := s.decls.push specDecl }
+      return mkAppN (.const specDecl.name usNew) argsNew
 
   partial def visitFunDecl (funDecl : FunDecl) : SpecializeM FunDecl := do
     let value ← withParams funDecl.params <| visitCode funDecl.value
@@ -330,11 +369,11 @@ mutual
   partial def visitCode (code : Code) : SpecializeM Code := do
     match code with
     | .let decl k =>
-      if let some k ← specializeApp? decl k then
-        visitCode k
-      else
-        let k ← withLetDecl decl <| visitCode k
-        return code.updateLet! decl k
+      let mut decl := decl
+      if let some value ← specializeApp? decl.value then
+        decl ← decl.updateValue value
+      let k ← withLetDecl decl <| visitCode k
+      return code.updateLet! decl k
     | .fun decl k | .jp decl k =>
       let decl ← visitFunDecl decl
       let k ← withFVar decl.fvarId <| visitCode k
@@ -372,5 +411,6 @@ def specialize : Pass where
 builtin_initialize
   registerTraceClass `Compiler.specialize (inherited := true)
   registerTraceClass `Compiler.specialize.candidate
+  registerTraceClass `Compiler.specialize.step
 
 end Lean.Compiler.LCNF

src/Lean/Compiler/LCNF/Types.lean

@@ -145,6 +145,13 @@ where
         result := mkApp result erasedExpr
     return result
 
+/--
+Save the LCNF type for the given declaration.
+-/
+def saveLCNFType (declName : Name) (type : Expr) : CoreM Unit := do
+  modifyEnv fun env =>
+    lcnfTypeExt.modifyState env fun s => { s with types := s.types.insert declName type }
+
 /--
 Return the LCNF type for the given declaration.
 -/
@@ -154,7 +161,7 @@ def getDeclLCNFType (declName : Name) : CoreM Expr := do
   | none =>
     let info ← getConstInfo declName
     let type ← Meta.MetaM.run' <| toLCNFType info.type
-    modifyEnv fun env => lcnfTypeExt.modifyState env fun s => { s with types := s.types.insert declName type }
+    saveLCNFType declName type
     return type
 
 /--

tests/lean/CompilerSimp.lean.expected.out

@@ -7,3 +7,7 @@
 [Compiler.test] Post condition test simpInlinesBinds for simp occurrence 1 successful
 [Compiler.test] Starting post condition test simpFix for simp occurrence 1
 [Compiler.test] Post condition test simpFix for simp occurrence 1 successful
+[Compiler.test] Starting post condition test simpInlinesBinds for simp occurrence 2
+[Compiler.test] Post condition test simpInlinesBinds for simp occurrence 2 successful
+[Compiler.test] Starting post condition test simpFix for simp occurrence 2
+[Compiler.test] Post condition test simpFix for simp occurrence 2 successful