Refactor btree iterator generation code into a base class rather than using ifdefs inside btree_iterator.

PiperOrigin-RevId: 490317784
Change-Id: I4ffe2a1ad2e39890790e278d82eec7223b67908c

1 files changed, 60 insertions, 49 deletions

diff --git a/absl/container/internal/btree.h b/absl/container/internal/btree.h
index ef4f5dc3..d734676a 100644
--- a/absl/container/internal/btree.h
+++ b/absl/container/internal/btree.h
@@ -1028,8 +1028,50 @@ bool AreNodesFromSameContainer(const Node *node_a, const Node *node_b) {
   return node_a == node_b;
 }
 
+class btree_iterator_generation_info_enabled {
+ public:
+  explicit btree_iterator_generation_info_enabled(uint32_t g)
+      : generation_(g) {}
+
+  // Updates the generation. For use internally right before we return an
+  // iterator to the user.
+  template <typename Node>
+  void update_generation(const Node *node) {
+    if (node != nullptr) generation_ = node->generation();
+  }
+  uint32_t generation() const { return generation_; }
+
+  template <typename Node>
+  void assert_valid_generation(const Node *node) const {
+    if (node != nullptr && node->generation() != generation_) {
+      ABSL_INTERNAL_LOG(
+          FATAL,
+          "Attempting to use an invalidated iterator. The corresponding b-tree "
+          "container has been mutated since this iterator was constructed.");
+    }
+  }
+
+ private:
+  // Used to check that the iterator hasn't been invalidated.
+  uint32_t generation_;
+};
+
+class btree_iterator_generation_info_disabled {
+ public:
+  explicit btree_iterator_generation_info_disabled(uint32_t) {}
+  void update_generation(const void *) {}
+  uint32_t generation() const { return 0; }
+  void assert_valid_generation(const void *) const {}
+};
+
+#ifdef ABSL_BTREE_ENABLE_GENERATIONS
+using btree_iterator_generation_info = btree_iterator_generation_info_enabled;
+#else
+using btree_iterator_generation_info = btree_iterator_generation_info_disabled;
+#endif
+
 template <typename Node, typename Reference, typename Pointer>
-class btree_iterator {
+class btree_iterator : private btree_iterator_generation_info {
   using field_type = typename Node::field_type;
   using key_type = typename Node::key_type;
   using size_type = typename Node::size_type;
@@ -1060,13 +1102,11 @@ class btree_iterator {
 
   btree_iterator() : btree_iterator(nullptr, -1) {}
   explicit btree_iterator(Node *n) : btree_iterator(n, n->start()) {}
-  btree_iterator(Node *n, int p) : node_(n), position_(p) {
-#ifdef ABSL_BTREE_ENABLE_GENERATIONS
-    // Use `~uint32_t{}` as a sentinel value for iterator generations so it
-    // doesn't match the initial value for the actual generation.
-    generation_ = n != nullptr ? n->generation() : ~uint32_t{};
-#endif
-  }
+  btree_iterator(Node *n, int p)
+      : btree_iterator_generation_info(n != nullptr ? n->generation()
+                                                    : ~uint32_t{}),
+        node_(n),
+        position_(p) {}
 
   // NOTE: this SFINAE allows for implicit conversions from iterator to
   // const_iterator, but it specifically avoids hiding the copy constructor so
@@ -1077,11 +1117,9 @@ class btree_iterator {
                     std::is_same<btree_iterator, const_iterator>::value,
                 int> = 0>
   btree_iterator(const btree_iterator<N, R, P> other)  // NOLINT
-      : node_(other.node_), position_(other.position_) {
-#ifdef ABSL_BTREE_ENABLE_GENERATIONS
-    generation_ = other.generation_;
-#endif
-  }
+      : btree_iterator_generation_info(other),
+        node_(other.node_),
+        position_(other.position_) {}
 
   bool operator==(const iterator &other) const {
     return Equals(other);
@@ -1109,7 +1147,7 @@ class btree_iterator {
   // Accessors for the key/value the iterator is pointing at.
   reference operator*() const {
     ABSL_HARDENING_ASSERT(node_ != nullptr);
-    assert_valid_generation();
+    assert_valid_generation(node_);
     ABSL_HARDENING_ASSERT(position_ >= node_->start());
     if (position_ >= node_->finish()) {
       ABSL_HARDENING_ASSERT(!IsEndIterator() && "Dereferencing end() iterator");
@@ -1165,12 +1203,9 @@ class btree_iterator {
                     std::is_same<btree_iterator, iterator>::value,
                 int> = 0>
   explicit btree_iterator(const btree_iterator<N, R, P> other)
-      : node_(const_cast<node_type *>(other.node_)),
-        position_(other.position_) {
-#ifdef ABSL_BTREE_ENABLE_GENERATIONS
-    generation_ = other.generation_;
-#endif
-  }
+      : btree_iterator_generation_info(other.generation()),
+        node_(const_cast<node_type *>(other.node_)),
+        position_(other.position_) {}
 
   bool Equals(const const_iterator other) const {
     ABSL_HARDENING_ASSERT(((node_ == nullptr && other.node_ == nullptr) ||
@@ -1182,8 +1217,8 @@ class btree_iterator {
     // N/M are sizes of the containers containing node_/other.node_.
     assert(AreNodesFromSameContainer(node_, other.node_) &&
            "Comparing iterators from different containers.");
-    assert_valid_generation();
-    other.assert_valid_generation();
+    assert_valid_generation(node_);
+    other.assert_valid_generation(other.node_);
     return node_ == other.node_ && position_ == other.position_;
   }
 
@@ -1204,7 +1239,7 @@ class btree_iterator {
 
   // Increment/decrement the iterator.
   void increment() {
-    assert_valid_generation();
+    assert_valid_generation(node_);
     if (node_->is_leaf() && ++position_ < node_->finish()) {
       return;
     }
@@ -1213,7 +1248,7 @@ class btree_iterator {
   void increment_slow();
 
   void decrement() {
-    assert_valid_generation();
+    assert_valid_generation(node_);
     if (node_->is_leaf() && --position_ >= node_->start()) {
       return;
     }
@@ -1228,29 +1263,9 @@ class btree_iterator {
     return node_->slot(static_cast<size_type>(position_));
   }
 
-// TODO(b/207380122): use an ifdef to select a base class instead of ifdefs
-// inside the class.
-#ifdef ABSL_BTREE_ENABLE_GENERATIONS
-  // Updates the generation. For use internally right before we return an
-  // iterator to the user.
   void update_generation() {
-    if (node_ != nullptr) generation_ = node_->generation();
+    btree_iterator_generation_info::update_generation(node_);
   }
-  uint32_t generation() const { return generation_; }
-
-  void assert_valid_generation() const {
-    if (node_ != nullptr && node_->generation() != generation_) {
-      ABSL_INTERNAL_LOG(
-          FATAL,
-          "Attempting to use an invalidated iterator. The corresponding b-tree "
-          "container has been mutated since this iterator was constructed.");
-    }
-  }
-#else
-  void update_generation() {}
-  uint32_t generation() const { return 0; }
-  void assert_valid_generation() const {}
-#endif
 
   // The node in the tree the iterator is pointing at.
   Node *node_;
@@ -1258,10 +1273,6 @@ class btree_iterator {
   // NOTE: this is an int rather than a field_type because iterators can point
   // to invalid positions (such as -1) in certain circumstances.
   int position_;
-#ifdef ABSL_BTREE_ENABLE_GENERATIONS
-  // Used to check that the iterator hasn't been invalidated.
-  uint32_t generation_;
-#endif
 };
 
 template <typename Params>