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 <[email protected]>

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,

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)));