chore: remove optimization incompatible with Mathlib OrderDual

This PR ensures the type resolution cache properly caches results for
type classe containing output parameters.

It ensures the cache key for a query like
```
HAppend.{0, 0, ?u} (BitVec 8) (BitVec 8) ?m
```
should be independent of the specific metavariable IDs in output
parameter positions. To achieve this, output parameter arguments are
erased from the cache key. Universe levels that only appear in output
parameter types (e.g., ?u corresponding to the result type's universe)
must also be erased to avoid cache misses when the same query is
issued with different universe metavariable IDs.

src/Lean/Compiler/InitAttr.lean

@@ -161,7 +161,9 @@ def declareBuiltin (forDecl : Name) (value : Expr) : CoreM Unit :=
 @[export lean_run_init_attrs]
 private unsafe def runInitAttrs (env : Environment) (opts : Options) : IO Unit := do
   if (← isInitializerExecutionEnabled) then
-    for mod in env.header.moduleNames, modIdx in 0...* do
+    -- **Note**: `ModuleIdx` is not an abbreviation, and we don't have instances for it.
+    -- Thus, we use `(modIdx : Nat)`
+    for mod in env.header.moduleNames, (modIdx : Nat) in 0...* do
       -- any native Lean code reachable by the interpreter (i.e. from shared
       -- libraries with their corresponding module in the Environment) must
       -- first be initialized

src/Lean/Meta/SynthInstance.lean

@@ -6,7 +6,6 @@ Authors: Daniel Selsam, Leonardo de Moura
 Type class instance synthesizer using tabled resolution.
 -/
 module
-
 prelude
 public import Init.Data.Array.InsertionSort
 public import Lean.Meta.Instances
@@ -15,7 +14,6 @@ public import Lean.Meta.Check
 import Init.While
 
 public section
-
 namespace Lean.Meta
 
 register_builtin_option synthInstance.maxHeartbeats : Nat := {
@@ -188,7 +186,7 @@ structure State where
   result?        : Option AbstractMVarsResult    := none
   generatorStack : Array GeneratorNode           := #[]
   resumeStack    : Array (ConsumerNode × Answer) := #[]
-  tableEntries   : Std.HashMap Expr TableEntry       := {}
+  tableEntries   : Std.HashMap Expr TableEntry   := {}
 
 abbrev SynthM := ReaderT Context $ StateRefT State MetaM
 
@@ -661,43 +659,120 @@ If it succeeds, and metavariables ?m_i have been assigned, we try to unify
 the original type `C a_1 ... a_n` with the normalized one.
 -/
 
-private def preprocess (type : Expr) : MetaM Expr :=
-  forallTelescopeReducing type fun xs type => do
-    let type ← whnf type
-    mkForallFVars xs type
+/-- Result kind for `preprocess` -/
+private inductive PreprocessKind where
+  | /--
+    Target type does not have metavariables.
+    We use the type to construct the cache key even if the class has output parameters.
+    Reason: we want to avoid the normalization step in this case.
+    -/
+    noMVars
+  | /-- Target type has metavariables, and class does not have output parameters. -/
+    mvarsNoOutputParams
+  | /-- Target type has metavariables, and class has output parameters. -/
+    mvarsOutputParams
 
-private partial def preprocessArgs (type : Expr) (i : Nat) (args : Array Expr) (outParamsPos : Array Nat) : MetaM (Array Expr) := do
-  if h : i < args.size then
-    let type ← whnf type
-    match type with
-    | .forallE _ d b _ => do
-      let arg := args[i]
-      /-
-      We should not simply check `d.isOutParam`. See `checkOutParam` and issue #1852.
-      If an instance implicit argument depends on an `outParam`, it is treated as an `outParam` too.
-      -/
-      let arg ← if outParamsPos.contains i then mkFreshExprMVar d else pure arg
-      let args := args.set i arg
-      preprocessArgs (b.instantiate1 arg) (i+1) args outParamsPos
-    | _ =>
-      throwError "type class resolution failed, insufficient number of arguments" -- TODO improve error message
-  else
-    return args
+/-- Return type for `preprocess` -/
+private structure PreprocessResult where
+  type         : Expr
+  cacheKeyType : Expr := type
+  kind         : PreprocessKind
 
-private def preprocessOutParam (type : Expr) : MetaM Expr :=
+/--
+Returns `{ type, cacheKeyType, hasOutParams }`, where `type` is the normalized type, and `cacheKeyType`
+is part of the key for the type class resolution cache. If the class associated with `type`
+does not have output parameters, then, `cacheKeyType` is `type`.
+If it has, we replace arguments corresponding with output parameters with wildcard terms.
+
+For example, the cache key for a query like
+`HAppend.{0, 0, ?u} (BitVec 8) (BitVec 8) ?m` should be independent of the specific
+metavariable IDs in output parameter positions. To achieve this, output parameter arguments
+are erased from the cache key. However, universe levels that only appear in output parameter
+types (e.g., `?u` corresponding to the result type's universe) must also be erased to avoid
+cache misses when the same query is issued with different universe metavariable IDs.
+-/
+private def preprocess (type : Expr) : MetaM PreprocessResult :=
+  let keyExprWildcard := mkFVar { name := `__wild__  }
+  let keyLevelWildcard := mkLevelParam `__wild__
+  forallTelescopeReducing type fun xs typeBody => do
+    let typeBody ← whnf typeBody
+    let type ← mkForallFVars xs typeBody
+    if !type.hasMVar then return { type, kind := .noMVars }
+    /-
+    **Note**: Workaround for classes such as `class ToLevel.{u}`. They do not have any parameters,
+    the universe parameter inference engine at `Class.lean` assumes `u` is an output parameter,
+    but this is not correct. We can remove this check after we update `Class.lean` and perform an
+    update stage0
+    -/
+    if typeBody.isConst then return { type, kind := .mvarsNoOutputParams }
+    let c := typeBody.getAppFn
+    let .const declName us := c | return { type, kind := .mvarsNoOutputParams }
+    let env ← getEnv
+    let some outParamsPos := getOutParamPositions? env declName | return { type, kind := .mvarsNoOutputParams }
+    let some outLevelParamPos := getOutLevelParamPositions? env declName | unreachable!
+    if outParamsPos.isEmpty && outLevelParamPos.isEmpty then return { type, kind := .mvarsNoOutputParams }
+    let c := if outLevelParamPos.isEmpty then c else
+      let rec normLevels (us : List Level) (i : Nat) : List Level :=
+        match us with
+        | [] => []
+        | u :: us =>
+          let u := if i ∈ outLevelParamPos then keyLevelWildcard else u
+          u :: normLevels us (i+1)
+      mkConst declName (normLevels us 0)
+    let rec norm (e : Expr) (i : Nat) : Expr :=
+      match e with
+      | .app f a =>
+        let a := if i ∈ outParamsPos then keyExprWildcard else a
+        mkApp (norm f (i-1)) a
+      | _ => c
+    let typeBody := norm typeBody (typeBody.getAppNumArgs - 1)
+    let cacheKeyType ← mkForallFVars xs typeBody
+    return { type, cacheKeyType, kind := .mvarsOutputParams }
+
+private partial def preprocessOutParam (type : Expr) : MetaM Expr :=
   forallTelescope type fun xs typeBody => do
-    match typeBody.getAppFn with
-    | c@(.const declName _) =>
-      let env ← getEnv
-      if let some outParamsPos := getOutParamPositions? env declName then
-        unless outParamsPos.isEmpty do
-          let args := typeBody.getAppArgs
-          let cType ← inferType c
-          let args ← preprocessArgs cType 0 args outParamsPos
-          return (← mkForallFVars xs (mkAppN c args))
-      return type
-    | _ =>
-      return type
+    /- **Note**: See similar test at preprocess. -/
+    if typeBody.isConst then return type
+    let c := typeBody.getAppFn
+    let .const declName us := c | return type
+    let env ← getEnv
+    let some outParamsPos := getOutParamPositions? env declName | return type
+    let some outLevelParamPos := getOutLevelParamPositions? env declName | unreachable!
+    if outParamsPos.isEmpty && outLevelParamPos.isEmpty then return type
+    let c ← if outLevelParamPos.isEmpty then pure c else
+      -- Replace universe parameters corresponding to output parameters with fresh universe metavariables.
+      let rec preprocessLevels (us : List Level) (i : Nat) : MetaM (List Level) := do
+        match us with
+        | [] => return []
+        | u :: us =>
+          let u ← if i ∈ outLevelParamPos then mkFreshLevelMVar else pure u
+          let us ← preprocessLevels us (i+1)
+          return u :: us
+      pure <| mkConst declName (← preprocessLevels us 0)
+    let rec preprocessArgs (type : Expr) (i : Nat) (args : Array Expr) : MetaM (Array Expr) := do
+      if h : i < args.size then
+        let type ← whnf type
+        match type with
+        | .forallE _ d b _ => do
+          let arg := args[i]
+          /-
+          We should not simply check `d.isOutParam`. See `checkOutParam` and issue #1852.
+          If an instance implicit argument depends on an `outParam`, it is treated as an `outParam` too.
+          -/
+          let arg ← if outParamsPos.contains i then mkFreshExprMVar d else pure arg
+          let args := args.set i arg
+          preprocessArgs (b.instantiate1 arg) (i+1) args
+        | _ =>
+          throwError "type class resolution failed, insufficient number of arguments" -- TODO improve error message
+      else
+        return args
+    let args := typeBody.getAppArgs
+    if outParamsPos.isEmpty then
+      mkForallFVars xs (mkAppN c args)
+    else
+      let cType ← inferType c
+      let args ← preprocessArgs cType 0 args
+      mkForallFVars xs (mkAppN c args)
 
 /-!
   Remark: when `maxResultSize? == none`, the configuration option `synthInstance.maxResultSize` is used.
@@ -706,9 +781,21 @@ private def preprocessOutParam (type : Expr) : MetaM Expr :=
 
 private def assignOutParams (type : Expr) (result : Expr) : MetaM Bool := do
   let resultType ← inferType result
-  /- Output parameters of local instances may be marked as `syntheticOpaque` by the application-elaborator.
-      We use `withAssignableSyntheticOpaque` to make sure this kind of parameter can be assigned by the following `isDefEq`.
-      TODO: rewrite this check to avoid `withAssignableSyntheticOpaque`. -/
+  /-
+  Output parameters of local instances may be marked as `syntheticOpaque` by the application-elaborator.
+  We use `withAssignableSyntheticOpaque` to make sure this kind of parameter can be assigned by the following `isDefEq`.
+  TODO: rewrite this check to avoid `withAssignableSyntheticOpaque`.
+
+  **Note**: We tried to remove `withDefault` at the following `isDefEq` because it was a potential performance footgun. TC is supposed to unfold only `reducible` definitions and `instances`.
+  We reverted the change because it triggered thousands of failures related to the `OrderDual` type. Example:
+  ```
+  variable {ι : Type}
+  def OrderDual (α : Type) : Type := α
+  instance [I : DecidableEq ι] : DecidableEq (OrderDual ι) := inferInstance -- Failure
+  ```
+  Mathlib developers are currently trying to refactor the `OrderDual` declaration,
+  but it will take time. We will try to remove the `withDefault` again after the refactoring.
+  -/
   let defEq ← withDefault <| withAssignableSyntheticOpaque <| isDefEq type resultType
   unless defEq do
     trace[Meta.synthInstance] "{crossEmoji} result type{indentExpr resultType}\nis not definitionally equal to{indentExpr type}"
@@ -769,15 +856,18 @@ private def applyCachedAbstractResult? (type : Expr) (abstResult? : Option Abstr
     applyAbstractResult? type abstResult?
 
 /-- Helper function for caching synthesized type class instances. -/
-private def cacheResult (cacheKey : SynthInstanceCacheKey) (abstResult? : Option AbstractMVarsResult) (result? : Option Expr) : MetaM Unit := do
-  match result? with
+private def cacheResult (cacheKey : SynthInstanceCacheKey) (kind : PreprocessKind) (abstResult? : Option AbstractMVarsResult) (result? : Option Expr) : MetaM Unit := do
+  -- **TODO**: simplify this function.
+  match abstResult? with
   | none => modify fun s => { s with cache.synthInstance := s.cache.synthInstance.insert cacheKey none }
-  | some result =>
-    let some abstResult := abstResult? | return ()
-    if abstResult.numMVars == 0 && abstResult.paramNames.isEmpty then
-      -- See `applyCachedAbstractResult?` If new metavariables have **not** been introduced,
-      -- we don't need to perform extra checks again when reusing result.
-      modify fun s => { s with cache.synthInstance := s.cache.synthInstance.insert cacheKey (some { expr := result, paramNames := #[], mvars := #[] }) }
+  | some abstResult =>
+    if abstResult.numMVars == 0 && abstResult.paramNames.isEmpty && kind matches .noMVars | .mvarsNoOutputParams then
+      match result? with
+      | none => modify fun s => { s with cache.synthInstance := s.cache.synthInstance.insert cacheKey none }
+      | some result =>
+        -- See `applyCachedAbstractResult?` If new metavariables have **not** been introduced,
+        -- we don't need to perform extra checks again when reusing result.
+        modify fun s => { s with cache.synthInstance := s.cache.synthInstance.insert cacheKey (some { expr := result, paramNames := #[], mvars := #[] }) }
     else
       modify fun s => { s with cache.synthInstance := s.cache.synthInstance.insert cacheKey (some abstResult) }
 
@@ -791,20 +881,44 @@ def synthInstanceCore? (type : Expr) (maxResultSize? : Option Nat := none) : Met
   withInTypeClassResolution do
     let localInsts ← getLocalInstances
     let type ← instantiateMVars type
-    let type ← preprocess type
-    let cacheKey := { localInsts, type, synthPendingDepth := (← read).synthPendingDepth }
+    let { type, cacheKeyType, kind } ← preprocess type
+    let cacheKey := { localInsts, type := cacheKeyType, synthPendingDepth := (← read).synthPendingDepth }
     match (← get).cache.synthInstance.find? cacheKey with
     | some abstResult? =>
+      trace[Meta.synthInstance.cache] "cached: {type}"
       let result? ← applyCachedAbstractResult? type abstResult?
       trace[Meta.synthInstance] "result {result?} (cached)"
       return result?
     | none =>
+      trace[Meta.synthInstance.cache] "new: {type}"
       let abstResult? ← withNewMCtxDepth (allowLevelAssignments := true) do
-        let normType ← preprocessOutParam type
-        SynthInstance.main normType maxResultSize
+        match kind with
+        | .noMVars =>
+          /-
+          **Note**: The expensive `preprocessOutParam` step is morally **not** needed here because
+          the output params should be uniquely determined by the input params. During type class
+          resolution, definitional equality only unfolds `[reducible]` and `[instance_reducible]`
+          declarations. This is a contract with our users to ensure performance is reasonable.
+          However, the same `OrderDual` declaration that creates problems for `assignOutParams`
+          also prevents us from using this optimization. As an example, suppose we are trying to
+          synthesize
+          ```
+          FunLike F (OrderDual α) (OrderDual β)
+          ```
+          where the last two arguments of `FunLike` are output parameters. This term has no
+          metavariables, and it seems natural to skip `preprocessOutParam`, which would replace
+          the last two arguments with metavariables. However, if we don't replace them,
+          TC resolution fails because it cannot unfold `OrderDual` since it is semireducible.
+
+          **Note**: We should remove `preprocessOutParam` from the following line as soon as
+          Mathlib refactors `OrderDual`.
+          -/
+          SynthInstance.main (← preprocessOutParam type) maxResultSize
+        | .mvarsNoOutputParams => SynthInstance.main type maxResultSize
+        | .mvarsOutputParams => SynthInstance.main (← preprocessOutParam type) maxResultSize
       let result? ← applyAbstractResult? type abstResult?
       trace[Meta.synthInstance] "result {result?}"
-      cacheResult cacheKey abstResult? result?
+      cacheResult cacheKey kind abstResult? result?
       return result?
 
 def synthInstance? (type : Expr) (maxResultSize? : Option Nat := none) : MetaM (Option Expr) := do profileitM Exception "typeclass inference" (← getOptions) (decl := type.getAppFn.constName?.getD .anonymous) do
@@ -874,5 +988,6 @@ builtin_initialize
   registerTraceClass `Meta.synthInstance.resume (inherited := true)
   registerTraceClass `Meta.synthInstance.unusedArgs
   registerTraceClass `Meta.synthInstance.newAnswer
+  registerTraceClass `Meta.synthInstance.cache
 
 end Lean.Meta

src/lake/Lake/Build/Fetch.lean

@@ -25,7 +25,7 @@ using the `fetch` function defined in this module.
 namespace Lake
 
 /-- A type alias for `Option Package` that assists monad type class synthesis. -/
-@[expose] public def CurrPackage := Option Package
+@[expose] public abbrev CurrPackage := Option Package
 
 /-- Run the action `x` with `pkg?` as the current package or no package if `none`. -/
 @[inline] public def withCurrPackage? [MonadWithReader CurrPackage m] (pkg? : Option Package) (x : m α): m α :=

src/lake/Lake/CLI/Shake.lean

@@ -257,7 +257,7 @@ def calcNeeds (s : State) (i : ModuleIdx) : Needs := Id.run do
       needs := visitExpr k e needs
 
   for use in getExtraModUses env i do
-    let j := env.getModuleIdx? use.module |>.get!
+    let j : Nat := env.getModuleIdx? use.module |>.get!
     needs := needs.union { use with } {j}
 
   return needs
@@ -268,11 +268,11 @@ where
     Lean.Expr.foldConsts e deps fun c deps => Id.run do
       let mut deps := deps
       if let some c := getDepConstName? s c then
-        if let some j := env.getModuleIdxFor? c then
+        if let some (j : Nat) := env.getModuleIdxFor? c then
           let k := { k with isMeta := k.isMeta && !isDeclMeta' env c }
           if j != i then
             deps := deps.union k {j}
-            for indMod in s.indirectModUses[c]?.getD #[] do
+            for (indMod : Nat) in s.indirectModUses[c]?.getD #[] do
               if s.transDeps[i]!.has k indMod then
                 deps := deps.union k {indMod}
       return deps
@@ -427,7 +427,7 @@ def visitModule (pkgs : Array Name) (srcSearchPath : SearchPath)
   -- Add additional preserved imports
   for impStx in imports do
     let imp := decodeImport impStx
-    let j := s.env.getModuleIdx? imp.module |>.get!
+    let j : Nat := s.env.getModuleIdx? imp.module |>.get!
     let k := NeedsKind.ofImport imp
     if addOnly ||
         -- TODO: allow per-library configuration instead of hardcoding `Init`
@@ -456,13 +456,14 @@ def visitModule (pkgs : Array Name) (srcSearchPath : SearchPath)
             deps := deps.sub k' (transDeps.sub k' {j} |>.get k')
 
   if prelude?.isNone then
-    deps := deps.union .pub {s.env.getModuleIdx? `Init |>.get!}
+    let j : Nat := s.env.getModuleIdx? `Init |>.get!
+    deps := deps.union .pub {j}
 
   -- Accumulate `transDeps` which is the non-reflexive transitive closure of the still-live imports
   let mut transDeps := Needs.empty
   let mut alwaysAdd : Array Import := #[]  -- to be added even if implied by other imports
   for imp in s.mods[i]!.imports do
-    let j := s.env.getModuleIdx? imp.module |>.get!
+    let j : Nat := s.env.getModuleIdx? imp.module |>.get!
     let k := NeedsKind.ofImport imp
     if deps.has k j || imp.importAll then
       transDeps := addTransitiveImps transDeps imp j s.transDeps[j]!

src/lake/Lake/Util/Cli.lean

@@ -17,8 +17,7 @@ Defines the abstract CLI interface for Lake.
 
 /-! # Types -/
 
-@[expose]  -- for codegen
-public def ArgList := List String
+public abbrev ArgList := List String
 
 @[inline] public def ArgList.mk (args : List String) : ArgList :=
   args

tests/lean/run/tc_cache.lean

@@ -0,0 +1,31 @@
+/-
+Type class resolution cache.
+Recall that we normalize keys for type class with output parameters only when the input type
+contains metavariables. This is why in the following example we sold
+```
+new: HAppend (List Nat) (List Nat) ?_
+```
+and
+```
+new: HAppend (List Nat) (List Nat) (List Nat)
+```
+-/
+
+set_option pp.mvars.anonymous false
+set_option trace.Meta.synthInstance.cache true
+/--
+trace: [Meta.synthInstance.cache] cached: HAppend (List Nat) (List Nat) (List Nat)
+[Meta.synthInstance.cache] cached: HAppend (List Nat) (List Nat) ?_
+---
+trace: [Meta.synthInstance.cache] cached: HAppend (List Nat) (List Nat) (List Nat)
+[Meta.synthInstance.cache] cached: HAppend (List Nat) (List Nat) ?_
+---
+trace: [Meta.synthInstance.cache] cached: HAppend (List Nat) (List Nat) (List Nat)
+[Meta.synthInstance.cache] new: HAppend (List Nat) (List Nat) ?_
+---
+trace: [Meta.synthInstance.cache] new: HAppend (List Nat) (List Nat) (List Nat)
+[Meta.synthInstance.cache] cached: HAppend (List Nat) (List Nat) ?_
+-/
+#guard_msgs (ordering := sorted) in
+def ex (a : List Nat) : List Nat :=
+  a ++ a ++ a ++ a ++ a