summaryrefslogtreecommitdiffstats
path: root/third_party/rust/cranelift-bforest/src/set.rs
diff options
context:
space:
mode:
Diffstat (limited to 'third_party/rust/cranelift-bforest/src/set.rs')
-rw-r--r--third_party/rust/cranelift-bforest/src/set.rs598
1 files changed, 598 insertions, 0 deletions
diff --git a/third_party/rust/cranelift-bforest/src/set.rs b/third_party/rust/cranelift-bforest/src/set.rs
new file mode 100644
index 0000000000..e7761a63d9
--- /dev/null
+++ b/third_party/rust/cranelift-bforest/src/set.rs
@@ -0,0 +1,598 @@
+//! Forest of sets.
+
+use super::{Comparator, Forest, Node, NodeData, NodePool, Path, SetValue, INNER_SIZE};
+use crate::packed_option::PackedOption;
+#[cfg(test)]
+use alloc::string::String;
+#[cfg(test)]
+use core::fmt;
+use core::marker::PhantomData;
+
+/// Tag type defining forest types for a set.
+struct SetTypes<K>(PhantomData<K>);
+
+impl<K> Forest for SetTypes<K>
+where
+ K: Copy,
+{
+ type Key = K;
+ type Value = SetValue;
+ type LeafKeys = [K; 2 * INNER_SIZE - 1];
+ type LeafValues = [SetValue; 2 * INNER_SIZE - 1];
+
+ fn splat_key(key: Self::Key) -> Self::LeafKeys {
+ [key; 2 * INNER_SIZE - 1]
+ }
+
+ fn splat_value(value: Self::Value) -> Self::LeafValues {
+ [value; 2 * INNER_SIZE - 1]
+ }
+}
+
+/// Memory pool for a forest of `Set` instances.
+pub struct SetForest<K>
+where
+ K: Copy,
+{
+ nodes: NodePool<SetTypes<K>>,
+}
+
+impl<K> SetForest<K>
+where
+ K: Copy,
+{
+ /// Create a new empty forest.
+ pub fn new() -> Self {
+ Self {
+ nodes: NodePool::new(),
+ }
+ }
+
+ /// Clear all sets in the forest.
+ ///
+ /// All `Set` instances belong to this forest are invalidated and should no longer be used.
+ pub fn clear(&mut self) {
+ self.nodes.clear();
+ }
+}
+
+/// B-tree representing an ordered set of `K`s using `C` for comparing elements.
+///
+/// This is not a general-purpose replacement for `BTreeSet`. See the [module
+/// documentation](index.html) for more information about design tradeoffs.
+///
+/// Sets can be cloned, but that operation should only be used as part of cloning the whole forest
+/// they belong to. *Cloning a set does not allocate new memory for the clone*. It creates an alias
+/// of the same memory.
+#[derive(Clone)]
+pub struct Set<K>
+where
+ K: Copy,
+{
+ root: PackedOption<Node>,
+ unused: PhantomData<K>,
+}
+
+impl<K> Set<K>
+where
+ K: Copy,
+{
+ /// Make an empty set.
+ pub fn new() -> Self {
+ Self {
+ root: None.into(),
+ unused: PhantomData,
+ }
+ }
+
+ /// Is this an empty set?
+ pub fn is_empty(&self) -> bool {
+ self.root.is_none()
+ }
+
+ /// Does the set contain `key`?.
+ pub fn contains<C: Comparator<K>>(&self, key: K, forest: &SetForest<K>, comp: &C) -> bool {
+ self.root
+ .expand()
+ .and_then(|root| Path::default().find(key, root, &forest.nodes, comp))
+ .is_some()
+ }
+
+ /// Try to insert `key` into the set.
+ ///
+ /// If the set did not contain `key`, insert it and return true.
+ ///
+ /// If `key` is already present, don't change the set and return false.
+ pub fn insert<C: Comparator<K>>(
+ &mut self,
+ key: K,
+ forest: &mut SetForest<K>,
+ comp: &C,
+ ) -> bool {
+ self.cursor(forest, comp).insert(key)
+ }
+
+ /// Remove `key` from the set and return true.
+ ///
+ /// If `key` was not present in the set, return false.
+ pub fn remove<C: Comparator<K>>(
+ &mut self,
+ key: K,
+ forest: &mut SetForest<K>,
+ comp: &C,
+ ) -> bool {
+ let mut c = self.cursor(forest, comp);
+ if c.goto(key) {
+ c.remove();
+ true
+ } else {
+ false
+ }
+ }
+
+ /// Remove all entries.
+ pub fn clear(&mut self, forest: &mut SetForest<K>) {
+ if let Some(root) = self.root.take() {
+ forest.nodes.free_tree(root);
+ }
+ }
+
+ /// Retains only the elements specified by the predicate.
+ ///
+ /// Remove all elements where the predicate returns false.
+ pub fn retain<F>(&mut self, forest: &mut SetForest<K>, mut predicate: F)
+ where
+ F: FnMut(K) -> bool,
+ {
+ let mut path = Path::default();
+ if let Some(root) = self.root.expand() {
+ path.first(root, &forest.nodes);
+ }
+ while let Some((node, entry)) = path.leaf_pos() {
+ if predicate(forest.nodes[node].unwrap_leaf().0[entry]) {
+ path.next(&forest.nodes);
+ } else {
+ self.root = path.remove(&mut forest.nodes).into();
+ }
+ }
+ }
+
+ /// Create a cursor for navigating this set. The cursor is initially positioned off the end of
+ /// the set.
+ pub fn cursor<'a, C: Comparator<K>>(
+ &'a mut self,
+ forest: &'a mut SetForest<K>,
+ comp: &'a C,
+ ) -> SetCursor<'a, K, C> {
+ SetCursor::new(self, forest, comp)
+ }
+
+ /// Create an iterator traversing this set. The iterator type is `K`.
+ pub fn iter<'a>(&'a self, forest: &'a SetForest<K>) -> SetIter<'a, K> {
+ SetIter {
+ root: self.root,
+ pool: &forest.nodes,
+ path: Path::default(),
+ }
+ }
+}
+
+impl<K> Default for Set<K>
+where
+ K: Copy,
+{
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+/// A position in a `Set` used to navigate and modify the ordered set.
+///
+/// A cursor always points at an element in the set, or "off the end" which is a position after the
+/// last element in the set.
+pub struct SetCursor<'a, K, C>
+where
+ K: 'a + Copy,
+ C: 'a + Comparator<K>,
+{
+ root: &'a mut PackedOption<Node>,
+ pool: &'a mut NodePool<SetTypes<K>>,
+ comp: &'a C,
+ path: Path<SetTypes<K>>,
+}
+
+impl<'a, K, C> SetCursor<'a, K, C>
+where
+ K: Copy,
+ C: Comparator<K>,
+{
+ /// Create a cursor with a default (invalid) location.
+ fn new(container: &'a mut Set<K>, forest: &'a mut SetForest<K>, comp: &'a C) -> Self {
+ Self {
+ root: &mut container.root,
+ pool: &mut forest.nodes,
+ comp,
+ path: Path::default(),
+ }
+ }
+
+ /// Is this cursor pointing to an empty set?
+ pub fn is_empty(&self) -> bool {
+ self.root.is_none()
+ }
+
+ /// Move cursor to the next element and return it.
+ ///
+ /// If the cursor reaches the end, return `None` and leave the cursor at the off-the-end
+ /// position.
+ #[cfg_attr(feature = "cargo-clippy", allow(clippy::should_implement_trait))]
+ pub fn next(&mut self) -> Option<K> {
+ self.path.next(self.pool).map(|(k, _)| k)
+ }
+
+ /// Move cursor to the previous element and return it.
+ ///
+ /// If the cursor is already pointing at the first element, leave it there and return `None`.
+ pub fn prev(&mut self) -> Option<K> {
+ self.root
+ .expand()
+ .and_then(|root| self.path.prev(root, self.pool).map(|(k, _)| k))
+ }
+
+ /// Get the current element, or `None` if the cursor is at the end.
+ pub fn elem(&self) -> Option<K> {
+ self.path
+ .leaf_pos()
+ .and_then(|(node, entry)| self.pool[node].unwrap_leaf().0.get(entry).cloned())
+ }
+
+ /// Move this cursor to `elem`.
+ ///
+ /// If `elem` is in the set, place the cursor at `elem` and return true.
+ ///
+ /// If `elem` is not in the set, place the cursor at the next larger element (or the end) and
+ /// return false.
+ pub fn goto(&mut self, elem: K) -> bool {
+ match self.root.expand() {
+ None => false,
+ Some(root) => {
+ if self.path.find(elem, root, self.pool, self.comp).is_some() {
+ true
+ } else {
+ self.path.normalize(self.pool);
+ false
+ }
+ }
+ }
+ }
+
+ /// Move this cursor to the first element.
+ pub fn goto_first(&mut self) -> Option<K> {
+ self.root.map(|root| self.path.first(root, self.pool).0)
+ }
+
+ /// Try to insert `elem` into the set and leave the cursor at the inserted element.
+ ///
+ /// If the set did not contain `elem`, insert it and return true.
+ ///
+ /// If `elem` is already present, don't change the set, place the cursor at `goto(elem)`, and
+ /// return false.
+ pub fn insert(&mut self, elem: K) -> bool {
+ match self.root.expand() {
+ None => {
+ let root = self.pool.alloc_node(NodeData::leaf(elem, SetValue()));
+ *self.root = root.into();
+ self.path.set_root_node(root);
+ true
+ }
+ Some(root) => {
+ // TODO: Optimize the case where `self.path` is already at the correct insert pos.
+ if self.path.find(elem, root, self.pool, self.comp).is_none() {
+ *self.root = self.path.insert(elem, SetValue(), self.pool).into();
+ true
+ } else {
+ false
+ }
+ }
+ }
+ }
+
+ /// Remove the current element (if any) and return it.
+ /// This advances the cursor to the next element after the removed one.
+ pub fn remove(&mut self) -> Option<K> {
+ let elem = self.elem();
+ if elem.is_some() {
+ *self.root = self.path.remove(self.pool).into();
+ }
+ elem
+ }
+}
+
+#[cfg(test)]
+impl<'a, K, C> SetCursor<'a, K, C>
+where
+ K: Copy + fmt::Display,
+ C: Comparator<K>,
+{
+ fn verify(&self) {
+ self.path.verify(self.pool);
+ self.root.map(|root| self.pool.verify_tree(root, self.comp));
+ }
+
+ /// Get a text version of the path to the current position.
+ fn tpath(&self) -> String {
+ use alloc::string::ToString;
+ self.path.to_string()
+ }
+}
+
+/// An iterator visiting the elements of a `Set`.
+pub struct SetIter<'a, K>
+where
+ K: 'a + Copy,
+{
+ root: PackedOption<Node>,
+ pool: &'a NodePool<SetTypes<K>>,
+ path: Path<SetTypes<K>>,
+}
+
+impl<'a, K> Iterator for SetIter<'a, K>
+where
+ K: 'a + Copy,
+{
+ type Item = K;
+
+ fn next(&mut self) -> Option<Self::Item> {
+ // We use `self.root` to indicate if we need to go to the first element. Reset to `None`
+ // once we've returned the first element. This also works for an empty tree since the
+ // `path.next()` call returns `None` when the path is empty. This also fuses the iterator.
+ match self.root.take() {
+ Some(root) => Some(self.path.first(root, self.pool).0),
+ None => self.path.next(self.pool).map(|(k, _)| k),
+ }
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::super::NodeData;
+ use super::*;
+ use alloc::vec::Vec;
+ use core::mem;
+
+ #[test]
+ fn node_size() {
+ // check that nodes are cache line sized when keys are 32 bits.
+ type F = SetTypes<u32>;
+ assert_eq!(mem::size_of::<NodeData<F>>(), 64);
+ }
+
+ #[test]
+ fn empty() {
+ let mut f = SetForest::<u32>::new();
+ f.clear();
+
+ let mut s = Set::<u32>::new();
+ assert!(s.is_empty());
+ s.clear(&mut f);
+ assert!(!s.contains(7, &f, &()));
+
+ // Iterator for an empty set.
+ assert_eq!(s.iter(&f).next(), None);
+
+ s.retain(&mut f, |_| unreachable!());
+
+ let mut c = SetCursor::new(&mut s, &mut f, &());
+ c.verify();
+ assert_eq!(c.elem(), None);
+
+ assert_eq!(c.goto_first(), None);
+ assert_eq!(c.tpath(), "<empty path>");
+ }
+
+ #[test]
+ fn simple_cursor() {
+ let mut f = SetForest::<u32>::new();
+ let mut s = Set::<u32>::new();
+ let mut c = SetCursor::new(&mut s, &mut f, &());
+
+ assert!(c.insert(50));
+ c.verify();
+ assert_eq!(c.elem(), Some(50));
+
+ assert!(c.insert(100));
+ c.verify();
+ assert_eq!(c.elem(), Some(100));
+
+ assert!(c.insert(10));
+ c.verify();
+ assert_eq!(c.elem(), Some(10));
+
+ // Basic movement.
+ assert_eq!(c.next(), Some(50));
+ assert_eq!(c.next(), Some(100));
+ assert_eq!(c.next(), None);
+ assert_eq!(c.next(), None);
+ assert_eq!(c.prev(), Some(100));
+ assert_eq!(c.prev(), Some(50));
+ assert_eq!(c.prev(), Some(10));
+ assert_eq!(c.prev(), None);
+ assert_eq!(c.prev(), None);
+
+ assert!(c.goto(50));
+ assert_eq!(c.elem(), Some(50));
+ assert_eq!(c.remove(), Some(50));
+ c.verify();
+
+ assert_eq!(c.elem(), Some(100));
+ assert_eq!(c.remove(), Some(100));
+ c.verify();
+ assert_eq!(c.elem(), None);
+ assert_eq!(c.remove(), None);
+ c.verify();
+ }
+
+ #[test]
+ fn two_level_sparse_tree() {
+ let mut f = SetForest::<u32>::new();
+ let mut s = Set::<u32>::new();
+ let mut c = SetCursor::new(&mut s, &mut f, &());
+
+ // Insert enough elements that we get a two-level tree.
+ // Each leaf node holds 8 elements
+ assert!(c.is_empty());
+ for i in 0..50 {
+ assert!(c.insert(i));
+ assert_eq!(c.elem(), Some(i));
+ }
+ assert!(!c.is_empty());
+
+ assert_eq!(c.goto_first(), Some(0));
+ assert_eq!(c.tpath(), "node2[0]--node0[0]");
+
+ assert_eq!(c.prev(), None);
+ for i in 1..50 {
+ assert_eq!(c.next(), Some(i));
+ }
+ assert_eq!(c.next(), None);
+ for i in (0..50).rev() {
+ assert_eq!(c.prev(), Some(i));
+ }
+ assert_eq!(c.prev(), None);
+
+ assert!(c.goto(25));
+ for i in 25..50 {
+ assert_eq!(c.remove(), Some(i));
+ assert!(!c.is_empty());
+ c.verify();
+ }
+
+ for i in (0..25).rev() {
+ assert!(!c.is_empty());
+ assert_eq!(c.elem(), None);
+ assert_eq!(c.prev(), Some(i));
+ assert_eq!(c.remove(), Some(i));
+ c.verify();
+ }
+ assert_eq!(c.elem(), None);
+ assert!(c.is_empty());
+ }
+
+ #[test]
+ fn three_level_sparse_tree() {
+ let mut f = SetForest::<u32>::new();
+ let mut s = Set::<u32>::new();
+ let mut c = SetCursor::new(&mut s, &mut f, &());
+
+ // Insert enough elements that we get a 3-level tree.
+ // Each leaf node holds 8 elements when filled up sequentially.
+ // Inner nodes hold 8 node pointers.
+ assert!(c.is_empty());
+ for i in 0..150 {
+ assert!(c.insert(i));
+ assert_eq!(c.elem(), Some(i));
+ }
+ assert!(!c.is_empty());
+
+ assert!(c.goto(0));
+ assert_eq!(c.tpath(), "node11[0]--node2[0]--node0[0]");
+
+ assert_eq!(c.prev(), None);
+ for i in 1..150 {
+ assert_eq!(c.next(), Some(i));
+ }
+ assert_eq!(c.next(), None);
+ for i in (0..150).rev() {
+ assert_eq!(c.prev(), Some(i));
+ }
+ assert_eq!(c.prev(), None);
+
+ assert!(c.goto(125));
+ for i in 125..150 {
+ assert_eq!(c.remove(), Some(i));
+ assert!(!c.is_empty());
+ c.verify();
+ }
+
+ for i in (0..125).rev() {
+ assert!(!c.is_empty());
+ assert_eq!(c.elem(), None);
+ assert_eq!(c.prev(), Some(i));
+ assert_eq!(c.remove(), Some(i));
+ c.verify();
+ }
+ assert_eq!(c.elem(), None);
+ assert!(c.is_empty());
+ }
+
+ // Generate a densely populated 4-level tree.
+ //
+ // Level 1: 1 root
+ // Level 2: 8 inner
+ // Level 3: 64 inner
+ // Level 4: 512 leafs, up to 7680 elements
+ //
+ // A 3-level tree can hold at most 960 elements.
+ fn dense4l(f: &mut SetForest<i32>) -> Set<i32> {
+ f.clear();
+ let mut s = Set::new();
+
+ // Insert 400 elements in 7 passes over the range to avoid the half-full leaf node pattern
+ // that comes from sequential insertion. This will generate a normal leaf layer.
+ for n in 0..4000 {
+ assert!(s.insert((n * 7) % 4000, f, &()));
+ }
+ s
+ }
+
+ #[test]
+ fn four_level() {
+ let mut f = SetForest::<i32>::new();
+ let mut s = dense4l(&mut f);
+
+ assert_eq!(
+ s.iter(&f).collect::<Vec<_>>()[0..10],
+ [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
+ );
+
+ let mut c = s.cursor(&mut f, &());
+
+ c.verify();
+
+ // Peel off a whole sub-tree of the root by deleting from the front.
+ // The 900 element is near the front of the second sub-tree.
+ assert!(c.goto(900));
+ assert_eq!(c.tpath(), "node48[1]--node47[0]--node26[0]--node20[4]");
+ assert!(c.goto(0));
+ for i in 0..900 {
+ assert!(!c.is_empty());
+ assert_eq!(c.remove(), Some(i));
+ }
+ c.verify();
+ assert_eq!(c.elem(), Some(900));
+
+ // Delete backwards from somewhere in the middle.
+ assert!(c.goto(3000));
+ for i in (2000..3000).rev() {
+ assert_eq!(c.prev(), Some(i));
+ assert_eq!(c.remove(), Some(i));
+ assert_eq!(c.elem(), Some(3000));
+ }
+ c.verify();
+
+ // Remove everything in a scattered manner, triggering many collapsing patterns.
+ for i in 0..4000 {
+ if c.goto((i * 7) % 4000) {
+ c.remove();
+ }
+ }
+ assert!(c.is_empty());
+ }
+
+ #[test]
+ fn four_level_clear() {
+ let mut f = SetForest::<i32>::new();
+ let mut s = dense4l(&mut f);
+ s.clear(&mut f);
+ }
+}
generated by cgit v1.2.3 (git 2.39.1) at 2025-02-02 00:03:46 +0000