In sanitizer mode, detect when invalidated iterators are compared.

PiperOrigin-RevId: 499964205
Change-Id: I45a1d62a5e093921946e7c3c7ab31480252b330e

2 files changed, 25 insertions, 13 deletions

diff --git a/absl/container/internal/raw_hash_set.h b/absl/container/internal/raw_hash_set.h
index 3762820d..0baca2a3 100644
--- a/absl/container/internal/raw_hash_set.h
+++ b/absl/container/internal/raw_hash_set.h
@@ -832,8 +832,8 @@ class HashSetIteratorGenerationInfoEnabled {
   void set_generation_ptr(const GenerationType* ptr) { generation_ptr_ = ptr; }
 
  private:
-  const GenerationType* generation_ptr_ = nullptr;
-  GenerationType generation_ = 0;
+  const GenerationType* generation_ptr_ = EmptyGeneration();
+  GenerationType generation_ = *generation_ptr_;
 };
 
 class HashSetIteratorGenerationInfoDisabled {
@@ -1021,12 +1021,17 @@ size_t SelectBucketCountForIterRange(InputIter first, InputIter last,
   } while (0)
 
 // Note that for comparisons, null/end iterators are valid.
-// TODO(b/254649633): when generations are enabled, detect cases of invalid
-// iterators being compared.
-inline void AssertIsValidForComparison(const ctrl_t* ctrl) {
+inline void AssertIsValidForComparison(const ctrl_t* ctrl,
+                                       GenerationType generation,
+                                       const GenerationType* generation_ptr) {
   ABSL_HARDENING_ASSERT((ctrl == nullptr || IsFull(*ctrl)) &&
                         "Invalid iterator comparison. The element might have "
                         "been erased or the table might have rehashed.");
+  if (SwisstableGenerationsEnabled() && generation != *generation_ptr) {
+    ABSL_INTERNAL_LOG(FATAL,
+                      "Invalid iterator comparison. The table could have "
+                      "rehashed since this iterator was initialized.");
+  }
 }
 
 // If the two iterators come from the same container, then their pointers will
@@ -1426,8 +1431,8 @@ class raw_hash_set {
 
     friend bool operator==(const iterator& a, const iterator& b) {
       AssertSameContainer(a.ctrl_, b.ctrl_, a.slot_, b.slot_);
-      AssertIsValidForComparison(a.ctrl_);
-      AssertIsValidForComparison(b.ctrl_);
+      AssertIsValidForComparison(a.ctrl_, a.generation(), a.generation_ptr());
+      AssertIsValidForComparison(b.ctrl_, b.generation(), b.generation_ptr());
       return a.ctrl_ == b.ctrl_;
     }
     friend bool operator!=(const iterator& a, const iterator& b) {
@@ -1444,6 +1449,9 @@ class raw_hash_set {
       // not equal to any end iterator.
       ABSL_ASSUME(ctrl != nullptr);
     }
+    // For end() iterators.
+    explicit iterator(const GenerationType* generation_ptr)
+        : HashSetIteratorGenerationInfo(generation_ptr) {}
 
     // Fixes up `ctrl_` to point to a full by advancing it and `slot_` until
     // they reach one.
@@ -1700,12 +1708,12 @@ class raw_hash_set {
     it.skip_empty_or_deleted();
     return it;
   }
-  iterator end() { return {}; }
+  iterator end() { return iterator(common().generation_ptr()); }
 
   const_iterator begin() const {
     return const_cast<raw_hash_set*>(this)->begin();
   }
-  const_iterator end() const { return {}; }
+  const_iterator end() const { return iterator(common().generation_ptr()); }
   const_iterator cbegin() const { return begin(); }
   const_iterator cend() const { return end(); }
 
diff --git a/absl/container/internal/raw_hash_set_test.cc b/absl/container/internal/raw_hash_set_test.cc
index 8096567c..801b758e 100644
--- a/absl/container/internal/raw_hash_set_test.cc
+++ b/absl/container/internal/raw_hash_set_test.cc
@@ -1819,7 +1819,6 @@ TEST(Table, HeterogeneousLookupOverloads) {
   EXPECT_TRUE((VerifyResultOf<CallCount, TransparentTable>()));
 }
 
-// TODO(alkis): Expand iterator tests.
 TEST(Iterator, IsDefaultConstructible) {
   StringTable::iterator i;
   EXPECT_TRUE(i == StringTable::iterator());
@@ -2249,7 +2248,8 @@ TEST(Table, AlignOne) {
 // Invalid iterator use can trigger heap-use-after-free in asan,
 // use-of-uninitialized-value in msan, or invalidated iterator assertions.
 constexpr const char* kInvalidIteratorDeathMessage =
-    "heap-use-after-free|use-of-uninitialized-value|invalidated iterator";
+    "heap-use-after-free|use-of-uninitialized-value|invalidated "
+    "iterator|Invalid iterator";
 
 #if defined(__clang__) && defined(_MSC_VER)
 constexpr bool kLexan = true;
@@ -2257,7 +2257,7 @@ constexpr bool kLexan = true;
 constexpr bool kLexan = false;
 #endif
 
-TEST(Table, InvalidIteratorUse) {
+TEST(Iterator, InvalidUseCrashesWithSanitizers) {
   if (!SwisstableGenerationsEnabled()) GTEST_SKIP() << "Generations disabled.";
   if (kLexan) GTEST_SKIP() << "Lexan doesn't support | in regexp.";
 
@@ -2268,10 +2268,12 @@ TEST(Table, InvalidIteratorUse) {
     auto it = t.begin();
     t.insert(i);
     EXPECT_DEATH_IF_SUPPORTED(*it, kInvalidIteratorDeathMessage);
+    EXPECT_DEATH_IF_SUPPORTED(void(it == t.begin()),
+                              kInvalidIteratorDeathMessage);
   }
 }
 
-TEST(Table, InvalidIteratorUseWithReserve) {
+TEST(Iterator, InvalidUseWithReserveCrashesWithSanitizers) {
   if (!SwisstableGenerationsEnabled()) GTEST_SKIP() << "Generations disabled.";
   if (kLexan) GTEST_SKIP() << "Lexan doesn't support | in regexp.";
 
@@ -2290,6 +2292,8 @@ TEST(Table, InvalidIteratorUseWithReserve) {
   // Unreserved growth can rehash.
   t.insert(10);
   EXPECT_DEATH_IF_SUPPORTED(*it, kInvalidIteratorDeathMessage);
+  EXPECT_DEATH_IF_SUPPORTED(void(it == t.begin()),
+                            kInvalidIteratorDeathMessage);
 }
 
 TEST(Table, ReservedGrowthUpdatesWhenTableDoesntGrow) {