util/set: Add a _mesa_set_search_or_add() function

Unlike _mesa_set_search_and_add(), it doesn't replace an entry if it's
found, returning it instead. This is useful for nir_instr_set, where
we have to know both the original original instruction and its
equivalent.

Reviewed-by: Eric Anholt <[email protected]>
Acked-by: Jason Ekstrand <[email protected]>

3 files changed, 107 insertions, 23 deletions

diff --git a/src/util/set.c b/src/util/set.c
index c2b97bdea08..599a44a707a 100644
--- a/src/util/set.c
+++ b/src/util/set.c
@@ -247,7 +247,7 @@ _mesa_set_search_pre_hashed(const struct set *set, uint32_t hash,
 }
 
 static struct set_entry *
-set_add(struct set *ht, uint32_t hash, const void *key, bool *replaced);
+set_add(struct set *ht, uint32_t hash, const void *key);
 
 static void
 set_rehash(struct set *ht, unsigned new_size_index)
@@ -274,7 +274,7 @@ set_rehash(struct set *ht, unsigned new_size_index)
    ht->deleted_entries = 0;
 
    set_foreach(&old_ht, entry) {
-      set_add(ht, entry->hash, entry->key, NULL);
+      set_add(ht, entry->hash, entry->key);
    }
 
    ralloc_free(old_ht.table);
@@ -295,13 +295,14 @@ _mesa_set_resize(struct set *set, uint32_t entries)
 }
 
 /**
- * Inserts the key with the given hash into the table.
+ * Find a matching entry for the given key, or insert it if it doesn't already
+ * exist.
  *
  * Note that insertion may rearrange the table on a resize or rehash,
  * so previously found hash_entries are no longer valid after this function.
  */
 static struct set_entry *
-set_add(struct set *ht, uint32_t hash, const void *key, bool *replaced)
+set_search_or_add(struct set *ht, uint32_t hash, const void *key, bool *found)
 {
    uint32_t hash_address;
    struct set_entry *available_entry = NULL;
@@ -325,22 +326,11 @@ set_add(struct set *ht, uint32_t hash, const void *key, bool *replaced)
             break;
       }
 
-      /* Implement replacement when another insert happens
-       * with a matching key.  This is a relatively common
-       * feature of hash tables, with the alternative
-       * generally being "insert the new value as well, and
-       * return it first when the key is searched for".
-       *
-       * Note that the hash table doesn't have a delete callback.
-       * If freeing of old keys is required to avoid memory leaks,
-       * perform a search before inserting.
-       */
       if (!entry_is_deleted(entry) &&
           entry->hash == hash &&
           ht->key_equals_function(key, entry->key)) {
-         entry->key = key;
-         if (replaced)
-            *replaced = true;
+         if (found)
+            *found = true;
          return entry;
       }
 
@@ -350,13 +340,14 @@ set_add(struct set *ht, uint32_t hash, const void *key, bool *replaced)
    } while (hash_address != hash % ht->size);
 
    if (available_entry) {
+      /* There is no matching entry, create it. */
       if (entry_is_deleted(available_entry))
          ht->deleted_entries--;
       available_entry->hash = hash;
       available_entry->key = key;
       ht->entries++;
-      if (replaced)
-         *replaced = false;
+      if (found)
+         *found = false;
       return available_entry;
    }
 
@@ -366,11 +357,38 @@ set_add(struct set *ht, uint32_t hash, const void *key, bool *replaced)
    return NULL;
 }
 
+/**
+ * Inserts the key with the given hash into the table.
+ *
+ * Note that insertion may rearrange the table on a resize or rehash,
+ * so previously found hash_entries are no longer valid after this function.
+ */
+static struct set_entry *
+set_add(struct set *ht, uint32_t hash, const void *key)
+{
+   struct set_entry *entry = set_search_or_add(ht, hash, key, NULL);
+
+   if (unlikely(!entry))
+      return NULL;
+
+   /* Note: If a matching entry already exists, this will replace it.  This is
+    * a relatively common feature of hash tables, with the alternative
+    * generally being "insert the new value as well, and return it first when
+    * the key is searched for".
+    *
+    * Note that the hash table doesn't have a delete callback.  If freeing of
+    * old keys is required to avoid memory leaks, use the alternative
+    * _mesa_set_search_or_add function and implement the replacement yourself.
+    */
+   entry->key = key;
+   return entry;
+}
+
 struct set_entry *
 _mesa_set_add(struct set *set, const void *key)
 {
    assert(set->key_hash_function);
-   return set_add(set, set->key_hash_function(key), key, NULL);
+   return set_add(set, set->key_hash_function(key), key);
 }
 
 struct set_entry *
@@ -378,14 +396,16 @@ _mesa_set_add_pre_hashed(struct set *set, uint32_t hash, const void *key)
 {
    assert(set->key_hash_function == NULL ||
           hash == set->key_hash_function(key));
-   return set_add(set, hash, key, NULL);
+   return set_add(set, hash, key);
 }
 
 struct set_entry *
 _mesa_set_search_and_add(struct set *set, const void *key, bool *replaced)
 {
    assert(set->key_hash_function);
-   return set_add(set, set->key_hash_function(key), key, replaced);
+   return _mesa_set_search_and_add_pre_hashed(set,
+                                              set->key_hash_function(key),
+                                              key, replaced);
 }
 
 struct set_entry *
@@ -394,7 +414,32 @@ _mesa_set_search_and_add_pre_hashed(struct set *set, uint32_t hash,
 {
    assert(set->key_hash_function == NULL ||
           hash == set->key_hash_function(key));
-   return set_add(set, hash, key, replaced);
+   struct set_entry *entry = set_search_or_add(set, hash, key, replaced);
+
+   if (unlikely(!entry))
+      return NULL;
+
+   /* This implements the replacement, same as _mesa_set_add(). The user will
+    * be notified if we're overwriting a found entry.
+    */
+   entry->key = key;
+   return entry;
+}
+
+struct set_entry *
+_mesa_set_search_or_add(struct set *set, const void *key)
+{
+   assert(set->key_hash_function);
+   return set_search_or_add(set, set->key_hash_function(key), key, NULL);
+}
+
+struct set_entry *
+_mesa_set_search_or_add_pre_hashed(struct set *set, uint32_t hash,
+                                   const void *key)
+{
+   assert(set->key_hash_function == NULL ||
+          hash == set->key_hash_function(key));
+   return set_search_or_add(set, hash, key, NULL);
 }
 
 /**
diff --git a/src/util/set.h b/src/util/set.h
index 5742c311a77..783c3d41c46 100644
--- a/src/util/set.h
+++ b/src/util/set.h
@@ -76,6 +76,12 @@ struct set_entry *
 _mesa_set_add_pre_hashed(struct set *set, uint32_t hash, const void *key);
 
 struct set_entry *
+_mesa_set_search_or_add(struct set *set, const void *key);
+struct set_entry *
+_mesa_set_search_or_add_pre_hashed(struct set *set, uint32_t hash,
+                                   const void *key);
+
+struct set_entry *
 _mesa_set_search(const struct set *set, const void *key);
 struct set_entry *
 _mesa_set_search_pre_hashed(const struct set *set, uint32_t hash,
diff --git a/src/util/tests/set/set_test.cpp b/src/util/tests/set/set_test.cpp
index a1eef0b3d98..0b4355af128 100644
--- a/src/util/tests/set/set_test.cpp
+++ b/src/util/tests/set/set_test.cpp
@@ -113,3 +113,36 @@ TEST(set, remove_key)
 
    _mesa_set_destroy(s, NULL);
 }
+
+static uint32_t hash_int(const void *p)
+{
+   int i = *(const int *)p;
+   return i;
+}
+
+static bool cmp_int(const void *p1, const void *p2)
+{
+   int i1 = *(const int *)p1, i2 = *(const int *)p2;
+   return i1 == i2;
+}
+
+TEST(set, search_or_add)
+{
+   struct set *s = _mesa_set_create(NULL, hash_int, cmp_int);
+
+   int a = 10, b = 20, c = 20, d = 30;
+
+   _mesa_set_add(s, &a);
+   _mesa_set_add(s, &b);
+   EXPECT_EQ(s->entries, 2);
+
+   struct set_entry *entry = _mesa_set_search_or_add(s, &c);
+   EXPECT_EQ(entry->key, (void *)&b);
+   EXPECT_EQ(s->entries, 2);
+
+   struct set_entry *entry3 = _mesa_set_search_or_add(s, &d);
+   EXPECT_EQ(entry3->key, &d);
+   EXPECT_EQ(s->entries, 3);
+
+   _mesa_set_destroy(s, NULL);
+}