From 9d95c6836dfe5c4d02728343a0c89508a2476c33 Mon Sep 17 00:00:00 2001
From: Chandler Carruth <chandlerc@gmail.com>
Date: Sat, 22 Jun 2024 02:57:41 +0200
Subject: [PATCH] Add a key-returning callback insert to `Set`. (#4072)

This turns out to be super useful now that we have the `key_context`
mechanism and can do much more meaningful heterogeneous lookups, where
the stored key can be *very* different from the lookup key.

---------

Co-authored-by: Richard Smith <richard@metafoo.co.uk>
---
 common/set.h        | 32 ++++++++++++++++++++++++++++++++
 common/set_test.cpp | 11 +++++++++--
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/common/set.h b/common/set.h
index 9ce86a51c7b7..6116dae8e376 100644
--- a/common/set.h
+++ b/common/set.h
@@ -204,6 +204,21 @@ class SetBase
   auto Insert(LookupKeyT lookup_key, KeyContextT key_context = KeyContextT())
       -> InsertResult;
 
+  // Insert a key into the map and call the provided callback if necessary to
+  // produce a new key when no existing value is found.
+  //
+  // Example: `m.Insert(key_equivalent, [] { return real_key; });`
+  //
+  // The point of this function is when the lookup key is _different_from the
+  // stored key. However, we don't restrict it in case that blocks generic
+  // usage.
+  template <typename LookupKeyT, typename KeyCallbackT>
+  auto Insert(LookupKeyT lookup_key, KeyCallbackT key_cb,
+              KeyContextT key_context = KeyContextT()) -> InsertResult
+    requires(
+        !std::same_as<KeyT, KeyCallbackT> &&
+        std::convertible_to<decltype(std::declval<KeyCallbackT>()()), KeyT>);
+
   // Insert a key into the set and call the provided callback to allow in-place
   // construction of the key if not already present. The lookup key is passed
   // through to the callback so it needn't be captured and can be kept in a
@@ -331,6 +346,23 @@ auto SetBase<InputKeyT, InputKeyContextT>::Insert(LookupKeyT lookup_key,
       key_context);
 }
 
+template <typename InputKeyT, typename InputKeyContextT>
+template <typename LookupKeyT, typename KeyCallbackT>
+auto SetBase<InputKeyT, InputKeyContextT>::Insert(LookupKeyT lookup_key,
+                                                  KeyCallbackT key_cb,
+                                                  KeyContextT key_context)
+    -> InsertResult
+  requires(!std::same_as<KeyT, KeyCallbackT> &&
+           std::convertible_to<decltype(std::declval<KeyCallbackT>()()), KeyT>)
+{
+  return Insert(
+      lookup_key,
+      [&key_cb](LookupKeyT /*lookup_key*/, void* key_storage) {
+        new (key_storage) KeyT(key_cb());
+      },
+      key_context);
+}
+
 template <typename InputKeyT, typename InputKeyContextT>
 template <typename LookupKeyT, typename InsertCallbackT>
 auto SetBase<InputKeyT, InputKeyContextT>::Insert(LookupKeyT lookup_key,
diff --git a/common/set_test.cpp b/common/set_test.cpp
index 2032d78adbb5..1a215898e680 100644
--- a/common/set_test.cpp
+++ b/common/set_test.cpp
@@ -343,12 +343,17 @@ TEST(SetContextTest, Basic) {
   auto i_result = s.Insert(1, IndexKeyContext<TestData>(keys));
   EXPECT_FALSE(i_result.is_inserted());
   EXPECT_TRUE(s.Contains(1, key_context));
+  EXPECT_TRUE(s.Insert(
+                   TestData(200), [] { return 2; }, key_context)
+                  .is_inserted());
+  EXPECT_TRUE(s.Contains(2, key_context));
+  EXPECT_TRUE(s.Contains(TestData(200), key_context));
 
   // Verify all the elements.
-  ExpectSetElementsAre(s, {1});
+  ExpectSetElementsAre(s, {1, 2});
 
   // Fill up a bunch to ensure we trigger growth a few times.
-  for (int i : llvm::seq(2, 512)) {
+  for (int i : llvm::seq(3, 512)) {
     SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
     EXPECT_TRUE(s.Insert(i, key_context).is_inserted());
   }
@@ -359,6 +364,8 @@ TEST(SetContextTest, Basic) {
   }
   EXPECT_FALSE(s.Contains(0, key_context));
   EXPECT_FALSE(s.Contains(512, key_context));
+  EXPECT_FALSE(s.Contains(TestData(0), key_context));
+  EXPECT_FALSE(s.Contains(TestData(51200), key_context));
 
   // Verify all the elements.
   ExpectSetElementsAre(s, MakeElements(llvm::seq(1, 512)));