feat(injector): add an extend method to Nucleo's injector

Cargo.toml

@@ -13,8 +13,8 @@ exclude = ["/typos.toml", "/tarpaulin.toml"]
 
 [dependencies]
 nucleo-matcher = { version = "0.3.1", path = "matcher" }
-parking_lot = { version = "0.12.1", features = ["send_guard", "arc_lock"]}
+parking_lot = { version = "0.12.1", features = ["send_guard", "arc_lock"] }
 rayon = "1.7.0"
 
 [workspace]
-members = [ "matcher", "bench" ]
+members = ["matcher", "bench"]

src/boxcar.rs

@@ -182,6 +182,83 @@ impl<T> Vec<T> {
         index
     }
 
+    /// Extends the vector by appending multiple elements at once.
+    pub fn extend<I>(&self, values: I, fill_columns: impl Fn(&T, &mut [Utf32String]))
+    where
+        I: IntoIterator<Item = T> + ExactSizeIterator,
+    {
+        let count: u32 = values
+            .len()
+            .try_into()
+            .expect("overflowed maximum capacity");
+        if count == 0 {
+            return;
+        }
+
+        // Reserve all indices at once
+        let start_index: u32 = self
+            .inflight
+            .fetch_add(u64::from(count), Ordering::Release)
+            .try_into()
+            .expect("overflowed maximum capacity");
+
+        // Compute first and last locations
+        let start_location = Location::of(start_index);
+        let end_location = Location::of(start_index + count - 1);
+
+        // Allocate necessary buckets upfront
+        if start_location.bucket != end_location.bucket {
+            for bucket in start_location.bucket..=end_location.bucket {
+                if let Some(bucket_ptr) = self.buckets.get(bucket as usize) {
+                    Vec::get_or_alloc(bucket_ptr, Location::bucket_len(bucket), self.columns);
+                }
+            }
+        }
+
+        let mut bucket = unsafe { self.buckets.get_unchecked(start_location.bucket as usize) };
+        let mut entries = bucket.entries.load(Ordering::Acquire);
+        if entries.is_null() {
+            entries = Vec::get_or_alloc(
+                bucket,
+                Location::bucket_len(start_location.bucket),
+                self.columns,
+            );
+        }
+        // Route each value to its corresponding bucket
+        let mut location;
+        for (i, v) in values.into_iter().enumerate() {
+            location =
+                Location::of(start_index + u32::try_from(i).expect("overflowed maximum capacity"));
+
+            unsafe {
+                let entry = Bucket::get(entries, location.entry, self.columns);
+
+                // Initialize matcher columns
+                for col in Entry::matcher_cols_raw(entry, self.columns) {
+                    col.get().write(MaybeUninit::new(Utf32String::default()));
+                }
+                fill_columns(&v, Entry::matcher_cols_mut(entry, self.columns));
+                (*entry).slot.get().write(MaybeUninit::new(v));
+                (*entry).active.store(true, Ordering::Release);
+            }
+
+            // if we are at the end of the bucket, move on to the next one
+            if location.entry == location.bucket_len - 1 {
+                // safety: `location.bucket + 1` is always in bounds
+                bucket = unsafe { self.buckets.get_unchecked((location.bucket + 1) as usize) };
+                entries = bucket.entries.load(Ordering::Acquire);
+
+                if entries.is_null() {
+                    entries = Vec::get_or_alloc(
+                        bucket,
+                        Location::bucket_len(location.bucket + 1),
+                        self.columns,
+                    );
+                }
+            }
+        }
+    }
+
     /// race to initialize a bucket
     fn get_or_alloc(bucket: &Bucket<T>, len: u32, cols: u32) -> *mut Entry<T> {
         let entries = unsafe { Bucket::alloc(len, cols) };

src/lib.rs

@@ -79,6 +79,23 @@ impl<T> Injector<T> {
         idx
     }
 
+    /// Appends multiple elements to the list of matched items.
+    /// This function is lock-free and wait-free.
+    ///
+    /// You should favor this function over `push` if at least one of the following is true:
+    /// - the number of items you're adding can be computed beforehand and is typically larger
+    ///     than 1k
+    /// - you're able to batch incoming items
+    /// - you're adding items from multiple threads concurrently (this function results in less
+    ///     contention)
+    pub fn extend<I>(&self, values: I, fill_columns: impl Fn(&T, &mut [Utf32String]))
+    where
+        I: IntoIterator<Item = T> + ExactSizeIterator,
+    {
+        self.items.extend(values, fill_columns);
+        (self.notify)();
+    }
+
     /// Returns the total number of items injected in the matcher. This might
     /// not match the number of items in the match snapshot (if the matcher
     /// is still running)