Add erase(iterator) to the Highway B-tree and document iterator invalidation

hwy/contrib/btree/btree-inl.h

@@ -1442,6 +1442,14 @@ void MergeLeaves(MapLeafNode<KeyT, ValueT>* leaf,
 // 2. Adapter Mode (Non-owning): Constructed with an external BTreeState pointer
 //    (e.g., from hwy::BTreeSet). Operates directly on the external state
 //    without copying data, and leaves node deallocation to the external owner.
+//
+// Iterator invalidation: as with absl::btree_*, any mutation (insert, erase,
+// insert_or_assign, operator[], clear, Build, swap, assignment) may invalidate
+// ALL outstanding iterators, pointers and references, including end(). Keys
+// are stored delta-compressed inside 512-byte leaves, so an insert can shift,
+// recompress or split a leaf, and an erase can merge or free one. To continue
+// iterating after a mutation, use the iterator returned by insert() or
+// erase(iterator); do not reuse iterators obtained before the mutation.
 template <typename Traits>
 class BTree {
  public:
@@ -1870,7 +1878,8 @@ class BTree {
   // ---------------------------------------------------------------------------
 
   // Constructs a BTreeSet from an array of pre-sorted, unique keys in
-  // O(N) time.
+  // O(N) time. Assigning the result to an existing tree invalidates all of
+  // that tree's iterators, pointers and references.
   //
   // Example usage:
   //   std::vector<uint32_t> sorted_keys = {10, 20, 30, 40, 50};
@@ -2203,6 +2212,9 @@ class BTree {
   // ---------------------------------------------------------------------------
   // Dynamic Mutations (Insertions & Deletions)
   // ---------------------------------------------------------------------------
+  // If an insertion or erasure occurs, every outstanding iterator, pointer and
+  // reference (including end()) may be invalidated. The returned iterator (if
+  // any) is valid until the next mutation.
 
   // Inserts a key into the Set. Returns pair of (iterator, bool_inserted).
   template <bool IsMap = Traits::kIsMap, typename = std::enable_if_t<!IsMap>>
@@ -2302,10 +2314,30 @@ class BTree {
     return EraseInternal(KeyCodec<KeyT>::ToStorage(key));
   }
 
+  // Erases the element at `pos`, which must be dereferenceable (not end()).
+  // Returns an iterator to the element following the erased one, or end().
+  // Like absl::btree, this is the only safe way to erase while iterating:
+  //   for (auto it = t.begin(); it != t.end();) {
+  //     if (ShouldErase(*it)) it = t.erase(it); else ++it;
+  //   }
+  // Cost is one additional root-to-leaf descent to re-locate the successor,
+  // because erasing may merge or free the leaf that `pos` pointed into.
+  iterator erase(const_iterator pos) {
+    HWY_DASSERT(pos.leaf() != nullptr && pos.slot() < pos.leaf()->NumKeys());
+    const StorageKeyT key = GetLeafKey(pos.leaf(), pos.slot());
+    const size_t erased = EraseInternal(key);
+    HWY_DASSERT(erased == 1);
+    (void)erased;
+    // The key is gone, so lower_bound yields the first element greater than it.
+    return LowerBoundInternal(key);
+  }
+
   // ---------------------------------------------------------------------------
   // Capacity & Iteration
   // ---------------------------------------------------------------------------
 
+  // Removes all elements and frees all nodes. Invalidates all iterators,
+  // pointers and references.
   void clear() {
     DestroySubtree(state_->root_, state_->tree_height_);
     state_->root_ = nullptr;

hwy/contrib/btree/btree_map.h

@@ -60,6 +60,8 @@ struct MapDispatch;
     static std::pair<std::pair<Leaf*, size_t>, bool> Insert(                  \
         State* state, KeyT key, StorageValueT value, bool assign_if_exists);  \
     static size_t Erase(State* state, KeyT key);                              \
+    static std::pair<Leaf*, size_t> EraseIter(State* state, const Leaf* leaf, \
+                                              size_t slot);                   \
     static void ContainsBatch(const State* state, const KeyT* keys,           \
                               size_t count, bool* out);                       \
     static void FindBatch(const State* state, const KeyT* keys, size_t count, \
@@ -98,6 +100,20 @@ struct ValueStorageTraits {
   using Type = std::conditional_t<sizeof(ValueT) == 4, uint32_t, uint64_t>;
 };
 
+// SIMD-accelerated ordered map from 32/64-bit integer keys to trivially
+// copyable 32/64-bit values, with an API modelled on std::map /
+// absl::btree_map. Keys are stored delta-compressed in 512-byte leaves; SIMD
+// kernels are selected at runtime via dynamic dispatch.
+//
+// Iterator invalidation: as with absl::btree_map, if an insertion or erasure
+// occurs (insert, insert_or_assign, emplace, operator[], erase, clear, Build,
+// swap, assignment), ALL outstanding iterators, pointers and references may be
+// invalidated, including end() and ValueT& obtained from operator[]/at/
+// FindValue. To continue iterating after a mutation, use the iterator returned
+// by insert() or erase(iterator):
+//   for (auto it = m.begin(); it != m.end();) {
+//     if (ShouldErase(it->second)) it = m.erase(it); else ++it;
+//   }
 template <typename KeyT, typename ValueT>
 class BTreeMap {
  public:
@@ -146,9 +162,11 @@ class BTreeMap {
     return *this;
   }
 
-  // Bulk-builds a tree from sorted keys and values. The contiguous values array
-  // is passed as const void* across dynamic dispatch to avoid strict-aliasing
-  // issues across different types sharing the same 32-bit or 64-bit size.
+  // Bulk-builds a tree from strictly ascending keys and their values. The
+  // contiguous values array is passed as const void* across dynamic dispatch
+  // to avoid strict-aliasing issues across different types sharing the same
+  // 32-bit or 64-bit size. Assigning the result to an existing map invalidates
+  // all of its iterators.
   static BTreeMap Build(const KeyT* sorted_keys, const ValueT* sorted_values,
                         size_t num_keys, float fill_ratio = 1.0f) {
     BTreeMap map;
@@ -157,6 +175,7 @@ class BTreeMap {
     return map;
   }
 
+  // Removes all elements. Invalidates all iterators, pointers and references.
   void clear() { Dispatch::Clear(&state_); }
 
   // ---------------------------------------------------------------------------
@@ -556,6 +575,8 @@ class BTreeMap {
 
   // Inserts key and value. Converts value to StorageValueT via BitCastScalar,
   // preventing strict aliasing issues while compiling down to a register move.
+  // Returns (iterator to the element, whether it was inserted). If an
+  // insertion occurs, all other iterators are invalidated.
   std::pair<iterator, bool> insert(KeyT key, const ValueT& value) {
     StorageValueT s_val = hwy::BitCastScalar<StorageValueT>(value);
     auto res =
@@ -576,7 +597,15 @@ class BTreeMap {
     return {iterator(res.first.first, res.first.second, state_.last_leaf_),
             res.second};
   }
+  // Erases `key` if present. Returns the number of elements erased (0 or 1).
+  // If an erasure occurs, all iterators are invalidated.
   size_t erase(KeyT key) { return Dispatch::Erase(&state_, key); }
+  // Erases the element at `pos` (must not be end()). Returns an iterator to the
+  // following element, or end(). All other iterators are invalidated.
+  iterator erase(const_iterator pos) {
+    auto res = Dispatch::EraseIter(&state_, pos.leaf(), pos.slot());
+    return iterator(res.first, res.second, state_.last_leaf_);
+  }
 
   const LeafT* last_leaf() const { return state_.last_leaf_; }
   LeafT* last_leaf() { return state_.last_leaf_; }

hwy/contrib/btree/btree_set.h

@@ -25,6 +25,18 @@
 
 namespace hwy {
 
+// SIMD-accelerated ordered set of 32/64-bit integer keys with an API modelled
+// on std::set / absl::btree_set. Keys are stored delta-compressed in 512-byte
+// leaves; SIMD kernels are selected at runtime via dynamic dispatch.
+//
+// Iterator invalidation: as with absl::btree_set, if an insertion or erasure
+// occurs (insert, emplace, erase, clear, Build, swap, assignment), ALL
+// outstanding iterators, pointers and references may be invalidated,
+// including end(). To continue iterating after a mutation, use the iterator
+// returned by insert() or erase(iterator):
+//   for (auto it = s.begin(); it != s.end();) {
+//     if (ShouldErase(*it)) it = s.erase(it); else ++it;
+//   }
 template <typename KeyT>
 class BTreeSet {
  public:
@@ -60,9 +72,12 @@ class BTreeSet {
     return *this;
   }
 
+  // Bulk-builds a set from `num_keys` strictly ascending keys in O(N).
+  // Assigning the result to an existing set invalidates all of its iterators.
   static BTreeSet Build(const KeyT* sorted_keys, size_t num_keys,
                         float fill_ratio = 1.0f);
 
+  // Removes all elements. Invalidates all iterators, pointers and references.
   void clear();
 
   // ---------------------------------------------------------------------------
@@ -380,12 +395,19 @@ class BTreeSet {
                     state_.last_leaf_);
   }
 
+  // Inserts `key`. Returns (iterator to the element, whether it was inserted).
+  // If an insertion occurs, all other iterators are invalidated.
   std::pair<iterator, bool> insert(KeyT key);
   template <typename... Args>
   std::pair<iterator, bool> emplace(Args&&... args) {
     return insert(KeyT(std::forward<Args>(args)...));
   }
+  // Erases `key` if present. Returns the number of elements erased (0 or 1).
+  // If an erasure occurs, all iterators are invalidated.
   size_t erase(KeyT key);
+  // Erases the element at `pos` (must not be end()). Returns an iterator to the
+  // following element, or end(). All other iterators are invalidated.
+  iterator erase(const_iterator pos);
 
   const LeafT* last_leaf() const { return state_.last_leaf_; }
   LeafT* last_leaf() { return state_.last_leaf_; }
@@ -436,6 +458,9 @@ class BTreeSet {
   template <>                                                                \
   HWY_CONTRIB_DLLEXPORT size_t BTreeSet<KeyT>::erase(KeyT key);              \
   template <>                                                                \
+  HWY_CONTRIB_DLLEXPORT BTreeSet<KeyT>::iterator BTreeSet<KeyT>::erase(      \
+      const_iterator pos);                                                   \
+  template <>                                                                \
   HWY_CONTRIB_DLLEXPORT bool BTreeSet<KeyT>::Contains(KeyT key) const;       \
   template <>                                                                \
   HWY_CONTRIB_DLLEXPORT void BTreeSet<KeyT>::ContainsBatch(                  \