diff --git a/src/Lean/Compiler/InitAttr.lean b/src/Lean/Compiler/InitAttr.lean index c2998f9db0..d49ae8f9a2 100644 --- a/src/Lean/Compiler/InitAttr.lean +++ b/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 diff --git a/src/Lean/Meta/SynthInstance.lean b/src/Lean/Meta/SynthInstance.lean index b7216e64fc..43d828e069 100644 --- a/src/Lean/Meta/SynthInstance.lean +++ b/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 diff --git a/src/lake/Lake/Build/Fetch.lean b/src/lake/Lake/Build/Fetch.lean index 8b8adf7a05..67b43f10fb 100644 --- a/src/lake/Lake/Build/Fetch.lean +++ b/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 α := diff --git a/src/lake/Lake/CLI/Shake.lean b/src/lake/Lake/CLI/Shake.lean index 326573ee5e..8cdf4db70d 100644 --- a/src/lake/Lake/CLI/Shake.lean +++ b/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]! diff --git a/src/lake/Lake/Util/Cli.lean b/src/lake/Lake/Util/Cli.lean index 79d4d76fb1..22be41acb3 100644 --- a/src/lake/Lake/Util/Cli.lean +++ b/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 diff --git a/tests/lean/run/tc_cache.lean b/tests/lean/run/tc_cache.lean new file mode 100644 index 0000000000..00bb01022b --- /dev/null +++ b/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