Merge remote-tracking branch 'origin/master' into defEqCache

Co-authored-by: Timo <timorcb@gmail.com>

src/Lean/Meta/Basic.lean

@@ -202,10 +202,15 @@ abbrev FunInfoCache   := PersistentHashMap InfoCacheKey FunInfo
 abbrev WhnfCache      := PersistentExprStructMap Expr
 
 /--
-  A mapping `(s, t) ↦ isDefEq s t`.
+  A mapping `(s, t) ↦ isDefEq s t` per transparency level.
   TODO: consider more efficient representations (e.g., a proper set) and caching policies (e.g., imperfect cache).
   We should also investigate the impact on memory consumption. -/
-abbrev DefEqCache := PersistentHashMap (Expr × Expr) Bool
+structure DefEqCache where
+  reducible : PersistentHashMap (Expr × Expr) Bool := {}
+  instances : PersistentHashMap (Expr × Expr) Bool := {}
+  default   : PersistentHashMap (Expr × Expr) Bool := {}
+  all       : PersistentHashMap (Expr × Expr) Bool := {}
+  deriving Inhabited
 
 /--
   Cache datastructures for type inference, type class resolution, whnf, and definitional equality.
@@ -216,7 +221,8 @@ structure Cache where
   synthInstance  : SynthInstanceCache := {}
   whnfDefault    : WhnfCache := {} -- cache for closed terms and `TransparencyMode.default`
   whnfAll        : WhnfCache := {} -- cache for closed terms and `TransparencyMode.all`
-  defEq          : DefEqCache := {}
+  defEqTrans     : DefEqCache := {} -- transient cache for terms containing mvars or using nonstandard configuration options, it is frequently reset.
+  defEqPerm      : DefEqCache := {} -- permanent cache for terms not containing mvars and using standard configuration options
   deriving Inhabited
 
 /--
@@ -363,10 +369,16 @@ variable [MonadControlT MetaM n] [Monad n]
   modify fun ⟨mctx, cache, zetaFVarIds, postponed⟩ => ⟨mctx, f cache, zetaFVarIds, postponed⟩
 
 @[inline] def modifyInferTypeCache (f : InferTypeCache → InferTypeCache) : MetaM Unit :=
-  modifyCache fun ⟨ic, c1, c2, c3, c4, c5⟩ => ⟨f ic, c1, c2, c3, c4, c5⟩
+  modifyCache fun ⟨ic, c1, c2, c3, c4, c5, c6⟩ => ⟨f ic, c1, c2, c3, c4, c5, c6⟩
 
-@[inline] def modifyDefEqCache (f : DefEqCache → DefEqCache) : MetaM Unit :=
-  modifyCache fun ⟨c1, c2, c3, c4, c5, defeq⟩ => ⟨c1, c2, c3, c4, c5, f defeq⟩
+@[inline] def modifyDefEqTransientCache (f : DefEqCache → DefEqCache) : MetaM Unit :=
+  modifyCache fun ⟨c1, c2, c3, c4, c5, defeqTrans, c6⟩ => ⟨c1, c2, c3, c4, c5, f defeqTrans, c6⟩
+
+@[inline] def modifyDefEqPermCache (f : DefEqCache → DefEqCache) : MetaM Unit :=
+  modifyCache fun ⟨c1, c2, c3, c4, c5, c6, defeqPerm⟩ => ⟨c1, c2, c3, c4, c5, c6, f defeqPerm⟩
+
+@[inline] def resetDefEqPermCaches : MetaM Unit :=
+  modifyDefEqPermCache fun _ => {}
 
 def getLocalInstances : MetaM LocalInstances :=
   return (← read).localInstances
@@ -1601,7 +1613,7 @@ partial def processPostponed (mayPostpone : Bool := true) (exceptionOnFailure :=
     See issue #1102 for an example that triggers an exponential blowup if we don't use this more
     aggressive form of caching.
   -/
-  modifyDefEqCache fun _ => {}
+  modifyDefEqTransientCache fun _ => {}
   let postponed ← getResetPostponed
   try
     if (← x) then
@@ -1628,6 +1640,26 @@ def isLevelDefEq (u v : Level) : MetaM Bool :=
 /-- See `isDefEq`. -/
 def isExprDefEq (t s : Expr) : MetaM Bool :=
   withReader (fun ctx => { ctx with defEqCtx? := some { lhs := t, rhs := s, lctx := ctx.lctx, localInstances := ctx.localInstances } }) do
+    /-
+    The following `resetDefEqPermCaches` is a workaround. Without it the test suite fails, and we probably cannot compile complex libraries such as Mathlib.
+    TODO: investigate why we need this reset.
+    Some conjectures:
+    - It is not enough to check whether `t` and `s` do not contain metavariables. We would need to check the type
+      of all local variables `t` and `s` depend on. If the local variables contain metavariables, the result of `isDefEq` may change if these
+      variables are instantiated.
+    - Related to the previous one: the operation
+      ```lean
+      _root_.Lean.MVarId.replaceLocalDeclDefEq (mvarId : MVarId) (fvarId : FVarId) (typeNew : Expr)
+      ```
+      is probably being misused. We are probably using it to replace a `type` with `typeNew` where these two types
+      are definitionally equal IFF we can assign the metavariables in `type`.
+
+    Possible fix: always generate new `FVarId`s when update the type of local variables.
+    Drawback: this operation can be quite expensive, and we must evaluate whether it is worth doing to remove the following `reset`.
+
+    Remark: the kernel does *not* update the type of variables in the local context.
+    -/
+    resetDefEqPermCaches
     checkpointDefEq (mayPostpone := true) <| Meta.isExprDefEqAux t s
 
 /--

src/Lean/Meta/ExprDefEq.lean

@@ -1792,22 +1792,61 @@ private def isExprDefEqExpensive (t : Expr) (s : Expr) : MetaM Bool := do
     if (← isDefEqUnitLike t s) then return true else
     isDefEqOnFailure t s
 
-private def mkCacheKey (t : Expr) (s : Expr) : Expr × Expr :=
-  if Expr.quickLt t s then (t, s) else (s, t)
+inductive DefEqCacheKind where
+  | transient -- problem has mvars or is using nonstandard configuration, we should use transient cache
+  | permanent -- problem does not have mvars and we are using standard config, we can use one persistent cache.
 
-private def getCachedResult (key : Expr × Expr) : MetaM LBool := do
-  match (← get).cache.defEq.find? key with
+private def getDefEqCacheKind (t s : Expr) : MetaM DefEqCacheKind := do
+  if t.hasMVar || s.hasMVar || (← read).canUnfold?.isSome then
+    return .transient
+  else
+    return .permanent
+
+/--
+Structure for storing defeq cache key information.
+-/
+structure DefEqCacheKeyInfo where
+  kind : DefEqCacheKind
+  key  : Expr × Expr
+
+private def mkCacheKey (t s : Expr) : MetaM DefEqCacheKeyInfo := do
+  let kind ← getDefEqCacheKind t s
+  let key := if Expr.quickLt t s then (t, s) else (s, t)
+  return { key, kind }
+
+private def getCachedResult (keyInfo : DefEqCacheKeyInfo) : MetaM LBool := do
+  let cache ← match keyInfo.kind with
+    | .transient => pure (← get).cache.defEqTrans
+    | .permanent => pure (← get).cache.defEqPerm
+  let cache := match (← getTransparency) with
+    | .reducible => cache.reducible
+    | .instances => cache.instances
+    | .default   => cache.default
+    | .all       => cache.all
+  match cache.find? keyInfo.key with
   | some val => return val.toLBool
   | none => return .undef
 
-private def cacheResult (key : Expr × Expr) (result : Bool) : MetaM Unit := do
-  /-
-  We must ensure that all assigned metavariables in the key are replaced by their current assignments.
-  Otherwise, the key is invalid after the assignment is "backtracked".
-  See issue #1870 for an example.
-  -/
-  let key := (← instantiateMVars key.1, ← instantiateMVars key.2)
-  modifyDefEqCache fun c => c.insert key result
+def DefEqCache.update (cache : DefEqCache) (mode : TransparencyMode) (key : Expr × Expr) (result : Bool) : DefEqCache :=
+  match mode with
+  | .reducible => { cache with reducible := cache.reducible.insert key result }
+  | .instances => { cache with instances := cache.instances.insert key result }
+  | .default   => { cache with default   := cache.default.insert key result }
+  | .all       => { cache with all       := cache.all.insert key result }
+
+private def cacheResult (keyInfo : DefEqCacheKeyInfo) (result : Bool) : MetaM Unit := do
+  let mode ← getTransparency
+  let key := keyInfo.key
+  match keyInfo.kind with
+  | .permanent => modifyDefEqPermCache fun c => c.update mode key result
+  | .transient =>
+    /-
+    We must ensure that all assigned metavariables in the key are replaced by their current assignments.
+    Otherwise, the key is invalid after the assignment is "backtracked".
+    See issue #1870 for an example.
+    -/
+    let key := (← instantiateMVars key.1, ← instantiateMVars key.2)
+    modifyDefEqTransientCache fun c => c.update mode key result
 
 @[export lean_is_expr_def_eq]
 partial def isExprDefEqAuxImpl (t : Expr) (s : Expr) : MetaM Bool := withIncRecDepth do
@@ -1839,7 +1878,7 @@ partial def isExprDefEqAuxImpl (t : Expr) (s : Expr) : MetaM Bool := withIncRecD
     let t ← instantiateMVars t
     let s ← instantiateMVars s
     let numPostponed ← getNumPostponed
-    let k := mkCacheKey t s
+    let k ← mkCacheKey t s
     match (← getCachedResult k) with
     | .true  =>
       trace[Meta.isDefEq.cache] "cache hit 'true' for {t} =?= {s}"