Fix panic in FromIterator impl

tokio-rs · Jun 6, 2021 · 7804ecb · 7804ecb
1 parent 6c5fe73
commit 7804ecb
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 0 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -1182,6 +1182,7 @@ impl<T> FromIterator<(usize, T)> for Slab<T> {
         let mut slab = Self::with_capacity(iterator.size_hint().0);
 
         let mut vacant_list_broken = false;
+        let mut first_vacant_index = None;
         for (key, value) in iterator {
             if key < slab.entries.len() {
                 // iterator is not sorted, might need to recreate vacant list
@@ -1193,6 +1194,9 @@ impl<T> FromIterator<(usize, T)> for Slab<T> {
                 // This is consistent with HashMap and BtreeMap
                 slab.entries[key] = Entry::Occupied(value);
             } else {
+                if first_vacant_index.is_none() && slab.entries.len() < key {
+                    first_vacant_index = Some(slab.entries.len());
+                }
                 // insert holes as necessary
                 while slab.entries.len() < key {
                     // add the entry to the start of the vacant list
@@ -1209,7 +1213,16 @@ impl<T> FromIterator<(usize, T)> for Slab<T> {
             slab.next = slab.entries.len();
         } else if vacant_list_broken {
             slab.recreate_vacant_list();
+        } else if let Some(first_vacant_index) = first_vacant_index {
+            let next = slab.entries.len();
+            match &mut slab.entries[first_vacant_index] {
+                Entry::Vacant(n) => *n = next,
+                _ => unreachable!(),
+            }
+        } else {
+            unreachable!()
         }
+
         slab
     }
 }

diff --git a/src/serde.rs b/src/serde.rs
@@ -43,6 +43,7 @@ where
 
         // same as FromIterator impl
         let mut vacant_list_broken = false;
+        let mut first_vacant_index = None;
         while let Some((key, value)) = map.next_entry()? {
             if key < slab.entries.len() {
                 // iterator is not sorted, might need to recreate vacant list
@@ -54,6 +55,9 @@ where
                 // This is consistent with HashMap and BtreeMap
                 slab.entries[key] = Entry::Occupied(value);
             } else {
+                if first_vacant_index.is_none() && slab.entries.len() < key {
+                    first_vacant_index = Some(slab.entries.len());
+                }
                 // insert holes as necessary
                 while slab.entries.len() < key {
                     // add the entry to the start of the vacant list
@@ -70,6 +74,12 @@ where
             slab.next = slab.entries.len();
         } else if vacant_list_broken {
             slab.recreate_vacant_list();
+        } else if let Some(first_vacant_index) = first_vacant_index {
+            let next = slab.entries.len();
+            match &mut slab.entries[first_vacant_index] {
+                Entry::Vacant(n) => *n = next,
+                _ => unreachable!(),
+            }
         }
 
         Ok(slab)

diff --git a/tests/slab.rs b/tests/slab.rs
@@ -403,6 +403,36 @@ fn from_iterator_unordered() {
     assert_eq!(iter.next(), None);
 }
 
+// https://github.com/tokio-rs/slab/issues/100
+#[test]
+fn from_iterator_issue_100() {
+    let mut slab: slab::Slab<()> = vec![(1, ())].into_iter().collect();
+    assert_eq!(slab.len(), 1);
+    assert_eq!(slab.insert(()), 0);
+    assert_eq!(slab.insert(()), 2);
+    assert_eq!(slab.insert(()), 3);
+
+    let mut slab: slab::Slab<()> = vec![(1, ()), (2, ())].into_iter().collect();
+    assert_eq!(slab.len(), 2);
+    assert_eq!(slab.insert(()), 0);
+    assert_eq!(slab.insert(()), 3);
+    assert_eq!(slab.insert(()), 4);
+
+    let mut slab: slab::Slab<()> = vec![(1, ()), (3, ())].into_iter().collect();
+    assert_eq!(slab.len(), 2);
+    assert_eq!(slab.insert(()), 2);
+    assert_eq!(slab.insert(()), 0);
+    assert_eq!(slab.insert(()), 4);
+
+    let mut slab: slab::Slab<()> = vec![(0, ()), (2, ()), (3, ()), (5, ())]
+        .into_iter()
+        .collect();
+    assert_eq!(slab.len(), 4);
+    assert_eq!(slab.insert(()), 4);
+    assert_eq!(slab.insert(()), 1);
+    assert_eq!(slab.insert(()), 6);
+}
+
 #[test]
 fn clear() {
     let mut slab = Slab::new();