Avoid extra usize array by unioning next pointer into elements

src/lib.rs

@@ -77,24 +77,37 @@ impl<IT, N> Key<IT, N> {
     }
 }
 
+/// Module to hide the Entry type which needs to be public due to the generic-array internals.
+mod entry {
+    pub enum Entry<IT> {
+        Used(IT),
+        EmptyNext(usize),
+        EmptyLast
+    }
+}
+
+use entry::Entry;
+
 // Data type that stores values and returns a key that can be used to manipulate
 // the stored values.
 // Values can be read by anyone but can only be modified using the key.
 pub struct Slots<IT, N>
-    where N: ArrayLength<Option<IT>> + ArrayLength<usize> + Unsigned {
-    items: GenericArray<Option<IT>, N>,
-    free_list: GenericArray<usize, N>,
+    where N: ArrayLength<Entry<IT>> + Unsigned {
+    items: GenericArray<Entry<IT>, N>,
+    // Could be optimized by making it just usize and relying on free_count to determine its
+    // validity
+    next_free: Option<usize>,
     free_count: usize
 }
 
 impl<IT, N> Slots<IT, N>
-    where N: ArrayLength<Option<IT>> + ArrayLength<usize> + Unsigned {
+    where N: ArrayLength<Entry<IT>> + Unsigned {
     pub fn new() -> Self {
         let size = N::to_usize();
 
         Self {
-            items: GenericArray::default(),
-            free_list: GenericArray::generate(|i: usize| size - i - 1),
+            items: GenericArray::generate(|i| i.checked_sub(1).map(Entry::EmptyNext).unwrap_or(Entry::EmptyLast)),
+            next_free: size.checked_sub(1),
             free_count: size
         }
     }
@@ -108,37 +121,41 @@ impl<IT, N> Slots<IT, N>
     }
 
     fn free(&mut self, idx: usize) {
-        self.free_list[self.free_count] = idx;
+        self.items[idx] = match self.next_free {
+            Some(n) => Entry::EmptyNext(n),
+            None => Entry::EmptyLast
+        };
+        self.next_free = Some(idx);
         self.free_count += 1;
     }
 
     fn alloc(&mut self) -> Option<usize> {
-        if self.count() == self.capacity() {
-            None
-        } else {
-            let i = self.free_list[self.free_count - 1];
-            self.free_count -= 1;
-            Some(i)
-        }
+        let index = self.next_free?;
+        self.next_free = match self.items[index] {
+            Entry::EmptyNext(n) => Some(n),
+            Entry::EmptyLast => None,
+            _ => unreachable!("Non-empty item in entry behind free chain"),
+        };
+        self.free_count -= 1;
+        Some(index)
     }
 
     pub fn store(&mut self, item: IT) -> Result<Key<IT, N>, IT> {
         match self.alloc() {
             Some(i) => {
-                self.items[i] = Some(item);
+                self.items[i] = Entry::Used(item);
                 Ok(Key::new(i))
             }
             None => Err(item)
         }
     }
 
     pub fn take(&mut self, key: Key<IT, N>) -> IT {
-        match self.items[key.index].take() {
-            Some(item) => {
-                self.free(key.index);
-                item
-            }
-            None => panic!()
+        let taken = core::mem::replace(&mut self.items[key.index], Entry::EmptyLast);
+        self.free(key.index);
+        match taken {
+            Entry::Used(item) => item,
+            _ => panic!()
         }
     }
 
@@ -151,15 +168,15 @@ impl<IT, N> Slots<IT, N>
 
     pub fn try_read<T, F>(&self, key: usize, function: F) -> Option<T> where F: Fn(&IT) -> T {
         match &self.items[key] {
-            Some(item) => Some(function(&item)),
-            None => None
+            Entry::Used(item) => Some(function(&item)),
+            _ => None
         }
     }
 
     pub fn modify<T, F>(&mut self, key: &Key<IT, N>, function: F) -> T where F: Fn(&mut IT) -> T {
         match self.items[key.index] {
-            Some(ref mut item) => function(item),
-            None => panic!()
+            Entry::Used(ref mut item) => function(item),
+            _ => panic!()
         }
     }
 }

tests/test.rs

@@ -33,9 +33,9 @@ fn index_can_be_used_to_read_value() {
     slots.store(6).unwrap();
     slots.store(7).unwrap();
 
-    assert_eq!(5, slots.try_read(0, |&w| w).unwrap());
-    assert_eq!(6, slots.try_read(1, |&w| w).unwrap());
-    assert_eq!(7, slots.try_read(2, |&w| w).unwrap());
+    assert_eq!(5, slots.try_read(7, |&w| w).unwrap());
+    assert_eq!(6, slots.try_read(6, |&w| w).unwrap());
+    assert_eq!(7, slots.try_read(5, |&w| w).unwrap());
 }
 
 #[test]
@@ -83,3 +83,54 @@ fn store_returns_err_when_full() {
 
     assert!(k2.is_err());
 }
+
+#[should_panic(expected = "assertion failed: `(left == right)`\n  left: `792`,\n right: `536`")]
+#[test]
+/// Verify some size bounds: an N long array over IT is not larger than 3 usize + N * IT (as long
+/// as IT is larger than two usize and has two niches)
+//
+// Fails until https://github.com/rust-lang/rust/issues/46213 is resolved (possibly,
+// https://github.com/rust-lang/rust/pull/70477 is sufficient). When this starts not failing any
+// more, be happy, remove the panic, and figure out how to skip the test on older Rust versions.
+// (If left just goes down but does not reach right, that should be investigated further, as it
+// indicates that the optimization was implemented incompletely, or it turns out it is not possible
+// for some reasons and needs fixing in the code).
+fn is_compact() {
+    #[allow(unused)]
+    struct TwoNichesIn16Byte {
+        n1: u64,
+        n2: u32,
+        n3: u16,
+        n4: u8,
+        b: bool,
+    }
+
+    assert_eq!(core::mem::size_of::<TwoNichesIn16Byte>(), 16);
+
+    assert_eq!(core::mem::size_of::<Slots<TwoNichesIn16Byte, U32>>(), 32 * 16 + 3 * core::mem::size_of::<usize>());
+}
+
+#[test]
+fn capacity_and_count() {
+    let mut slots: Slots<u8, U4> = Slots::new();
+
+    assert_eq!(slots.capacity(), 4);
+    assert_eq!(slots.count(), 0);
+
+    let k1 = slots.store(1).unwrap();
+    let k2 = slots.store(2).unwrap();
+
+    assert_eq!(slots.count(), 2);
+
+    let k3 = slots.store(3).unwrap();
+    let k4 = slots.store(4).unwrap();
+
+    assert_eq!(slots.count(), 4);
+
+    slots.take(k1);
+    slots.take(k2);
+    slots.take(k3);
+    slots.take(k4);
+
+    assert_eq!(slots.count(), 0);
+}