From 27134f25b11e3119a2814988c3979fdc033e54e1 Mon Sep 17 00:00:00 2001
From: Abseil Team <absl-team@google.com>
Date: Fri, 12 Jan 2024 12:31:42 -0800
Subject: Speed up `raw_hash_map::[]` with ABSL hardening enabled by unchecking
 dereference of iterator returned by `try_emplace`.

PiperOrigin-RevId: 597920257
Change-Id: I1b2e8f10a2f1efa763a6f0760294beafdb6fd9c0
---
 absl/container/internal/raw_hash_map.h | 11 +++++++++--
 absl/container/internal/raw_hash_set.h |  8 +++++++-
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/absl/container/internal/raw_hash_map.h b/absl/container/internal/raw_hash_map.h
index c1d4d4bc..97182bc7 100644
--- a/absl/container/internal/raw_hash_map.h
+++ b/absl/container/internal/raw_hash_map.h
@@ -177,13 +177,20 @@ class raw_hash_map : public raw_hash_set<Policy, Hash, Eq, Alloc> {
   template <class K = key_type, class P = Policy, K* = nullptr>
   MappedReference<P> operator[](key_arg<K>&& key)
       ABSL_ATTRIBUTE_LIFETIME_BOUND {
-    return Policy::value(&*try_emplace(std::forward<K>(key)).first);
+    // It is safe to use unchecked_deref here because try_emplace
+    // will always return an iterator pointing to a valid item in the table,
+    // since it inserts if nothing is found for the given key.
+    return Policy::value(
+        &this->unchecked_deref(try_emplace(std::forward<K>(key)).first));
   }
 
   template <class K = key_type, class P = Policy>
   MappedReference<P> operator[](const key_arg<K>& key)
       ABSL_ATTRIBUTE_LIFETIME_BOUND {
-    return Policy::value(&*try_emplace(key).first);
+    // It is safe to use unchecked_deref here because try_emplace
+    // will always return an iterator pointing to a valid item in the table,
+    // since it inserts if nothing is found for the given key.
+    return Policy::value(&this->unchecked_deref(try_emplace(key).first));
   }
 
  private:
diff --git a/absl/container/internal/raw_hash_set.h b/absl/container/internal/raw_hash_set.h
index b6dc7548..3518bc34 100644
--- a/absl/container/internal/raw_hash_set.h
+++ b/absl/container/internal/raw_hash_set.h
@@ -1943,7 +1943,7 @@ class raw_hash_set {
     // PRECONDITION: not an end() iterator.
     reference operator*() const {
       AssertIsFull(ctrl_, generation(), generation_ptr(), "operator*()");
-      return PolicyTraits::element(slot_);
+      return unchecked_deref();
     }
 
     // PRECONDITION: not an end() iterator.
@@ -2023,6 +2023,10 @@ class raw_hash_set {
     // Should be used when the lifetimes of the iterators are well-enough
     // understood to prove that they cannot be invalid.
     bool unchecked_equals(const iterator& b) { return ctrl_ == b.control(); }
+
+    // Dereferences the iterator without ABSL Hardening iterator invalidation
+    // checks.
+    reference unchecked_deref() const { return PolicyTraits::element(slot_); }
   };
 
   class const_iterator {
@@ -3159,6 +3163,8 @@ class raw_hash_set {
     return {control() + i, slot_array() + i, common().generation_ptr()};
   }
 
+  reference unchecked_deref(iterator it) { return it.unchecked_deref(); }
+
  private:
   friend struct RawHashSetTestOnlyAccess;
 
-- 
cgit v1.2.3