From c16f7be2d28a88fbf05a25a042478f330757b996 Mon Sep 17 00:00:00 2001
From: Orson Peters <orsonpeters@gmail.com>
Date: Thu, 1 Jul 2021 18:58:57 +0200
Subject: [PATCH] Fix issue #69.

---
 src/basic.rs | 10 +++++++++-
 src/dense.rs | 10 +++++++++-
 src/hop.rs   | 47 ++++++++++++++++++++++++-----------------------
 3 files changed, 42 insertions(+), 25 deletions(-)

diff --git a/src/basic.rs b/src/basic.rs
index 0666ed6..9037195 100644
--- a/src/basic.rs
+++ b/src/basic.rs
@@ -1246,7 +1246,7 @@ mod serialize {
 mod tests {
     use super::*;
     use quickcheck::quickcheck;
-    use std::collections::HashMap;
+    use std::collections::{HashSet, HashMap};
 
     #[derive(Clone)]
     struct CountDrop<'a>(&'a std::cell::RefCell<usize>);
@@ -1364,6 +1364,14 @@ mod tests {
 
                     // Delete.
                     1 => {
+                        // 10% of the time test clear.
+                        if val % 10 == 0 {
+                            let hmvals: HashSet<_> = hm.drain().map(|(_, v)| v).collect();
+                            let smvals: HashSet<_> = sm.drain().map(|(_, v)| v).collect();
+                            if hmvals != smvals {
+                                return false;
+                            }
+                        }
                         if hm_keys.is_empty() { continue; }
 
                         let idx = val as usize % hm_keys.len();
diff --git a/src/dense.rs b/src/dense.rs
index 926e3bc..6f0b771 100644
--- a/src/dense.rs
+++ b/src/dense.rs
@@ -1114,7 +1114,7 @@ mod serialize {
 mod tests {
     use super::*;
     use quickcheck::quickcheck;
-    use std::collections::HashMap;
+    use std::collections::{HashSet, HashMap};
 
     #[derive(Clone)]
     struct CountDrop<'a>(&'a core::cell::RefCell<usize>);
@@ -1231,6 +1231,14 @@ mod tests {
 
                     // Delete.
                     1 => {
+                        // 10% of the time test clear.
+                        if val % 10 == 0 {
+                            let hmvals: HashSet<_> = hm.drain().map(|(_, v)| v).collect();
+                            let smvals: HashSet<_> = sm.drain().map(|(_, v)| v).collect();
+                            if hmvals != smvals {
+                                return false;
+                            }
+                        }
                         if hm_keys.is_empty() { continue; }
 
                         let idx = val as usize % hm_keys.len();
diff --git a/src/hop.rs b/src/hop.rs
index 57e7857..3e35d04 100644
--- a/src/hop.rs
+++ b/src/hop.rs
@@ -622,8 +622,7 @@ impl<K: Key, V> HopSlotMap<K, V> {
     /// ```
     pub fn drain(&mut self) -> Drain<K, V> {
         Drain {
-            cur: 0,
-            num_left: self.len(),
+            cur: unsafe { self.slots.get_unchecked(0).u.free.other_end as usize + 1 },
             sm: self,
         }
     }
@@ -977,7 +976,6 @@ impl<K: Key, V> IndexMut<K> for HopSlotMap<K, V> {
 #[derive(Debug)]
 pub struct Drain<'a, K: Key + 'a, V: 'a> {
     cur: usize,
-    num_left: usize,
     sm: &'a mut HopSlotMap<K, V>,
 }
 
@@ -1043,33 +1041,28 @@ impl<'a, K: Key, V> Iterator for Drain<'a, K, V> {
     type Item = (K, V);
 
     fn next(&mut self) -> Option<(K, V)> {
-        if self.cur >= self.sm.slots.len() {
+        // All unchecked indices are safe due to the invariants of the freelist
+        // and that self.sm.len() guarantees there is another element.
+        if self.sm.len() == 0 {
             return None;
         }
 
-        let (idx, version) = {
-            let slot = &self.sm.slots[self.cur];
-            match slot.get() {
-                Occupied(_) => (self.cur, slot.version),
-                Vacant(free) => {
-                    // Skip block of contiguous vacant slots.
-                    let idx = free.other_end as usize + 1;
-                    if idx >= self.sm.slots.len() {
-                        return None;
-                    }
-                    (idx, self.sm.slots[idx].version)
-                }
-            }
+        // Skip ahead to next element. Must do this before removing.
+        let idx = self.cur;
+        self.cur = match self.sm.slots.get(idx + 1).map(|s| s.get()) {
+            Some(Occupied(_)) => idx + 1,
+            Some(Vacant(free)) => free.other_end as usize + 1,
+            None => 0,
         };
 
-        self.cur = idx + 1;
-        self.num_left -= 1;
-        let key = KeyData::new(idx as u32, version).into();
-        Some((key, unsafe { self.sm.remove_from_slot(idx) }))
+        let key = KeyData::new(idx as u32, unsafe {
+            self.sm.slots.get_unchecked(idx).version
+        });
+        Some((key.into(), unsafe { self.sm.remove_from_slot(idx) }))
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        (self.num_left, Some(self.num_left))
+        (self.sm.len(), Some(self.sm.len()))
     }
 }
 
@@ -1397,7 +1390,7 @@ mod serialize {
 mod tests {
     use super::*;
     use quickcheck::quickcheck;
-    use std::collections::HashMap;
+    use std::collections::{HashSet, HashMap};
 
     #[derive(Clone)]
     struct CountDrop<'a>(&'a std::cell::RefCell<usize>);
@@ -1515,6 +1508,14 @@ mod tests {
 
                     // Delete.
                     1 => {
+                        // 10% of the time test clear.
+                        if val % 10 == 0 {
+                            let hmvals: HashSet<_> = hm.drain().map(|(_, v)| v).collect();
+                            let smvals: HashSet<_> = sm.drain().map(|(_, v)| v).collect();
+                            if hmvals != smvals {
+                                return false;
+                            }
+                        }
                         if hm_keys.is_empty() { continue; }
 
                         let idx = val as usize % hm_keys.len();