`cranelift-entity`: Implement `EntitySet` in terms of `cranelift_bitset::CompoundBitSet` (#8834)

* cranelift-entity: Implement `EntitySet` in terms of `cranelift_bitset::CompoundBitSet`

* Shrink the size of `CompoundBitSet` so we don't perturb vmctx size test expectations

* Update vmctx size test expectations anyways because we shrunk "too much"

* Move `cranelift-bitset` to the front of `CRATES_TO_PUBLISH`

Cargo.lock

@@ -657,6 +657,7 @@ dependencies = [
 name = "cranelift-entity"
 version = "0.110.0"
 dependencies = [
+ "cranelift-bitset",
  "serde",
  "serde_derive",
 ]

cranelift/bitset/src/compound.rs

@@ -1,8 +1,8 @@
 //! Compound bit sets.
 
 use crate::scalar::{self, ScalarBitSet};
-use alloc::{vec, vec::Vec};
-use core::mem;
+use alloc::boxed::Box;
+use core::{cmp, iter, mem};
 
 /// A large bit set backed by dynamically-sized storage.
 ///
@@ -40,15 +40,14 @@ use core::mem;
 /// let elems: Vec<_> = bitset.iter().collect();
 /// assert_eq!(elems, [444, 555]);
 /// ```
-#[derive(Clone, PartialEq, Eq)]
+#[derive(Clone, Default, PartialEq, Eq)]
 #[cfg_attr(
     feature = "enable-serde",
     derive(serde_derive::Serialize, serde_derive::Deserialize)
 )]
 pub struct CompoundBitSet {
-    elems: Vec<ScalarBitSet<usize>>,
-    len: usize,
-    max: Option<usize>,
+    elems: Box<[ScalarBitSet<usize>]>,
+    max: Option<u32>,
 }
 
 impl core::fmt::Debug for CompoundBitSet {
@@ -58,13 +57,6 @@ impl core::fmt::Debug for CompoundBitSet {
     }
 }
 
-impl Default for CompoundBitSet {
-    #[inline]
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 const BITS_PER_WORD: usize = mem::size_of::<usize>() * 8;
 
 impl CompoundBitSet {
@@ -81,11 +73,7 @@ impl CompoundBitSet {
     /// ```
     #[inline]
     pub fn new() -> Self {
-        CompoundBitSet {
-            elems: vec![],
-            len: 0,
-            max: None,
-        }
+        CompoundBitSet::default()
     }
 
     /// Construct a new, empty bit set with space reserved to store any element
@@ -129,7 +117,7 @@ impl CompoundBitSet {
     /// ```
     #[inline]
     pub fn len(&self) -> usize {
-        self.len
+        self.elems.iter().map(|sub| usize::from(sub.len())).sum()
     }
 
     /// Get `n + 1` where `n` is the largest value that can be stored inside
@@ -156,7 +144,7 @@ impl CompoundBitSet {
     /// assert!(bitset.capacity() >= 999);
     ///```
     pub fn capacity(&self) -> usize {
-        self.elems.capacity() * BITS_PER_WORD
+        self.elems.len() * BITS_PER_WORD
     }
 
     /// Is this bitset empty?
@@ -176,7 +164,7 @@ impl CompoundBitSet {
     /// ```
     #[inline]
     pub fn is_empty(&self) -> bool {
-        self.len == 0
+        self.len() == 0
     }
 
     /// Convert an element `i` into the `word` that can be used to index into
@@ -255,7 +243,23 @@ impl CompoundBitSet {
     pub fn ensure_capacity(&mut self, n: usize) {
         let (word, _bit) = Self::word_and_bit(n);
         if word >= self.elems.len() {
-            self.elems.resize_with(word + 1, ScalarBitSet::new);
+            assert!(word < usize::try_from(isize::MAX).unwrap());
+
+            let delta = word - self.elems.len();
+            let to_grow = delta + 1;
+
+            // Amortize the cost of growing.
+            let to_grow = cmp::max(to_grow, self.elems.len() * 2);
+            // Don't make ridiculously small allocations.
+            let to_grow = cmp::max(to_grow, 4);
+
+            let new_elems = self
+                .elems
+                .iter()
+                .copied()
+                .chain(iter::repeat(ScalarBitSet::new()).take(to_grow))
+                .collect::<Box<[_]>>();
+            self.elems = new_elems;
         }
     }
 
@@ -290,10 +294,13 @@ impl CompoundBitSet {
     #[inline]
     pub fn insert(&mut self, i: usize) -> bool {
         self.ensure_capacity(i + 1);
+
         let (word, bit) = Self::word_and_bit(i);
         let is_new = self.elems[word].insert(bit);
-        self.len += is_new as usize;
-        self.max = self.max.map(|max| core::cmp::max(max, i)).or(Some(i));
+
+        let i = u32::try_from(i).unwrap();
+        self.max = self.max.map(|max| cmp::max(max, i)).or(Some(i));
+
         is_new
     }
 
@@ -323,8 +330,7 @@ impl CompoundBitSet {
         if word < self.elems.len() {
             let sub = &mut self.elems[word];
             let was_present = sub.remove(bit);
-            self.len -= was_present as usize;
-            if was_present && self.max == Some(i) {
+            if was_present && self.max() == Some(i) {
                 self.update_max(word);
             }
             was_present
@@ -341,7 +347,7 @@ impl CompoundBitSet {
             .rev()
             .filter_map(|(word, sub)| {
                 let bit = sub.max()?;
-                Some(Self::elem(word, bit))
+                Some(u32::try_from(Self::elem(word, bit)).unwrap())
             })
             .next();
     }
@@ -367,7 +373,7 @@ impl CompoundBitSet {
     /// ```
     #[inline]
     pub fn max(&self) -> Option<usize> {
-        self.max
+        self.max.map(|m| usize::try_from(m).unwrap())
     }
 
     /// Removes and returns the largest value in this set.
@@ -396,7 +402,7 @@ impl CompoundBitSet {
     /// ```
     #[inline]
     pub fn pop(&mut self) -> Option<usize> {
-        let max = self.max?;
+        let max = self.max()?;
         self.remove(max);
         Some(max)
     }
@@ -420,8 +426,15 @@ impl CompoundBitSet {
     /// ```
     #[inline]
     pub fn clear(&mut self) {
-        self.elems.clear();
-        self.len = 0;
+        let max = match self.max() {
+            Some(max) => max,
+            None => return,
+        };
+        let (word, _bit) = Self::word_and_bit(max);
+        debug_assert!(self.elems[word + 1..].iter().all(|sub| sub.is_empty()));
+        for sub in &mut self.elems[..=word] {
+            *sub = ScalarBitSet::new();
+        }
         self.max = None;
     }
 

cranelift/entity/Cargo.toml

@@ -15,8 +15,9 @@ edition.workspace = true
 workspace = true
 
 [dependencies]
+cranelift-bitset = { workspace=true }
 serde = { workspace = true, optional = true }
 serde_derive = { workspace = true, optional = true }
 
 [features]
-enable-serde = ["serde", "serde_derive"]
+enable-serde = ["serde", "serde_derive", "cranelift-bitset/enable-serde"]

cranelift/entity/src/set.rs

@@ -2,11 +2,8 @@
 
 use crate::keys::Keys;
 use crate::EntityRef;
-use alloc::vec::Vec;
 use core::marker::PhantomData;
-
-// How many bits are used to represent a single element in `EntitySet`.
-const BITS: usize = core::mem::size_of::<usize>() * 8;
+use cranelift_bitset::CompoundBitSet;
 
 /// A set of `K` for densely indexed entity references.
 ///
@@ -17,16 +14,14 @@ pub struct EntitySet<K>
 where
     K: EntityRef,
 {
-    elems: Vec<usize>,
-    len: usize,
+    bitset: CompoundBitSet,
     unused: PhantomData<K>,
 }
 
 impl<K: EntityRef> Default for EntitySet<K> {
     fn default() -> Self {
         Self {
-            elems: Vec::new(),
-            len: 0,
+            bitset: CompoundBitSet::default(),
             unused: PhantomData,
         }
     }
@@ -45,91 +40,49 @@ where
     /// Creates a new empty set with the specified capacity.
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
-            elems: Vec::with_capacity((capacity + (BITS - 1)) / BITS),
-            ..Self::new()
+            bitset: CompoundBitSet::with_capacity(capacity),
+            unused: PhantomData,
         }
     }
 
     /// Get the element at `k` if it exists.
     pub fn contains(&self, k: K) -> bool {
         let index = k.index();
-        if index < self.len {
-            (self.elems[index / BITS] & (1 << (index % BITS))) != 0
-        } else {
-            false
-        }
+        self.bitset.contains(index)
     }
 
     /// Is this set completely empty?
     pub fn is_empty(&self) -> bool {
-        if self.len != 0 {
-            false
-        } else {
-            self.elems.iter().all(|&e| e == 0)
-        }
+        self.bitset.is_empty()
     }
 
     /// Remove all entries from this set.
     pub fn clear(&mut self) {
-        self.len = 0;
-        self.elems.clear()
+        self.bitset.clear()
     }
 
     /// Iterate over all the keys in this set.
     pub fn keys(&self) -> Keys<K> {
-        Keys::with_len(self.len)
-    }
-
-    /// Resize the set to have `n` entries by adding default entries as needed.
-    pub fn resize(&mut self, n: usize) {
-        self.elems.resize((n + (BITS - 1)) / BITS, 0);
-        self.len = n
+        Keys::with_len(self.bitset.max().map_or(0, |x| x + 1))
     }
 
     /// Insert the element at `k`.
     pub fn insert(&mut self, k: K) -> bool {
         let index = k.index();
-        if index >= self.len {
-            self.resize(index + 1)
-        }
-        let result = !self.contains(k);
-        self.elems[index / BITS] |= 1 << (index % BITS);
-        result
+        self.bitset.insert(index)
     }
 
     /// 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 / BITS] &= !(1 << (last_index % BITS));
-
-        // 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(|(_, &elem)| elem != 0)
-            // Map `i` from `elem` index to bit level index.
-            // `(i + 1) * BITS` = Last bit in `elem`.
-            // `last - elem.leading_zeros()` = last set bit in `elem`.
-            // `as usize` won't ever truncate as the potential range is `0..=8`.
-            .map_or(0, |(i, elem)| {
-                ((i + 1) * BITS) - elem.leading_zeros() as usize
-            });
-
-        Some(K::new(last_index))
+        let index = self.bitset.pop()?;
+        Some(K::new(index))
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use alloc::vec::Vec;
     use core::u32;
 
     // `EntityRef` impl for testing.
@@ -168,7 +121,6 @@ mod tests {
         let v: Vec<E> = m.keys().collect();
         assert_eq!(v, [r0, r1, r2]);
 
-        m.resize(20);
         assert!(!m.contains(E(3)));
         assert!(!m.contains(E(4)));
         assert!(!m.contains(E(8)));
@@ -229,7 +181,7 @@ mod tests {
         for &block in &blocks {
             m.insert(block);
         }
-        assert_eq!(m.len, 13);
+        assert_eq!(m.bitset.max(), Some(12));
         blocks.sort();
 
         for &block in blocks.iter().rev() {

scripts/publish.rs

@@ -18,11 +18,11 @@ use std::time::Duration;
 // note that this list must be topologically sorted by dependencies
 const CRATES_TO_PUBLISH: &[&str] = &[
     // cranelift
+    "cranelift-bitset",
     "cranelift-isle",
     "cranelift-entity",
     "wasmtime-types",
     "cranelift-bforest",
-    "cranelift-bitset",
     "cranelift-codegen-shared",
     "cranelift-codegen-meta",
     "cranelift-egraph",

tests/all/pooling_allocator.rs

@@ -644,19 +644,19 @@ fn instance_too_large() -> Result<()> {
     let engine = Engine::new(&config)?;
     let expected = if cfg!(feature = "wmemcheck") {
         "\
-        instance allocation for this module requires 352 bytes which exceeds the \
+        instance allocation for this module requires 336 bytes which exceeds the \
 configured maximum of 16 bytes; breakdown of allocation requirement:
 
- * 72.73% - 256 bytes - instance state management
- * 25.00% - 88 bytes - static vmctx data
+ * 71.43% - 240 bytes - instance state management
+ * 26.19% - 88 bytes - static vmctx data
 "
     } else {
         "\
-instance allocation for this module requires 256 bytes which exceeds the \
+        instance allocation for this module requires 240 bytes which exceeds the \
 configured maximum of 16 bytes; breakdown of allocation requirement:
 
- * 62.50% - 160 bytes - instance state management
- * 34.38% - 88 bytes - static vmctx data
+ * 60.00% - 144 bytes - instance state management
+ * 36.67% - 88 bytes - static vmctx data
 "
     };
     match Module::new(&engine, "(module)") {
@@ -672,19 +672,19 @@ configured maximum of 16 bytes; breakdown of allocation requirement:
 
     let expected = if cfg!(feature = "wmemcheck") {
         "\
-instance allocation for this module requires 1952 bytes which exceeds the \
+instance allocation for this module requires 1936 bytes which exceeds the \
 configured maximum of 16 bytes; breakdown of allocation requirement:
 
- * 13.11% - 256 bytes - instance state management
- * 81.97% - 1600 bytes - defined globals
+ * 12.40% - 240 bytes - instance state management
+ * 82.64% - 1600 bytes - defined globals
 "
     } else {
         "\
-instance allocation for this module requires 1856 bytes which exceeds the \
+instance allocation for this module requires 1840 bytes which exceeds the \
 configured maximum of 16 bytes; breakdown of allocation requirement:
 
- * 8.62% - 160 bytes - instance state management
- * 86.21% - 1600 bytes - defined globals
+ * 7.83% - 144 bytes - instance state management
+ * 86.96% - 1600 bytes - defined globals
 "
     };
     match Module::new(&engine, &lots_of_globals) {