[libc] Add erase function to blockstore

[michaelrj-google]

Reland of #97641 with sanitizer fixes

This adds the ability to erase a value from a blockstore based on an
iterator. For usability/testing purposes it also includes an addition
operator for blockstore's iterator.

[llvmbot]

@llvm/pr-subscribers-libc

Author: Michael Jones (michaelrj-google)

Changes

Reland of #97641 with sanitizer fixes

This adds the ability to erase a value from a blockstore based on an
iterator. For usability/testing purposes it also includes an addition
operator for blockstore's iterator.

---

Full diff: https://github.com/llvm/llvm-project/pull/98674.diff

2 Files Affected:

• (modified) libc/src/__support/blockstore.h (+55-2)
• (modified) libc/test/src/__support/blockstore_test.cpp (+98)

diff --git a/libc/src/__support/blockstore.h b/libc/src/__support/blockstore.h
index 8d13e0ed290df..efe2234eace59 100644
--- a/libc/src/__support/blockstore.h
+++ b/libc/src/__support/blockstore.h
@@ -9,9 +9,11 @@
 #ifndef LLVM_LIBC_SRC___SUPPORT_BLOCKSTORE_H
 #define LLVM_LIBC_SRC___SUPPORT_BLOCKSTORE_H
 
+#include "src/__support/CPP/array.h"
+#include "src/__support/CPP/new.h"
+#include "src/__support/CPP/type_traits.h"
+#include "src/__support/libc_assert.h"
 #include "src/__support/macros/config.h"
-#include <src/__support/CPP/new.h>
-#include <src/__support/libc_assert.h>
 
 #include <stddef.h>
 #include <stdint.h>
@@ -98,6 +100,16 @@ class BlockStore {
       return *reinterpret_cast<T *>(block->data + sizeof(T) * true_index);
     }
 
+    LIBC_INLINE Iterator operator+(int i) {
+      LIBC_ASSERT(i >= 0 &&
+                  "BlockStore iterators only support incrementation.");
+      auto other = *this;
+      for (int j = 0; j < i; ++j)
+        ++other;
+
+      return other;
+    }
+
     LIBC_INLINE bool operator==(const Iterator &rhs) const {
       return block == rhs.block && index == rhs.index;
     }
@@ -176,6 +188,47 @@ class BlockStore {
     else
       return Iterator(current, fill_count);
   }
+
+  // Removes the element at pos, then moves all the objects after back by one to
+  // fill the hole. It's assumed that pos is a valid iterator to somewhere in
+  // this block_store.
+  LIBC_INLINE void erase(Iterator pos) {
+    const Iterator last_item = Iterator(current, fill_count);
+    if (pos == last_item) {
+      pop_back();
+      return;
+    }
+
+    if constexpr (REVERSE_ORDER) {
+      // REVERSE: Iterate from begin to pos
+      const Iterator range_end = pos;
+      Iterator cur = begin();
+      T prev_val = *cur;
+      ++cur;
+      T cur_val = *cur;
+
+      for (; cur != range_end; ++cur) {
+        cur_val = *cur;
+        *cur = prev_val;
+        prev_val = cur_val;
+      }
+      // As long as this isn't the end we will always need to move at least one
+      // item (since we know that pos isn't the last item due to the check
+      // above).
+      if (range_end != end())
+        *cur = prev_val;
+    } else {
+      // FORWARD: Iterate from pos to end
+      const Iterator range_end = end();
+      Iterator cur = pos;
+      Iterator prev = cur;
+      ++cur;
+
+      for (; cur != range_end; prev = cur, ++cur)
+        *prev = *cur;
+    }
+    pop_back();
+  }
 };
 
 template <typename T, size_t BLOCK_SIZE, bool REVERSE_ORDER>
diff --git a/libc/test/src/__support/blockstore_test.cpp b/libc/test/src/__support/blockstore_test.cpp
index 5fe8fef1b6edc..dd74ea18f2c02 100644
--- a/libc/test/src/__support/blockstore_test.cpp
+++ b/libc/test/src/__support/blockstore_test.cpp
@@ -64,6 +64,99 @@ class LlvmLibcBlockStoreTest : public LIBC_NAMESPACE::testing::Test {
     }
     block_store.destroy(&block_store);
   }
+
+  template <bool REVERSE> void erase_test() {
+    using LIBC_NAMESPACE::BlockStore;
+    BlockStore<int, 2, REVERSE> block_store;
+    int i;
+
+    constexpr int ARR_SIZE = 6;
+
+    ASSERT_TRUE(block_store.empty());
+    for (int i = 0; i < ARR_SIZE; i++) {
+      ASSERT_TRUE(block_store.push_back(i + 1));
+    }
+
+    // block_store state should be {1,2,3,4,5,6}
+
+    block_store.erase(block_store.begin());
+
+    // FORWARD: block_store state should be {2,3,4,5,6}
+    // REVERSE: block_store state should be {1,2,3,4,5}
+
+    auto iter = block_store.begin();
+    for (i = 0; iter != block_store.end(); ++i, ++iter) {
+      if (!REVERSE) {
+        ASSERT_EQ(*iter, i + 2);
+      } else {
+        ASSERT_EQ(*iter, (ARR_SIZE - 1) - i);
+      }
+    }
+
+    // Assert that there were the correct number of elements
+    ASSERT_EQ(i, ARR_SIZE - 1);
+
+    block_store.erase(block_store.end());
+
+    // BOTH: block_store state should be {2,3,4,5}
+
+    iter = block_store.begin();
+    for (i = 0; iter != block_store.end(); ++i, ++iter) {
+      if (!REVERSE) {
+        ASSERT_EQ(*iter, i + 2);
+      } else {
+        ASSERT_EQ(*iter, (ARR_SIZE - 1) - i);
+      }
+    }
+
+    ASSERT_EQ(i, ARR_SIZE - 2);
+
+    block_store.erase(block_store.begin() + 1);
+
+    // FORWARD: block_store state should be {2,4,5}
+    // REVERSE: block_store state should be {2,3,5}
+
+    const int FORWARD_RESULTS[] = {2, 4, 5};
+    const int REVERSE_RESULTS[] = {2, 3, 5};
+
+    iter = block_store.begin();
+    for (i = 0; iter != block_store.end(); ++i, ++iter) {
+      if (!REVERSE) {
+        ASSERT_EQ(*iter, FORWARD_RESULTS[i]);
+      } else {
+        ASSERT_EQ(*iter, REVERSE_RESULTS[ARR_SIZE - 4 - i]); // reversed
+      }
+    }
+
+    ASSERT_EQ(i, ARR_SIZE - 3);
+
+    block_store.erase(block_store.begin() + 1);
+    // BOTH: block_store state should be {2,5}
+
+    iter = block_store.begin();
+    if (!REVERSE) {
+      ASSERT_EQ(*iter, 2);
+      ASSERT_EQ(*(iter + 1), 5);
+    } else {
+      ASSERT_EQ(*iter, 5);
+      ASSERT_EQ(*(iter + 1), 2);
+    }
+
+    block_store.erase(block_store.begin());
+    // FORWARD: block_store state should be {5}
+    // REVERSE: block_store state should be {2}
+    iter = block_store.begin();
+    if (!REVERSE) {
+      ASSERT_EQ(*iter, 5);
+    } else {
+      ASSERT_EQ(*iter, 2);
+    }
+
+    block_store.erase(block_store.begin());
+    // BOTH: block_store state should be {}
+
+    block_store.destroy(&block_store);
+  }
 };
 
 TEST_F(LlvmLibcBlockStoreTest, PopulateAndIterate4) {
@@ -100,3 +193,8 @@ TEST_F(LlvmLibcBlockStoreTest, Empty) {
   empty_test<false>();
   empty_test<true>();
 }
+
+TEST_F(LlvmLibcBlockStoreTest, Erase) {
+  erase_test<false>();
+  erase_test<true>();
+}