From 969115bc29cb7c5e3caa4a05fa9bd92716cc3fcd Mon Sep 17 00:00:00 2001
From: Stiopa Koltsov <stepan.koltsov@gmail.com>
Date: Sun, 22 Aug 2021 22:14:51 +0100
Subject: [PATCH] IndexMap::insert_unique_unchecked

Insert a key-value pair into the map without checking if the key
already exists in the map.

This operation is safe if a key does not exist in the map.

However, if a key exists in the map already, the behavior is
unspecified: this operation may panic, or any following operation
with the map may panic or return arbitrary result.

This operation is faster than regular insert, because it does not
perform lookup before insertion.

This operation is useful during initial population of the map.  For
example, when constructing a map from another map, we know that
keys are unique.

Simple benchmark of `insert` vs `insert_unique_unchecked` included:

```
test insert                  ... bench:      14,929 ns/iter (+/- 2,222)
test insert_unique_unchecked ... bench:      11,272 ns/iter (+/- 1,172)
```
---
 benches/insert_unique_unchecked.rs | 29 ++++++++++++++++++++++++
 src/map.rs                         | 36 ++++++++++++++++++++++++++++++
 src/map/core.rs                    |  2 +-
 src/set.rs                         | 21 +++++++++++++++++
 4 files changed, 87 insertions(+), 1 deletion(-)
 create mode 100644 benches/insert_unique_unchecked.rs

diff --git a/benches/insert_unique_unchecked.rs b/benches/insert_unique_unchecked.rs
new file mode 100644
index 00000000..2bc18d79
--- /dev/null
+++ b/benches/insert_unique_unchecked.rs
@@ -0,0 +1,29 @@
+#![feature(test)]
+
+extern crate test;
+
+use test::Bencher;
+
+use indexmap::IndexMap;
+
+#[bench]
+fn insert(b: &mut Bencher) {
+    b.iter(|| {
+        let mut m = IndexMap::with_capacity(1000);
+        for i in 0..1000 {
+            m.insert(i, i);
+        }
+        m
+    });
+}
+
+#[bench]
+fn insert_unique_unchecked(b: &mut Bencher) {
+    b.iter(|| {
+        let mut m = IndexMap::with_capacity(1000);
+        for i in 0..1000 {
+            m.insert_unique_unchecked(i, i);
+        }
+        m
+    });
+}
diff --git a/src/map.rs b/src/map.rs
index 3a43a79b..637a9bb0 100644
--- a/src/map.rs
+++ b/src/map.rs
@@ -373,6 +373,30 @@ where
         self.core.insert_full(hash, key, value)
     }
 
+    /// Insert a key-value pair into the map without checking
+    /// if the key already exists in the map.
+    ///
+    /// Returns a reference to the key and value just inserted.
+    ///
+    /// This operation is memory-safe.
+    ///
+    /// However, **if a key exists in the map already, the behavior is unspecified**:
+    /// this operation may panic, loop forever, or any following operation with the map
+    /// may panic, loop forever or return arbitrary result. But it won't violate memory safety.
+    ///
+    /// This operation is faster than regular insert because it does not perform
+    /// lookup before insertion.
+    ///
+    /// This operation is useful during initial population of the map.
+    /// For example, when constructing a map from another map, we know
+    /// that keys are unique.
+    pub fn insert_unique_unchecked(&mut self, key: K, value: V) -> (&K, &mut V) {
+        let hash = self.hash(&key);
+        let index = self.core.push(hash, key, value);
+        let entry = &mut self.core.as_entries_mut()[index];
+        (&entry.key, &mut entry.value)
+    }
+
     /// Get the given key’s corresponding entry in the map for insertion and/or
     /// in-place manipulation.
     ///
@@ -1834,4 +1858,16 @@ mod tests {
         assert!(values.contains(&'b'));
         assert!(values.contains(&'c'));
     }
+
+    #[test]
+    fn insert_unique_unchecked() {
+        let mut map = IndexMap::new();
+        let (k1, v1) = map.insert_unique_unchecked(10, 11);
+        assert_eq!((&10, &mut 11), (k1, v1));
+        let (k2, v2) = map.insert_unique_unchecked(20, 21);
+        assert_eq!((&20, &mut 21), (k2, v2));
+        assert_eq!(Some(&11), map.get(&10));
+        assert_eq!(Some(&21), map.get(&20));
+        assert_eq!(None, map.get(&30));
+    }
 }
diff --git a/src/map/core.rs b/src/map/core.rs
index c4e725c9..2858331c 100644
--- a/src/map/core.rs
+++ b/src/map/core.rs
@@ -209,7 +209,7 @@ impl<K, V> IndexMapCore<K, V> {
 
     /// Append a key-value pair, *without* checking whether it already exists,
     /// and return the pair's new index.
-    fn push(&mut self, hash: HashValue, key: K, value: V) -> usize {
+    pub(crate) fn push(&mut self, hash: HashValue, key: K, value: V) -> usize {
         let i = self.entries.len();
         self.indices.insert(hash.get(), i, get_hash(&self.entries));
         if i == self.entries.capacity() {
diff --git a/src/set.rs b/src/set.rs
index 134c9045..a2f53c55 100644
--- a/src/set.rs
+++ b/src/set.rs
@@ -277,6 +277,27 @@ where
         self.map.insert(value, ()).is_none()
     }
 
+    /// Insert a value in the set without checking if the value already exists in the set.
+    ///
+    /// Returns a reference to the value just inserted.
+    ///
+    /// This operation is memory-safe.
+    ///
+    /// However, **if a value exists in the set already, the behavior is unspecified**:
+    /// this operation may panic, loop forever, or any following operation with the set
+    /// may panic, loop forever or return arbitrary result. But it won't violate memory safety.
+    ///
+    /// This operation is faster than regular insert because it does not perform
+    /// lookup before insertion.
+    ///
+    /// This operation is useful during initial population of the set.
+    /// For example, when constructing a set from another set, we know
+    /// that values are unique.
+    #[cfg_attr(feature = "inline-more", inline)]
+    pub fn insert_unique_unchecked(&mut self, value: T) -> &T {
+        self.map.insert_unique_unchecked(value, ()).0
+    }
+
     /// Insert the value into the set, and get its index.
     ///
     /// If an equivalent item already exists in the set, it returns