Added in-place dedup, dedup_by and dedup_by_key

src/arrayvec.rs

@@ -378,7 +378,7 @@ impl<T, const CAP: usize> ArrayVec<T, CAP> {
 
     /// Remove the element at `index` and swap the last element into its place.
     ///
-    /// This is a checked version of `.swap_remove`.  
+    /// This is a checked version of `.swap_remove`.
     /// This operation is O(1).
     ///
     /// Return `Some(` *element* `)` if the index is in bounds, else `None`.
@@ -740,6 +740,167 @@ impl<T, const CAP: usize> ArrayVec<T, CAP> {
     pub fn as_mut_ptr(&mut self) -> *mut T {
         ArrayVecImpl::as_mut_ptr(self)
     }
+
+    /// Removes all but the first of consecutive elements in the vector that resolve to the same
+    /// key.
+    ///
+    /// If the vector is sorted, this removes all duplicates.
+    #[inline]
+    pub fn dedup_by_key<F, K>(&mut self, mut key: F)
+    where
+        F: FnMut(&mut T) -> K,
+        K: PartialEq,
+    {
+        self.dedup_by(|a, b| key(a) == key(b))
+    }
+
+    /// Removes all but the first of consecutive elements in the vector satisfying a given equality
+    /// relation.
+    ///
+    /// The `same_bucket` function is passed references to two elements from the vector and
+    /// must determine if the elements compare equal. The elements are passed in opposite order
+    /// from their order in the slice, so if `same_bucket(a, b)` returns `true`, `a` is removed.
+    ///
+    /// If the vector is sorted, this removes all duplicates.
+    #[inline]
+    pub fn dedup_by<F>(&mut self, mut same_bucket: F)
+    where
+        F: FnMut(&mut T, &mut T) -> bool,
+    {
+        let len = self.len();
+        if len <= 1 {
+            return;
+        }
+
+        // Check if we ever want to remove anything.
+        // This allows to use copy_non_overlapping in next cycle.
+        // And avoids any memory writes if we don't need to remove anything.
+        let mut first_duplicate_idx: usize = 1;
+        let start = self.as_mut_ptr();
+        while first_duplicate_idx != len {
+            let found_duplicate = unsafe {
+                // SAFETY: first_duplicate always in range [1..len)
+                // Note that we start iteration from 1 so we never overflow.
+                let prev = start.add(first_duplicate_idx.wrapping_sub(1));
+                let current = start.add(first_duplicate_idx);
+                // We explicitly say in docs that references are reversed.
+                same_bucket(&mut *current, &mut *prev)
+            };
+            if found_duplicate {
+                break;
+            }
+            first_duplicate_idx += 1;
+        }
+        // Don't need to remove anything.
+        // We cannot get bigger than len.
+        if first_duplicate_idx == len {
+            return;
+        }
+
+        /* INVARIANT: vec.len() > read > write > write-1 >= 0 */
+        struct FillGapOnDrop<'a, T, const CAP: usize> {
+            /* Offset of the element we want to check if it is duplicate */
+            read: usize,
+
+            /* Offset of the place where we want to place the non-duplicate
+             * when we find it. */
+            write: usize,
+
+            /* The ArrayVec that would need correction if `same_bucket` panicked */
+            vec: &'a mut ArrayVec<T, CAP>,
+        }
+
+        impl<'a, T, const CAP: usize> Drop for FillGapOnDrop<'a, T, CAP> {
+            fn drop(&mut self) {
+                /* This code gets executed when `same_bucket` panics */
+
+                /* SAFETY: invariant guarantees that `read - write`
+                 * and `len - read` never overflow and that the copy is always
+                 * in-bounds. */
+                unsafe {
+                    let ptr = self.vec.as_mut_ptr();
+                    let len = self.vec.len();
+
+                    /* How many items were left when `same_bucket` panicked.
+                     * Basically vec[read..].len() */
+                    let items_left = len.wrapping_sub(self.read);
+
+                    /* Pointer to first item in vec[write..write+items_left] slice */
+                    let dropped_ptr = ptr.add(self.write);
+                    /* Pointer to first item in vec[read..] slice */
+                    let valid_ptr = ptr.add(self.read);
+
+                    /* Copy `vec[read..]` to `vec[write..write+items_left]`.
+                     * The slices can overlap, so `copy_nonoverlapping` cannot be used */
+                    ptr::copy(valid_ptr, dropped_ptr, items_left);
+
+                    /* How many items have been already dropped
+                     * Basically vec[read..write].len() */
+                    let dropped = self.read.wrapping_sub(self.write);
+
+                    self.vec.set_len(len - dropped);
+                }
+            }
+        }
+
+        /* Drop items while going through Vec, it should be more efficient than
+         * doing slice partition_dedup + truncate */
+
+        // Construct gap first and then drop item to avoid memory corruption if `T::drop` panics.
+        let mut gap =
+            FillGapOnDrop { read: first_duplicate_idx + 1, write: first_duplicate_idx, vec: self };
+        unsafe {
+            // SAFETY: we checked that first_duplicate_idx in bounds before.
+            // If drop panics, `gap` would remove this item without drop.
+            ptr::drop_in_place(start.add(first_duplicate_idx));
+        }
+
+        /* SAFETY: Because of the invariant, read_ptr, prev_ptr and write_ptr
+         * are always in-bounds and read_ptr never aliases prev_ptr */
+        unsafe {
+            while gap.read < len {
+                let read_ptr = start.add(gap.read);
+                let prev_ptr = start.add(gap.write.wrapping_sub(1));
+
+                // We explicitly say in docs that references are reversed.
+                let found_duplicate = same_bucket(&mut *read_ptr, &mut *prev_ptr);
+                if found_duplicate {
+                    // Increase `gap.read` now since the drop may panic.
+                    gap.read += 1;
+                    /* We have found duplicate, drop it in-place */
+                    ptr::drop_in_place(read_ptr);
+                } else {
+                    let write_ptr = start.add(gap.write);
+
+                    /* read_ptr cannot be equal to write_ptr because at this point
+                     * we guaranteed to skip at least one element (before loop starts).
+                     */
+                    ptr::copy_nonoverlapping(read_ptr, write_ptr, 1);
+
+                    /* We have filled that place, so go further */
+                    gap.write += 1;
+                    gap.read += 1;
+                }
+            }
+
+            /* Technically we could let `gap` clean up with its Drop, but
+             * when `same_bucket` is guaranteed to not panic, this bloats a little
+             * the codegen, so we just do it manually */
+            gap.vec.set_len(gap.write);
+            mem::forget(gap);
+        }
+    }
+}
+
+impl<T: PartialEq, const CAP: usize> ArrayVec<T, CAP> {
+    /// Removes consecutive repeated elements in the vector according to the
+    /// [`PartialEq`] trait implementation.
+    ///
+    /// If the vector is sorted, this removes all duplicates.
+    #[inline]
+    pub fn dedup(&mut self) {
+        self.dedup_by(|a, b| a == b)
+    }
 }
 
 impl<T, const CAP: usize> ArrayVecImpl for ArrayVec<T, CAP> {
@@ -1084,11 +1245,11 @@ impl Drop for WritebackGuard<'_> {
 
 
 /// Extend the `ArrayVec` with an iterator.
-/// 
+///
 /// ***Panics*** if extending the vector exceeds its capacity.
 impl<T, const CAP: usize> Extend<T> for ArrayVec<T, CAP> {
     /// Extend the `ArrayVec` with an iterator.
-    /// 
+    ///
     /// ***Panics*** if extending the vector exceeds its capacity.
     #[track_caller]
     fn extend<I: IntoIterator<Item=T>>(&mut self, iter: I) {
@@ -1154,11 +1315,11 @@ impl<T, const CAP: usize> ArrayVec<T, CAP> {
 }
 
 /// Create an `ArrayVec` from an iterator.
-/// 
+///
 /// ***Panics*** if the number of elements in the iterator exceeds the arrayvec's capacity.
 impl<T, const CAP: usize> iter::FromIterator<T> for ArrayVec<T, CAP> {
     /// Create an `ArrayVec` from an iterator.
-    /// 
+    ///
     /// ***Panics*** if the number of elements in the iterator exceeds the arrayvec's capacity.
     fn from_iter<I: IntoIterator<Item=T>>(iter: I) -> Self {
         let mut array = ArrayVec::new();

tests/tests.rs

@@ -814,10 +814,44 @@ fn test_arraystring_const_constructible() {
     assert_eq!(var, *"hello");
 }
 
-
 #[test]
 fn test_arraystring_zero_filled_has_some_sanity_checks() {
     let string = ArrayString::<4>::zero_filled();
     assert_eq!(string.as_str(), "\0\0\0\0");
     assert_eq!(string.len(), 4);
 }
+
+#[test]
+fn test_dedup_unordered() {
+    let mut vec: ArrayVec<usize, 8> = ArrayVec::new_const();
+
+    vec.push(0);
+    vec.push(10);
+    vec.push(10);
+    vec.push(5);
+    vec.push(0);
+    vec.push(10);
+    vec.push(7);
+    vec.push(7);
+
+    let () = vec.dedup();
+    assert_eq!(&vec[..], &[0, 10, 5, 0, 10, 7]);
+}
+
+#[test]
+fn test_dedup_ordered() {
+    let mut vec: ArrayVec<usize, 8> = ArrayVec::new_const();
+
+    vec.push(0);
+    vec.push(10);
+    vec.push(10);
+    vec.push(5);
+    vec.push(0);
+    vec.push(10);
+    vec.push(7);
+    vec.push(7);
+
+    let () = vec.sort_unstable();
+    let () = vec.dedup();
+    assert_eq!(&vec[..], &[0, 5, 7, 10]);
+}