Allow removing current iter from ListRange without invalidating

LoopModels · Jul 20, 2023 · 38a5da2 · 38a5da2
1 parent 6064800
commit 38a5da2
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/include/Utilities/ListRanges.hpp b/include/Utilities/ListRanges.hpp
@@ -43,8 +43,10 @@ class End {
   }
 };
 
+/// Safe for removing current iter from `list` while iterating.
 template <typename T, class Op, class Proj> class ListIterator {
   T *state_;
+  T *next_{nullptr};
   [[no_unique_address]] Op op_{};
   [[no_unique_address]] Proj p_{};
 
@@ -54,12 +56,14 @@ template <typename T, class Op, class Proj> class ListIterator {
   // constexpr auto operator->() const noexcept -> T * { return state_; }
   constexpr auto getState() const noexcept -> T * { return state_; }
   constexpr auto operator++() noexcept -> ListIterator & {
-    state_ = op_(state_);
+    state_ = next_;
+    if (next_) next_ = op_(next_);
     return *this;
   }
   constexpr auto operator++(int) noexcept -> ListIterator {
     ListIterator tmp{*this};
-    state_ = op_(state_);
+    state_ = next_;
+    if (next_) next_ = op_(next_);
     return tmp;
   }
   constexpr auto operator-(ListIterator const &other) const noexcept
@@ -74,9 +78,10 @@ template <typename T, class Op, class Proj> class ListIterator {
   constexpr auto operator==(End) const noexcept -> bool {
     return state_ == nullptr;
   }
-  constexpr ListIterator(T *state) noexcept : state_{state} {}
+  constexpr ListIterator(T *state) noexcept
+    : state_{state}, next_{state ? Op{}(state) : nullptr} {}
   constexpr ListIterator(T *state, Op op, Proj p) noexcept
-    : state_{state}, op_{op}, p_{p} {}
+    : state_{state}, next_{state ? op(state) : nullptr}, op_{op}, p_{p} {}
   constexpr ListIterator() noexcept = default;
   constexpr ListIterator(const ListIterator &) noexcept = default;
   constexpr ListIterator(ListIterator &&) noexcept = default;