2 files changed, 60 insertions, 1 deletions

diff --git a/src/support/small_vector.h b/src/support/small_vector.h
index c798dc594..4f2dc0b81 100644
--- a/src/support/small_vector.h
+++ b/src/support/small_vector.h
@@ -148,9 +148,13 @@ public:
   // iteration
 
   template<typename Parent, typename Iterator> struct IteratorBase {
+    // TODO: Add remaining things from
+    //       https://en.cppreference.com/w/cpp/named_req/RandomAccessIterator
+    using iterator_category = std::random_access_iterator_tag;
     using value_type = T;
     using difference_type = long;
     using reference = T&;
+    using pointer = T*;
 
     Parent* parent;
     size_t index;
@@ -161,7 +165,16 @@ public:
       return index != other.index || parent != other.parent;
     }
 
-    void operator++() { index++; }
+    Iterator& operator++() {
+      Iterator& self = *static_cast<Iterator*>(this);
+      index++;
+      return self;
+    }
+    Iterator operator++(int) {
+      Iterator self = *static_cast<Iterator*>(this);
+      index++;
+      return self;
+    }
 
     Iterator& operator+=(difference_type off) {
       index += off;
@@ -171,6 +184,12 @@ public:
     const Iterator operator+(difference_type off) const {
       return Iterator(*this) += off;
     }
+
+    off_t operator-(const Iterator& other) const { return index - other.index; }
+
+    bool operator==(const Iterator& other) const {
+      return parent == other.parent && index == other.index;
+    }
   };
 
   struct Iterator : IteratorBase<SmallVector<T, N>, Iterator> {
@@ -189,6 +208,12 @@ public:
   Iterator end() { return Iterator(this, size()); }
   ConstIterator begin() const { return ConstIterator(this, 0); }
   ConstIterator end() const { return ConstIterator(this, size()); }
+
+  void erase(Iterator a, Iterator b) {
+    // Atm we only support erasing at the end, which is very efficient.
+    assert(b == end());
+    resize(a.index);
+  }
 };
 
 // A SmallVector for which some values may be read before they are written, and
diff --git a/test/example/small_vector.cpp b/test/example/small_vector.cpp
index a1e394e23..35ce38d16 100644
--- a/test/example/small_vector.cpp
+++ b/test/example/small_vector.cpp
@@ -100,6 +100,39 @@ template<typename T> void test(size_t N) {
     t.clear();
     assert(t.empty());
   }
+  {
+    // Test iteration.
+    T t = {0, 1, 2};
+
+    // Pre-and-postfix ++.
+    auto iter = t.begin();
+    assert(*iter == 0);
+    iter++;
+    assert(*iter == 1);
+    ++iter;
+    assert(*iter == 2);
+
+    // Subtraction.
+    assert(t.begin() - t.begin() == 0);
+    assert(t.end() - t.begin() == 3);
+    iter = t.begin();
+    iter++;
+    assert(iter - t.begin() == 1);
+
+    // Comparison.
+    assert(t.begin() != t.end());
+    assert(iter != t.end());
+    iter++;
+    iter++;
+    assert(iter == t.end());
+
+    // Erasing at the end.
+    iter = t.begin();
+    iter++;
+    t.erase(iter, t.end());
+    assert(t.size() == 1);
+    assert(t[0] == 0);
+  }
 }
 
 int main() {
@@ -108,5 +141,6 @@ int main() {
   test<SmallVector<int, 2>>(2);
   test<SmallVector<int, 3>>(3);
   test<SmallVector<int, 10>>(10);
+
   std::cout << "ok.\n";
 }