Replaced instances of SparseSet with EntitySet

cranelift/bforest/src/pool.rs

@@ -83,7 +83,7 @@ impl<F: Forest> NodePool<F> {
         NodeData<F>: fmt::Display,
         F::Key: fmt::Display,
     {
-        use crate::entity::SparseSet;
+        use crate::entity::EntitySet;
         use core::borrow::Borrow;
         use core::cmp::Ordering;
         use std::vec::Vec;
@@ -94,7 +94,13 @@ impl<F: Forest> NodePool<F> {
             assert!(size > 0, "Root must have more than one sub-tree");
         }
 
-        let mut done = SparseSet::new();
+        let mut done = match self[node] {
+            NodeData::Inner { size, .. } | NodeData::Leaf { size, .. } => {
+                EntitySet::with_capacity(size.into())
+            }
+            _ => EntitySet::new(),
+        };
+
         let mut todo = Vec::new();
 
         // Todo-list entries are:
@@ -104,11 +110,7 @@ impl<F: Forest> NodePool<F> {
         todo.push((None, node, None));
 
         while let Some((lkey, node, rkey)) = todo.pop() {
-            assert_eq!(
+            assert!(done.insert(node), "Node appears more than once in tree");
-                done.insert(node),
-                None,
-                "Node appears more than once in tree"
-            );
             let mut lower = lkey;
 
             match self[node] {

cranelift/codegen/src/topo_order.rs

@@ -1,7 +1,7 @@
 //! Topological order of EBBs, according to the dominator tree.
 
 use crate::dominator_tree::DominatorTree;
-use crate::entity::SparseSet;
+use crate::entity::EntitySet;
 use crate::ir::{Ebb, Layout};
 use std::vec::Vec;
 
@@ -19,7 +19,7 @@ pub struct TopoOrder {
     next: usize,
 
     /// Set of visited EBBs.
-    visited: SparseSet<Ebb>,
+    visited: EntitySet<Ebb>,
 
     /// Stack of EBBs to be visited next, already in `visited`.
     stack: Vec<Ebb>,
@@ -31,7 +31,7 @@ impl TopoOrder {
         Self {
             preferred: Vec::new(),
             next: 0,
-            visited: SparseSet::new(),
+            visited: EntitySet::new(),
             stack: Vec::new(),
         }
     }
@@ -64,6 +64,7 @@ impl TopoOrder {
     /// - All EBBs in the `preferred` iterator given to `reset` will be returned.
     /// - All dominators are visited before the EBB returned.
     pub fn next(&mut self, layout: &Layout, domtree: &DominatorTree) -> Option<Ebb> {
+        self.visited.resize(layout.ebb_capacity());
         // Any entries in `stack` should be returned immediately. They have already been added to
         // `visited`.
         while self.stack.is_empty() {
@@ -73,7 +74,7 @@ impl TopoOrder {
                     // We have the next EBB in the preferred order.
                     self.next += 1;
                     // Push it along with any non-visited dominators.
-                    while self.visited.insert(ebb).is_none() {
+                    while self.visited.insert(ebb) {
                         self.stack.push(ebb);
                         match domtree.idom(ebb) {
                             Some(idom) => ebb = layout.inst_ebb(idom).expect("idom not in layout"),

cranelift/codegen/src/verifier/flags.rs

@@ -1,6 +1,6 @@
 //! Verify CPU flags values.
 
-use crate::entity::{SecondaryMap, SparseSet};
+use crate::entity::{EntitySet, SecondaryMap};
 use crate::flowgraph::{BasicBlock, ControlFlowGraph};
 use crate::ir;
 use crate::ir::instructions::BranchInfo;
@@ -50,7 +50,7 @@ impl<'a> FlagsVerifier<'a> {
     fn check(&mut self, errors: &mut VerifierErrors) -> VerifierStepResult<()> {
         // List of EBBs that need to be processed. EBBs may be re-added to this list when we detect
         // that one of their successor blocks needs a live-in flags value.
-        let mut worklist = SparseSet::new();
+        let mut worklist = EntitySet::with_capacity(self.func.layout.ebb_capacity());
         for ebb in self.func.layout.ebbs() {
             worklist.insert(ebb);
         }

cranelift/entity/src/set.rs

@@ -33,6 +33,14 @@ where
         }
     }
 
+    /// Creates a new empty set with the specified capacity.
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            elems: Vec::with_capacity((capacity + 7) / 8),
+            ..Self::new()
+        }
+    }
+
     /// Get the element at `k` if it exists.
     pub fn contains(&self, k: K) -> bool {
         let index = k.index();
@@ -45,9 +53,11 @@ where
 
     /// Is this set completely empty?
     pub fn is_empty(&self) -> bool {
-        // Note that this implementation will become incorrect should it ever become possible
+        if self.len != 0 {
-        // to remove elements from an `EntitySet`.
+            false
-        self.len == 0
+        } else {
+            self.elems.iter().all(|&e| e == 0)
+        }
     }
 
     /// Returns the cardinality of the set.  More precisely, it returns the number of calls to
@@ -93,6 +103,34 @@ where
         self.elems[index / 8] |= 1 << (index % 8);
         result
     }
+
+    /// Removes and returns the entity from the set if it exists.
+    pub fn pop(&mut self) -> Option<K> {
+        if self.len == 0 {
+            return None;
+        }
+
+        // Clear the last known entity in the list.
+        let last_index = self.len - 1;
+        self.elems[last_index / 8] &= !(1 << (last_index % 8));
+
+        // Set the length to the next last stored entity or zero if we pop'ed
+        // the last entity.
+        self.len = self
+            .elems
+            .iter()
+            .enumerate()
+            .rev()
+            .find(|(_, &byte)| byte != 0)
+            // Map `i` from byte index to bit level index.
+            // `(i + 1) * 8` = Last bit in byte.
+            // `last - byte.leading_zeros()` = last set bit in byte.
+            // `as usize` won't ever truncate as the potential range is `0..=8`.
+            .map(|(i, byte)| ((i + 1) * 8) - byte.leading_zeros() as usize)
+            .unwrap_or(0);
+
+        Some(K::new(last_index))
+    }
 }
 
 #[cfg(test)]
@@ -101,7 +139,7 @@ mod tests {
     use core::u32;
 
     // `EntityRef` impl for testing.
-    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+    #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
     struct E(u32);
 
     impl EntityRef for E {
@@ -159,4 +197,51 @@ mod tests {
         m.clear();
         assert!(m.is_empty());
     }
+
+    #[test]
+    fn pop_ordered() {
+        let r0 = E(0);
+        let r1 = E(1);
+        let r2 = E(2);
+        let mut m = EntitySet::new();
+        m.insert(r0);
+        m.insert(r1);
+        m.insert(r2);
+
+        assert_eq!(r2, m.pop().unwrap());
+        assert_eq!(r1, m.pop().unwrap());
+        assert_eq!(r0, m.pop().unwrap());
+        assert!(m.pop().is_none());
+        assert!(m.pop().is_none());
+    }
+
+    #[test]
+    fn pop_unordered() {
+        let mut ebbs = [
+            E(0),
+            E(1),
+            E(6),
+            E(7),
+            E(5),
+            E(9),
+            E(10),
+            E(2),
+            E(3),
+            E(11),
+            E(12),
+        ];
+
+        let mut m = EntitySet::new();
+        for &ebb in &ebbs {
+            m.insert(ebb);
+        }
+        assert_eq!(m.len, 13);
+        ebbs.sort();
+
+        for &ebb in ebbs.iter().rev() {
+            assert_eq!(ebb, m.pop().unwrap());
+        }
+
+        assert!(m.is_empty());
+    }
 }