Add compact()

src/lib.rs

@@ -401,6 +401,102 @@ impl<T> Slab<T> {
         }
     }
 
+    /// Reduce the capacity as much as possible, changing the key for elements when necessary.
+    ///
+    /// To allow updating references to the elements which must be moved to a new key,
+    /// this function takes a closure which is called before moving each element.
+    /// The second and third parameters to the closure are the current key and
+    /// new key respectively.
+    /// In case changing the key for one element turns out not to be possible,
+    /// the move can be cancelled by returning `false` from the closure.
+    /// In that case no further attempts at relocating elements is made.
+    /// If the closure unwinds, the slab will be left in a consistent state,
+    /// but the value that the closure panicked on might be removed.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use slab::*;
+    ///
+    /// let mut slab = Slab::with_capacity(10);
+    /// let a = slab.insert('a');
+    /// slab.insert('b');
+    /// slab.insert('c');
+    /// slab.remove(a);
+    /// slab.compact(|&mut value, from, to| {
+    ///     assert_eq!((value, from, to), ('c', 2, 0));
+    ///     true
+    /// });
+    /// assert!(slab.capacity() >= 2 && slab.capacity() < 10);
+    /// ```
+    ///
+    /// The value is not moved when the closure returns `Err`:
+    ///
+    /// ```
+    /// # use slab::*;
+    ///
+    /// let mut slab = Slab::with_capacity(100);
+    /// let a = slab.insert('a');
+    /// let b = slab.insert('b');
+    /// slab.remove(a);
+    /// slab.compact(|&mut value, from, to| false);
+    /// assert_eq!(slab.iter().next(), Some((b, &'b')));
+    /// ```
+    pub fn compact<F>(&mut self, mut rekey: F)
+    where
+        F: FnMut(&mut T, usize, usize) -> bool,
+    {
+        // If the closure unwinds, we need to restore a valid list of vacant entries
+        struct CleanupGuard<'a, T: 'a> {
+            slab: &'a mut Slab<T>,
+            decrement: bool,
+        }
+        impl<'a, T: 'a> Drop for CleanupGuard<'a, T> {
+            fn drop(&mut self) {
+                if self.decrement {
+                    // Value was popped and not pushed back on
+                    self.slab.len -= 1;
+                }
+                self.slab.recreate_vacant_list();
+            }
+        }
+        let mut guard = CleanupGuard {
+            slab: self,
+            decrement: true,
+        };
+
+        let mut occupied_until = 0;
+        // While there are vacant entries
+        while guard.slab.entries.len() > guard.slab.len {
+            // Find a value that needs to be moved,
+            // by popping entries until we find an occopied one.
+            // (entries cannot be empty because 0 is not greater than anything)
+            if let Some(Entry::Occupied(mut value)) = guard.slab.entries.pop() {
+                // Found one, now find a vacant entry to move it to
+                while let Some(&Entry::Occupied(_)) = guard.slab.entries.get(occupied_until) {
+                    occupied_until += 1;
+                }
+                // Let the caller try to update references to the key
+                if !rekey(&mut value, guard.slab.entries.len(), occupied_until) {
+                    // Changing the key failed, so push the entry back on at its old index.
+                    guard.slab.entries.push(Entry::Occupied(value));
+                    guard.decrement = false;
+                    guard.slab.entries.shrink_to_fit();
+                    return;
+                    // Guard drop handles cleanup
+                }
+                // Put the value in its new spot
+                guard.slab.entries[occupied_until] = Entry::Occupied(value);
+                // ... and mark it as occupied (this is optional)
+                occupied_until += 1;
+            }
+        }
+        guard.slab.next = guard.slab.len;
+        guard.slab.entries.shrink_to_fit();
+        // Normal cleanup is not necessary
+        mem::forget(guard);
+    }
+
     /// Clear the slab of all values.
     ///
     /// # Examples

tests/slab.rs

@@ -2,6 +2,8 @@ extern crate slab;
 
 use slab::*;
 
+use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe};
+
 #[test]
 fn insert_get_remove_one() {
     let mut slab = Slab::new();
@@ -336,6 +338,137 @@ fn shrink_to_fit_doesnt_recreate_list_when_nothing_can_be_done() {
     assert_eq!(slab.vacant_entry().key(), 1);
 }
 
+#[test]
+fn compact_empty() {
+    let mut slab = Slab::new();
+    slab.compact(|_, _, _| panic!());
+    assert_eq!(slab.len(), 0);
+    assert_eq!(slab.capacity(), 0);
+    slab.reserve(20);
+    slab.compact(|_, _, _| panic!());
+    assert_eq!(slab.len(), 0);
+    assert_eq!(slab.capacity(), 0);
+    slab.insert(0);
+    slab.insert(1);
+    slab.insert(2);
+    slab.remove(1);
+    slab.remove(2);
+    slab.remove(0);
+    slab.compact(|_, _, _| panic!());
+    assert_eq!(slab.len(), 0);
+    assert_eq!(slab.capacity(), 0);
+}
+
+#[test]
+fn compact_no_moves_needed() {
+    let mut slab = Slab::new();
+    for i in 0..10 {
+        slab.insert(i);
+    }
+    slab.remove(8);
+    slab.remove(9);
+    slab.remove(6);
+    slab.remove(7);
+    slab.compact(|_, _, _| panic!());
+    assert_eq!(slab.len(), 6);
+    for ((index, &value), want) in slab.iter().zip(0..6) {
+        assert!(index == value);
+        assert_eq!(index, want);
+    }
+    assert!(slab.capacity() >= 6 && slab.capacity() < 10);
+}
+
+#[test]
+fn compact_moves_successfully() {
+    let mut slab = Slab::with_capacity(20);
+    for i in 0..10 {
+        slab.insert(i);
+    }
+    for &i in &[0, 5, 9, 6, 3] {
+        slab.remove(i);
+    }
+    let mut moved = 0;
+    slab.compact(|&mut v, from, to| {
+        assert!(from > to);
+        assert!(from >= 5);
+        assert!(to < 5);
+        assert_eq!(from, v);
+        moved += 1;
+        true
+    });
+    assert_eq!(slab.len(), 5);
+    assert_eq!(moved, 2);
+    assert_eq!(slab.vacant_entry().key(), 5);
+    assert!(slab.capacity() >= 5 && slab.capacity() < 20);
+    let mut iter = slab.iter();
+    assert_eq!(iter.next(), Some((0, &8)));
+    assert_eq!(iter.next(), Some((1, &1)));
+    assert_eq!(iter.next(), Some((2, &2)));
+    assert_eq!(iter.next(), Some((3, &7)));
+    assert_eq!(iter.next(), Some((4, &4)));
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn compact_doesnt_move_if_closure_errors() {
+    let mut slab = Slab::with_capacity(20);
+    for i in 0..10 {
+        slab.insert(i);
+    }
+    for &i in &[9, 3, 1, 4, 0] {
+        slab.remove(i);
+    }
+    slab.compact(|&mut v, from, to| {
+        assert!(from > to);
+        assert_eq!(from, v);
+        v != 6
+    });
+    assert_eq!(slab.len(), 5);
+    assert!(slab.capacity() >= 7 && slab.capacity() < 20);
+    assert_eq!(slab.vacant_entry().key(), 3);
+    let mut iter = slab.iter();
+    assert_eq!(iter.next(), Some((0, &8)));
+    assert_eq!(iter.next(), Some((1, &7)));
+    assert_eq!(iter.next(), Some((2, &2)));
+    assert_eq!(iter.next(), Some((5, &5)));
+    assert_eq!(iter.next(), Some((6, &6)));
+    assert_eq!(iter.next(), None);
+}
+
+#[test]
+fn compact_handles_closure_panic() {
+    let mut slab = Slab::new();
+    for i in 0..10 {
+        slab.insert(i);
+    }
+    for i in 1..6 {
+        slab.remove(i);
+    }
+    let result = catch_unwind(AssertUnwindSafe(|| {
+        slab.compact(|&mut v, from, to| {
+            assert!(from > to);
+            assert_eq!(from, v);
+            if v == 7 {
+                panic!(());
+            }
+            true
+        })
+    }));
+    match result {
+        Err(ref payload) if payload.is::<()>() => {}
+        Err(bug) => resume_unwind(bug),
+        Ok(()) => unreachable!(),
+    }
+    assert_eq!(slab.len(), 5 - 1);
+    assert_eq!(slab.vacant_entry().key(), 3);
+    let mut iter = slab.iter();
+    assert_eq!(iter.next(), Some((0, &0)));
+    assert_eq!(iter.next(), Some((1, &9)));
+    assert_eq!(iter.next(), Some((2, &8)));
+    assert_eq!(iter.next(), Some((6, &6)));
+    assert_eq!(iter.next(), None);
+}
+
 #[test]
 fn fully_consumed_drain() {
     let mut slab = Slab::new();