feat: modify withNewMCtxDepth

src/Lean/Meta/AbstractMVars.lean

@@ -51,7 +51,7 @@ private partial def abstractLevelMVars (u : Level) : M Level := do
     | Level.mvar mvarId =>
       let s ← get
       let depth := s.mctx.getLevelDepth mvarId;
-      if depth != s.mctx.depth then
+      if depth < s.mctx.depth then
         return u -- metavariables from lower depths are treated as constants
       else
         match s.lmap[mvarId]? with

src/Lean/Meta/Basic.lean

@@ -758,7 +758,7 @@ Return true if the given metavariable is "read-only".
 That is, its `depth` is different from the current metavariable context depth.
 -/
 def _root_.Lean.MVarId.isReadOnly (mvarId : MVarId) : MetaM Bool := do
-  return (← mvarId.getDecl).depth != (← getMCtx).depth
+  return (← mvarId.getDecl).depth < (← getMCtx).depth
 
 /--
 Returns true if `mvarId.isReadOnly` returns true or if `mvarId` is a synthetic opaque metavariable.
@@ -767,7 +767,7 @@ Recall `isDefEq` will not assign a value to `mvarId` if `mvarId.isReadOnlyOrSynt
 -/
 def _root_.Lean.MVarId.isReadOnlyOrSyntheticOpaque (mvarId : MVarId) : MetaM Bool := do
   let mvarDecl ← mvarId.getDecl
-  if mvarDecl.depth != (← getMCtx).depth then
+  if mvarDecl.depth < (← getMCtx).depth then
     return true
   else
     match mvarDecl.kind with
@@ -1509,12 +1509,12 @@ def withExistingLocalDecls (decls : List LocalDecl) : n α → n α :=
   mapMetaM <| withExistingLocalDeclsImp decls
 
 private def withNewMCtxDepthImp (allowLevelAssignments : Bool) (x : MetaM α) : MetaM α := do
-  let saved ← get
+  let { mctx := { depth, levelAssignDepth, .. }, postponed, .. } ← get
   modify fun s => { s with mctx := s.mctx.incDepth allowLevelAssignments, postponed := {} }
   try
     x
   finally
-    modify fun s => { s with mctx := saved.mctx, postponed := saved.postponed }
+    modify fun s => { s with mctx := { s.mctx with depth, levelAssignDepth }, postponed }
 
 /--
 Removes `fvarId` from the local context, and replaces occurrences of it with `e`.

src/Lean/Meta/ExprDefEq.lean

@@ -755,7 +755,7 @@ mutual
       /- The local context of `mvar` - free variables being abstracted is a subprefix of the metavariable being assigned.
          We "subtract" variables being abstracted because we use `elimMVarDeps` -/
       return mvar
-    if mvarDecl.depth != (← getMCtx).depth || mvarDecl.kind.isSyntheticOpaque then
+    if mvarDecl.depth < (← getMCtx).depth || mvarDecl.kind.isSyntheticOpaque then
       traceM `Meta.isDefEq.assign.readOnlyMVarWithBiggerLCtx <| addAssignmentInfo (mkMVar mvarId)
       throwCheckAssignmentFailure
     let ctxMeta ← readThe Meta.Context

src/Lean/Meta/SynthInstance.lean

@@ -116,7 +116,7 @@ partial def normLevel (u : Level) : M Level := do
     | .max v w     => return u.updateMax! (← normLevel v) (← normLevel w)
     | .imax v w    => return u.updateIMax! (← normLevel v) (← normLevel w)
     | .mvar mvarId =>
-      if (← getMCtx).getLevelDepth mvarId != (← getMCtx).depth then
+      if (← getMCtx).getLevelDepth mvarId < (← getMCtx).depth then
         return u
       else
         let s ← get

src/Lean/MetavarContext.lean

@@ -320,6 +320,7 @@ For more information on specifics see the comment in the file that `MetavarConte
 structure MetavarContext where
   /-- Depth is used to control whether an mvar can be assigned in unification. -/
   depth          : Nat := 0
+  maxDepth       : Nat := 0
   /-- At what depth level mvars can be assigned. -/
   levelAssignDepth : Nat := 0
   /-- Counter for setting the field `index` at `MetavarDecl` -/
@@ -434,7 +435,7 @@ def MetavarContext.getDecl (mctx : MetavarContext) (mvarId : MVarId) : MetavarDe
 def _root_.Lean.MVarId.isAssignable [Monad m] [MonadMCtx m] (mvarId : MVarId) : m Bool := do
   let mctx ← getMCtx
   let decl := mctx.getDecl mvarId
-  return decl.depth == mctx.depth
+  return decl.depth >= mctx.depth
 
 /-- Return true iff the given level contains an assigned metavariable. -/
 def hasAssignedLevelMVar [Monad m] [MonadMCtx m] : Level → m Bool
@@ -874,10 +875,10 @@ def isAnonymousMVar (mctx : MetavarContext) (mvarId : MVarId) : Bool :=
   | some mvarDecl => mvarDecl.userName.isAnonymous
 
 def incDepth (mctx : MetavarContext) (allowLevelAssignments := false) : MetavarContext :=
-  let depth := mctx.depth + 1
+  let depth := mctx.maxDepth + 1
   let levelAssignDepth :=
     if allowLevelAssignments then mctx.levelAssignDepth else depth
-  { mctx with depth, levelAssignDepth }
+  { mctx with depth, maxDepth := depth, levelAssignDepth }
 
 instance : MonadMCtx (StateRefT MetavarContext (ST ω)) where
   getMCtx    := get

src/Lean/PrettyPrinter/Delaborator/TopDownAnalyze.lean

@@ -165,13 +165,13 @@ partial def hasMVarAtCurrDepth (e : Expr) : MetaM Bool := do
   let mctx ← getMCtx
   return Option.isSome <| e.findMVar? fun mvarId =>
     match mctx.findDecl? mvarId with
-    | some mdecl => mdecl.depth == mctx.depth
+    | some mdecl => mdecl.depth >= mctx.depth
     | _ => false
 
 partial def hasLevelMVarAtCurrDepth (e : Expr) : MetaM Bool := do
   let mctx ← getMCtx
   return Option.isSome <| e.findLevelMVar? fun mvarId =>
-    mctx.findLevelDepth? mvarId == some mctx.depth
+    (mctx.findLevelDepth? mvarId).elim false fun depth => depth >= mctx.depth
 
 private def valUnknown (e : Expr) : MetaM Bool := do
   hasMVarAtCurrDepth (← instantiateMVars e)

src/Std/Data/DHashMap/Basic.lean

@@ -152,6 +152,18 @@ variable {β : Type v}
 
 end
 
+@[inline, inherit_doc Raw.getKey?] def getKey? (m : DHashMap α β) (a : α) : Option α :=
+  Raw₀.getKey? ⟨m.1, m.2.size_buckets_pos⟩ a
+
+@[inline, inherit_doc Raw.getKey] def getKey (m : DHashMap α β) (a : α) (h : a ∈ m) : α :=
+  Raw₀.getKey ⟨m.1, m.2.size_buckets_pos⟩ a h
+
+@[inline, inherit_doc Raw.getKey!] def getKey! [Inhabited α] (m : DHashMap α β) (a : α) : α :=
+  Raw₀.getKey! ⟨m.1, m.2.size_buckets_pos⟩ a
+
+@[inline, inherit_doc Raw.getKeyD] def getKeyD (m : DHashMap α β) (a : α) (fallback : α) : α :=
+  Raw₀.getKeyD ⟨m.1, m.2.size_buckets_pos⟩ a fallback
+
 @[inline, inherit_doc Raw.size] def size (m : DHashMap α β) : Nat :=
   m.1.size
 

src/Std/Data/DHashMap/Internal/AssocList/Basic.lean

@@ -102,16 +102,31 @@ def getCast [BEq α] [LawfulBEq α] (a : α) : (l : AssocList α β) → l.conta
   | cons k v es, h => if hka : k == a then cast (congrArg β (eq_of_beq hka)) v
       else es.getCast a (by rw [← h, contains, Bool.of_not_eq_true hka, Bool.false_or])
 
+/-- Internal implementation detail of the hash map -/
+def getKey [BEq α] (a : α) : (l : AssocList α β) → l.contains a → α
+  | cons k _ es, h => if hka : k == a then k
+      else es.getKey a (by rw [← h, contains, Bool.of_not_eq_true hka, Bool.false_or])
+
 /-- Internal implementation detail of the hash map -/
 def getCast! [BEq α] [LawfulBEq α] (a : α) [Inhabited (β a)] : AssocList α β → β a
   | nil => panic! "key is not present in hash table"
   | cons k v es => if h : k == a then cast (congrArg β (eq_of_beq h)) v else es.getCast! a
 
+/-- Internal implementation detail of the hash map -/
+def getKey? [BEq α] (a : α) : AssocList α β → Option α
+  | nil => none
+  | cons k _ es => if k == a then some k else es.getKey? a
+
 /-- Internal implementation detail of the hash map -/
 def get! {β : Type v} [BEq α] [Inhabited β] (a : α) : AssocList α (fun _ => β) → β
   | nil => panic! "key is not present in hash table"
   | cons k v es => bif k == a then v else es.get! a
 
+/-- Internal implementation detail of the hash map -/
+def getKey! [BEq α] [Inhabited α] (a : α) : AssocList α β → α
+  | nil => panic! "key is not present in hash table"
+  | cons k _ es => if k == a then k else es.getKey! a
+
 /-- Internal implementation detail of the hash map -/
 def getCastD [BEq α] [LawfulBEq α] (a : α) (fallback : β a) : AssocList α β → β a
   | nil => fallback
@@ -123,6 +138,11 @@ def getD {β : Type v} [BEq α] (a : α) (fallback : β) : AssocList α (fun _ =
   | nil => fallback
   | cons k v es => bif k == a then v else es.getD a fallback
 
+/-- Internal implementation detail of the hash map -/
+def getKeyD [BEq α] (a : α) (fallback : α) : AssocList α β → α
+  | nil => fallback
+  | cons k _ es => if k == a then k else es.getKeyD a fallback
+
 /-- Internal implementation detail of the hash map -/
 def replace [BEq α] (a : α) (b : β a) : AssocList α β → AssocList α β
   | nil => nil

src/Std/Data/DHashMap/Internal/AssocList/Lemmas.lean

@@ -112,6 +112,32 @@ theorem get!_eq {β : Type v} [BEq α] [Inhabited β] {l : AssocList α (fun _ =
   · simp_all [get!, List.getValue!, List.getValue!, List.getValue?_cons,
       Bool.apply_cond Option.get!]
 
+@[simp]
+theorem getKey?_eq [BEq α] {l : AssocList α β} {a : α} :
+    l.getKey? a = List.getKey? a l.toList := by
+  induction l <;> simp_all [getKey?]
+
+@[simp]
+theorem getKey_eq [BEq α] {l : AssocList α β} {a : α} {h} :
+    l.getKey a h = List.getKey a l.toList (contains_eq.symm.trans h) := by
+  induction l
+  · simp [contains] at h
+  · next k v t ih => simp only [getKey, toList_cons, List.getKey_cons, ih]
+
+@[simp]
+theorem getKeyD_eq [BEq α] {l : AssocList α β} {a fallback : α} :
+    l.getKeyD a fallback = List.getKeyD a l.toList fallback := by
+  induction l
+  · simp [getKeyD, List.getKeyD]
+  · simp_all [getKeyD, List.getKeyD, Bool.apply_cond (fun x => Option.getD x fallback)]
+
+@[simp]
+theorem getKey!_eq [BEq α] [Inhabited α] {l : AssocList α β} {a : α} :
+    l.getKey! a = List.getKey! a l.toList := by
+  induction l
+  · simp [getKey!, List.getKey!]
+  · simp_all [getKey!, List.getKey!, Bool.apply_cond Option.get!]
+
 @[simp]
 theorem toList_replace [BEq α] {l : AssocList α β} {a : α} {b : β a} :
     (l.replace a b).toList = replaceEntry a b l.toList := by

src/Std/Data/DHashMap/Internal/Defs.lean

@@ -100,9 +100,9 @@ Here is a summary of the steps required to add and verify a new operation:
   * Connect the implementation on lists and associative lists in `Internal.AssocList.Lemmas` via a
     lemma `AssocList.operation_eq`.
 3. Write the model implementation
-  * Write the model implementation `Raw₀.operationₘ` in `Internal.List.Model`
+  * Write the model implementation `Raw₀.operationₘ` in `Internal.Model`
   * Prove that the model implementation is equal to the actual implementation in
-    `Internal.List.Model` via a lemma `operation_eq_operationₘ`.
+    `Internal.Model` via a lemma `operation_eq_operationₘ`.
 4. Verify the model implementation
   * In `Internal.WF`, prove `operationₘ_eq_List.operation` (for access operations) or
     `wfImp_operationₘ` and `toListModel_operationₘ`
@@ -121,18 +121,18 @@ Here is a summary of the steps required to add and verify a new operation:
     might also have to prove that your list operation is invariant under permutation and add that to
     the tactic.
 7. State and prove the user-facing lemmas
-  * Restate all of your lemmas for `DHashMap.Raw` in `DHashMap.Lemmas` and prove them using the
+  * Restate all of your lemmas for `DHashMap.Raw` in `DHashMap.RawLemmas` and prove them using the
     provided tactic after hooking in your `operation_eq` and `operation_val` from step 5.
   * Restate all of your lemmas for `DHashMap` in `DHashMap.Lemmas` and prove them by reducing to
     `Raw₀`.
-  * Restate all of your lemmas for `HashMap.Raw` in `HashMap.Lemmas` and prove them by reducing to
+  * Restate all of your lemmas for `HashMap.Raw` in `HashMap.RawLemmas` and prove them by reducing to
     `DHashMap.Raw`.
   * Restate all of your lemmas for `HashMap` in `HashMap.Lemmas` and prove them by reducing to
     `DHashMap`.
-  * Restate all of your lemmas for `HashSet.Raw` in `HashSet.Lemmas` and prove them by reducing to
-    `DHashSet.Raw`.
+  * Restate all of your lemmas for `HashSet.Raw` in `HashSet.RawLemmas` and prove them by reducing to
+    `HashMap.Raw`.
   * Restate all of your lemmas for `HashSet` in `HashSet.Lemmas` and prove them by reducing to
-    `DHashSet`.
+    `HashMap`.
 
 This sounds like a lot of work (and it is if you have to add a lot of user-facing lemmas), but the
 framework is set up in such a way that each step is really easy and the proofs are all really short
@@ -420,6 +420,30 @@ variable {β : Type v}
 
 end
 
+/-- Internal implementation detail of the hash map -/
+@[inline] def getKey? [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) : Option α :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let ⟨i, h⟩ := mkIdx buckets.size h (hash a)
+  buckets[i].getKey? a
+
+/-- Internal implementation detail of the hash map -/
+@[inline] def getKey [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (hma : m.contains a) : α :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getKey a hma
+
+/-- Internal implementation detail of the hash map -/
+@[inline] def getKeyD [BEq α] [Hashable α] (m : Raw₀ α β) (a : α) (fallback : α) : α :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getKeyD a fallback
+
+/-- Internal implementation detail of the hash map -/
+@[inline] def getKey! [BEq α] [Hashable α] [Inhabited α] (m : Raw₀ α β) (a : α) : α :=
+  let ⟨⟨_, buckets⟩, h⟩ := m
+  let idx := mkIdx buckets.size h (hash a)
+  buckets[idx.1].getKey! a
+
 end Raw₀
 
 /-- Internal implementation detail of the hash map -/