Add Raw iterators

Author: wkrozowski

src/Std/Data/DTreeMap/Internal/Def.lean

@@ -78,4 +78,11 @@ def toListModel : Impl α β → List ((a : α) × β a)
 @[simp] theorem toListModel_inner {sz k v l r} :
   (.inner sz k v l r : Impl α β).toListModel = l.toListModel ++ ⟨k, v⟩ :: r.toListModel := rfl
 
+/--
+  Computes the size of the tree. Used for verification of iterators.
+-/
+def treeSize : Internal.Impl α β → Nat
+| .leaf => 0
+| .inner _ _ _ l r => 1 + l.treeSize + treeSize r
+
 end Std.DTreeMap.Internal.Impl

src/Std/Data/DTreeMap/Internal/Lemmas.lean

@@ -7945,4 +7945,114 @@ end Const
 
 end map
 
+theorem prune_LE_eq_filter {α β} [Ord α] [TransOrd α] (t : Internal.Impl α β) (ord_t : t.Ordered) (lower_bound : α) :
+    (t.prune_LE ord_t lower_bound).toList = t.toList.filter (fun e => (compare e.fst lower_bound).isGE) := by
+  induction t
+  case leaf =>
+    simp only [prune_LE, toList_eq_toListModel, toListModel_leaf, List.filter_nil]
+  case inner _ k v l r l_ih r_ih =>
+    simp only [prune_LE, toList_eq_toListModel, toListModel_inner, List.filter_append]
+    generalize heq : compare lower_bound k = x
+    cases x
+    case lt =>
+      simp only [toListModel_inner]
+      specialize l_ih (Internal.Impl.Ordered.left ord_t)
+      rw [toList_eq_toListModel] at l_ih
+      simp only [l_ih, toList_eq_toListModel, List.filter, List.append_cancel_left_eq]
+      rw [OrientedOrd.eq_swap, Ordering.swap_eq_lt] at heq
+      simp only [heq, Ordering.isGE_gt, List.cons.injEq, true_and]
+      symm
+      apply List.filter_eq_self.2
+      intro a mem
+      apply Ordering.isGE_of_eq_gt
+      apply TransCmp.gt_trans ?_ heq
+      rw [OrientedOrd.eq_swap, Ordering.swap_eq_gt]
+      exact Internal.Impl.Ordered.compare_right ord_t mem
+    case eq =>
+      simp only [toListModel_inner, toListModel_leaf, List.nil_append, List.filter]
+      rw [OrientedCmp.eq_comm] at heq
+      simp only [heq, Ordering.isGE_eq]
+      suffices new_goal : List.filter (fun e => (compare e.fst lower_bound).isGE) l.toListModel = [] from by
+        simp only [new_goal, List.nil_append, List.cons.injEq, true_and]
+        symm
+        apply List.filter_eq_self.2
+        intro a mem
+        apply Ordering.isGE_of_eq_gt
+        apply TransCmp.gt_of_gt_of_eq ?_ heq
+        rw [OrientedOrd.eq_swap, Ordering.swap_eq_gt]
+        apply Internal.Impl.Ordered.compare_right ord_t mem
+      rw [List.filter_eq_nil_iff]
+      intro a mem
+      simp only [Bool.not_eq_true, Ordering.isGE_eq_false]
+      exact TransCmp.lt_of_lt_of_eq (Internal.Impl.Ordered.compare_left ord_t mem) heq
+    case gt =>
+      simp only [List.filter]
+      rw [OrientedOrd.eq_swap, Ordering.swap_eq_gt] at heq
+      rw [heq]
+      suffices new_goal : List.filter (fun e => (compare e.fst lower_bound).isGE) l.toListModel = [] from by
+        simp only [new_goal, Ordering.isGE_lt, List.nil_append]
+        simp only [toList_eq_toListModel] at r_ih
+        apply r_ih
+      rw [List.filter_eq_nil_iff]
+      intro a mem
+      simp only [Bool.not_eq_true, Ordering.isGE_eq_false]
+      exact TransCmp.lt_trans (Internal.Impl.Ordered.compare_left ord_t mem) heq
+
+theorem prune_LT_eq_filter {α β} [Ord α] [TransOrd α] (t : Internal.Impl α β) (ord_t : t.Ordered) (lower_bound : α) :
+    (t.prune_LT ord_t lower_bound).toList = t.toList.filter (fun e => (compare e.fst lower_bound).isGT) := by
+  induction t
+  case leaf =>
+    simp only [prune_LT, toList_eq_toListModel, toListModel_leaf, List.filter_nil]
+  case inner _ k v l r l_ih r_ih =>
+    simp only [prune_LT, toList_eq_toListModel, toListModel_inner, List.filter_append]
+    generalize heq : compare lower_bound k = x
+    cases x
+    case lt =>
+      simp
+      specialize l_ih (Internal.Impl.Ordered.left ord_t)
+      rw [toList_eq_toListModel] at l_ih
+      simp only [l_ih, toList_eq_toListModel, List.filter, List.append_cancel_left_eq]
+      rw [OrientedOrd.eq_swap, Ordering.swap_eq_lt] at heq
+      simp only [heq, Ordering.isGT_gt, List.cons.injEq, true_and]
+      symm
+      apply List.filter_eq_self.2
+      intro a mem
+      rw [Ordering.isGT_iff_eq_gt]
+      apply TransCmp.gt_trans ?_ heq
+      rw [OrientedOrd.eq_swap, Ordering.swap_eq_gt]
+      exact Internal.Impl.Ordered.compare_right ord_t mem
+    case eq =>
+      simp only [List.filter]
+      rw [OrientedCmp.eq_comm] at heq
+      rw [heq]
+      suffices new_goal : List.filter (fun e => (compare e.fst lower_bound).isGT) l.toListModel = [] ∧
+          List.filter (fun e => (compare e.fst lower_bound).isGT) r.toListModel = r.toListModel from by
+        simp only [new_goal, Ordering.isGT_eq, List.nil_append]
+      apply And.intro
+      . rw [List.filter_eq_nil_iff]
+        intro a mem
+        simp only [Ordering.isGT_iff_eq_gt, ← Ordering.isLE_iff_ne_gt]
+        apply TransOrd.isLE_trans _ (Ordering.isLE_of_eq_eq heq)
+        apply Ordering.isLE_of_eq_lt
+        exact Internal.Impl.Ordered.compare_left ord_t mem
+      . apply List.filter_eq_self.2
+        intro a mem
+        rw [Ordering.isGT_iff_eq_gt]
+        apply TransCmp.gt_of_gt_of_eq ?_ heq
+        rw [OrientedOrd.eq_swap, Ordering.swap_eq_gt]
+        exact Internal.Impl.Ordered.compare_right ord_t mem
+    case gt =>
+      simp only [List.filter]
+      rw [OrientedOrd.eq_swap] at heq
+      rw [Ordering.swap_eq_gt] at heq
+      simp only [heq, Ordering.isGT_lt]
+      specialize r_ih (Ordered.right ord_t)
+      rw [toList_eq_toListModel] at r_ih
+      simp only [r_ih, toList_eq_toListModel, List.self_eq_append_left, List.filter_eq_nil_iff,
+        Ordering.isGT_iff_eq_gt]
+      intro a mem
+      rw [← Ordering.isLE_iff_ne_gt]
+      apply Ordering.isLE_of_eq_lt
+      exact TransCmp.lt_trans (Internal.Impl.Ordered.compare_left ord_t mem) heq
+
 end Std.DTreeMap.Internal.Impl

src/Std/Data/DTreeMap/Internal/Operations.lean

@@ -512,6 +512,38 @@ def insertMany! [Ord α] {ρ : Type w} [ForIn Id ρ ((a : α) × β a)] (t : Imp
     r := ⟨r.val.insert! a b, fun h₀ h₁ => h₁ _ _ _ (r.2 h₀ h₁)⟩
   return r
 
+/-!
+## Operations for verification of iterators and slices
+-/
+
+/--
+  Removes all elements with key less than or equal to `lower_bound`.
+  Does not modify size information stored in the tree.
+-/
+def prune_LE {α β} [Ord α] (t : Internal.Impl α β)
+    (ord_t : t.Ordered) (lower_bound : α) : Internal.Impl α β :=
+  match t with
+  | .leaf => .leaf
+  | .inner sz k v l r =>
+    match compare lower_bound k with
+    | .lt => .inner sz k v (l.prune_LE (Internal.Impl.Ordered.left ord_t) lower_bound) r
+    | .eq => .inner sz k v .leaf r
+    | .gt => r.prune_LE (Internal.Impl.Ordered.right ord_t) lower_bound
+
+/--
+  Removes all elements with key less than to `lower_bound`.
+  Does not modify size information stored in the tree.
+-/
+def prune_LT {α β} [Ord α] (t : Internal.Impl α β)
+    (ord_t : t.Ordered) (lower_bound : α) : Internal.Impl α β :=
+  match t with
+  | .leaf => .leaf
+  | .inner sz k v l r =>
+    match compare lower_bound k with
+    | .lt => .inner sz k v (l.prune_LT (Internal.Impl.Ordered.left ord_t) lower_bound) r
+    | .eq => r
+    | .gt => r.prune_LT (Internal.Impl.Ordered.right ord_t) lower_bound
+
 namespace Const
 
 variable {β : Type v}