LazySet implementation

Author: smoelius

clippy_utils/src/mir/bit_matrix.rs

@@ -0,0 +1,57 @@
+use super::lazy_set::Set;
+use rustc_hir::def_id::DefId;
+use rustc_index::{bit_set::BitSet, vec::IndexVec};
+use rustc_middle::mir::{self, Local};
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct Param {
+    def_id: DefId,
+    domain_size: usize,
+}
+
+impl Param {
+    pub fn new(body: &mir::Body<'_>) -> Self {
+        Self {
+            def_id: body.source.def_id(),
+            domain_size: body.local_decls.len(),
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct BitMatrix(pub IndexVec<Local, Option<BitSet<Local>>>);
+
+impl Set for BitMatrix {
+    type Param = Param;
+    type Elem = (Local, Local);
+
+    fn batch_id(param: &Self::Param) -> DefId {
+        param.def_id
+    }
+
+    fn tree_size_max(param: &Self::Param) -> usize {
+        param.domain_size * param.domain_size
+    }
+
+    fn new(param: &Self::Param) -> Self {
+        Self(IndexVec::from_elem_n(None, param.domain_size))
+    }
+
+    fn insert(&mut self, elem: (Local, Local), param: &Self::Param) -> bool {
+        self.0[elem.0]
+            .get_or_insert_with(|| BitSet::new_empty(param.domain_size))
+            .insert(elem.1)
+    }
+
+    fn union(&mut self, other: &BitMatrix, param: &Self::Param) -> bool {
+        let mut changed = false;
+        for (borrowed, borrowers) in other.0.iter_enumerated() {
+            if let Some(borrowers) = borrowers {
+                changed |= self.0[borrowed]
+                    .get_or_insert_with(|| BitSet::new_empty(param.domain_size))
+                    .union(borrowers);
+            }
+        }
+        changed
+    }
+}

clippy_utils/src/mir/lazy_set.rs

@@ -0,0 +1,269 @@
+//! "Lazy" set-like data structure
+//!
+//! This data structure supports an in-place union operation similar to `BitSet`'s [`union`].
+//! However, unlike `BitSet`'s `union`, which returns true only when new bits were added, this data
+//! structure's `union` will sometimes delay computation and return `true` immediately even when no
+//! new elements were added. The time savings are paid for with inaccuracy in the return value.
+//!
+//! Delayed computations are represented in a tree-like data structure, and each such tree belongs
+//! to a "batch." When a tree's number of nodes reaches a predetermined maximum, the tree is
+//! "flattened" into an instance of the underlying set type. When this condition occurs for one tree
+//! within a batch, the other trees in the batch are similarly flattened as needed.
+//!
+//! When the underlying set is a `BitSet` (or is `BitSet`-like), it is expected that the maximum
+//! tree size is the domain size. The intuition is: the cost of maintaining a tree of this size is
+//! proportional to the cost of maintaining a set instance. So when a batch's trees reach this size,
+//! delaying computations is likely not helping us. We should therefore cut our losses and do things
+//! the traditional way.
+//!
+//! Why is this data structure needed? @mwkmwkmwk exhibited a Rust program that caused
+//! `possible_borrower` to OOM/timeout ([10134]). The reasons for this were essentially the
+//! following:
+//!
+//! - Her example involved a large number of basic blocks.
+//! - A program's number of locals tends to be proportional to its number of basic blocks (as was
+//!   the case in her example).
+//! - The work done by `possible_borrower` at each `join` in [`iterate_to_fixpoint`] (essentially,
+//!   at each basic block) was proportional to the number of locals.
+//!
+//! One way out of this situation is to make the work done at each `join` not proportional to the
+//! number of locals. The current data structure achieves this for certain usage patterns, e.g.,
+//! certain MIR dataflow problems. In particular, it avoids an OOM/timeout on @mwkmwkmwk's example
+//! and does not seem to degrade performance on Clippy's existing tests. But it is not sub-linear in
+//! the number of locals in a complexity-theoretic sense. At present, it is not clear whether such a
+//! solution exists.
+//!
+//! [10134]: https://github.com/rust-lang/rust-clippy/issues/10134
+//! [`iterate_to_fixpoint`]: https://doc.rust-lang.org/nightly/nightly-rustc/rustc_mir_dataflow/framework/engine/struct.Engine.html#method.iterate_to_fixpoint
+//! [`union`]: https://doc.rust-lang.org/nightly/nightly-rustc/rustc_index/bit_set/struct.BitSet.html#method.union
+
+use rustc_data_structures::{fx::FxHashSet, sync::Lrc};
+use rustc_hir::def_id::DefId;
+use std::borrow::Cow;
+use std::fmt::Debug;
+use std::sync::{LazyLock, RwLock};
+
+static FLATTENING: LazyLock<RwLock<FxHashSet<DefId>>> = LazyLock::new(|| RwLock::new(FxHashSet::default()));
+
+/// A set-like type with `insert` and `union` operations similar to those of `BitSet`
+///
+/// [`BitSet`]: https://doc.rust-lang.org/nightly/nightly-rustc/rustc_index/bit_set/struct.BitSet.html
+pub trait Set {
+    /// A value of type `Self::Param` is passed to each trait method invocation.
+    type Param: Clone + Debug + Eq;
+
+    /// The type of the elements in the set.
+    type Elem: Copy + Debug + Eq;
+
+    /// Returns the id of the "batch" to which this set belongs. When the set corresponds to a
+    /// `mir::Body`, it is expected that the batch id is the `DefId` of the corresponding function.
+    fn batch_id(param: &Self::Param) -> DefId;
+
+    /// Returns the set's maximum tree size. Trees that exceed this size are flattened into
+    /// instances of `Self`.
+    fn tree_size_max(param: &Self::Param) -> usize;
+
+    /// Creates a new instance of `Self`.
+    fn new(param: &Self::Param) -> Self;
+
+    /// Inserts `elem`. Returns whether the set has changed.
+    fn insert(&mut self, elem: Self::Elem, param: &Self::Param) -> bool;
+
+    /// Inserts all elements from `other` and returns `true` if `self` changed (i.e., if new
+    /// elements were added).
+    fn union(&mut self, other: &Self, param: &Self::Param) -> bool;
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct LazySet<S: Set> {
+    inner: Inner<S>,
+    param: <S as Set>::Param,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+enum Inner<S: Set> {
+    Flattened(S),
+    Tree(Option<Lrc<Tree<<S as Set>::Elem>>>),
+}
+
+impl<S: Set> Drop for Inner<S> {
+    fn drop(&mut self) {
+        if let Inner::Tree(tree) = self && let Some(tree) = tree.take() {
+            let mut stack = Vec::new();
+            stack.push(tree);
+            while let Some(tree) = stack.pop() {
+                if let Ok(Tree::Union { left, right, .. }) = Lrc::try_unwrap(tree) {
+                    stack.push(left);
+                    stack.push(right);
+                }
+            }
+        }
+    }
+}
+
+impl<S> LazySet<S>
+where
+    S: Clone + Debug + Set,
+{
+    pub fn new(param: <S as Set>::Param) -> Self {
+        Self {
+            inner: Inner::Tree(None),
+            param,
+        }
+    }
+
+    /// Similar to `BitSet`'s [`insert`]. However, this `insert` returns `true` when the set changed
+    /// _or_ computation of the new set was delayed.
+    ///
+    /// [`insert`]: https://doc.rust-lang.org/nightly/nightly-rustc/rustc_index/bit_set/struct.BitSet.html#method.insert
+    pub fn insert(&mut self, elem: <S as Set>::Elem) -> bool {
+        match &mut self.inner {
+            Inner::Flattened(flattened) => flattened.insert(elem, &self.param),
+            Inner::Tree(tree) => {
+                let tree = if let Some(tree) = tree {
+                    Tree::insert(tree, elem)
+                } else {
+                    Lrc::new(Tree::Singleton(elem))
+                };
+                self.set_inner(Inner::Tree(Some(tree)));
+                true
+            },
+        }
+    }
+
+    /// Similar to `BitSet`'s [`union`]. However, this `union` returns `true` when new elements were
+    /// added _or_  computation of the new set was delayed.
+    ///
+    /// [`union`]: https://doc.rust-lang.org/nightly/nightly-rustc/rustc_index/bit_set/struct.BitSet.html#method.union
+    pub fn union(&mut self, other: &Self) -> bool {
+        if matches!((&self.inner, &other.inner), (Inner::Tree(Some(_)), Inner::Flattened(_))) {
+            self.flatten_in_place();
+        }
+
+        match &mut self.inner {
+            Inner::Flattened(our_flattened) => match &other.inner {
+                Inner::Flattened(their_flattened) => our_flattened.union(their_flattened, &self.param),
+                Inner::Tree(None) => false,
+                Inner::Tree(Some(their_tree)) => {
+                    let mut changed = false;
+                    their_tree.for_each(&mut |elem: &<S as Set>::Elem| {
+                        changed |= our_flattened.insert(*elem, &self.param);
+                    });
+                    changed
+                },
+            },
+            Inner::Tree(None) => match &other.inner {
+                Inner::Flattened(_) | Inner::Tree(Some(_)) => {
+                    self.set_inner(other.inner.clone());
+                    true
+                },
+                Inner::Tree(None) => false,
+            },
+            Inner::Tree(Some(our_tree)) => match &other.inner {
+                Inner::Flattened(_) => unreachable!(),
+                Inner::Tree(None) => false,
+                Inner::Tree(Some(their_tree)) => {
+                    let (tree, changed) = Tree::union(our_tree, their_tree);
+                    self.set_inner(Inner::Tree(Some(tree)));
+                    changed
+                },
+            },
+        }
+    }
+
+    fn set_inner(&mut self, inner: Inner<S>) {
+        self.inner = inner;
+        if let Inner::Tree(Some(tree)) = &self.inner {
+            if tree.size() <= S::tree_size_max(&self.param) {
+                let lock = FLATTENING.read().unwrap();
+                if !lock.contains(&S::batch_id(&self.param)) {
+                    return;
+                }
+                drop(lock);
+                self.flatten_in_place();
+            } else {
+                let mut lock = FLATTENING.write().unwrap();
+                lock.insert(S::batch_id(&self.param));
+                drop(lock);
+                self.flatten_in_place();
+            }
+        }
+    }
+
+    fn flatten_in_place(&mut self) {
+        if !matches!(self.inner, Inner::Flattened(_)) {
+            self.inner = Inner::Flattened(self.flatten().into_owned());
+        }
+    }
+
+    pub fn flatten(&self) -> Cow<'_, S> {
+        match &self.inner {
+            Inner::Flattened(flattened) => Cow::Borrowed(flattened),
+            Inner::Tree(tree) => {
+                let mut flattened = S::new(&self.param);
+                if let Some(tree) = tree {
+                    tree.for_each(&mut |elem: &<S as Set>::Elem| {
+                        flattened.insert(*elem, &self.param);
+                    });
+                }
+                Cow::Owned(flattened)
+            },
+        }
+    }
+}
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+enum Tree<T> {
+    Singleton(T),
+    Union {
+        size: usize,
+        left: Lrc<Tree<T>>,
+        right: Lrc<Tree<T>>,
+    },
+}
+
+impl<T> Tree<T> {
+    fn size(&self) -> usize {
+        match self {
+            Self::Singleton(_) => 1,
+            Self::Union { size, .. } => *size,
+        }
+    }
+
+    fn insert(this: &Lrc<Tree<T>>, elem: T) -> Lrc<Self> {
+        Lrc::new(Self::Union {
+            size: 1 + this.size() + 1,
+            left: this.clone(),
+            right: Lrc::new(Self::Singleton(elem)),
+        })
+    }
+
+    fn union(this: &Lrc<Tree<T>>, other: &Lrc<Tree<T>>) -> (Lrc<Self>, bool) {
+        if Lrc::ptr_eq(this, other) {
+            (this.clone(), false)
+        } else {
+            (
+                Lrc::new(Tree::Union {
+                    size: 1 + this.size() + other.size(),
+                    left: this.clone(),
+                    right: other.clone(),
+                }),
+                true,
+            )
+        }
+    }
+
+    fn for_each(&self, f: &mut impl FnMut(&T)) {
+        let mut stack = Vec::new();
+        stack.push(self);
+        while let Some(tree) = stack.pop() {
+            match tree {
+                Self::Singleton(elem) => f(elem),
+                Self::Union { left, right, .. } => {
+                    stack.push(left);
+                    stack.push(right);
+                },
+            }
+        }
+    }
+}

clippy_utils/src/mir/mod.rs

@@ -5,6 +5,10 @@ use rustc_middle::mir::{
 };
 use rustc_middle::ty::TyCtxt;
 
+mod bit_matrix;
+
+mod lazy_set;
+
 mod possible_borrower;
 pub use possible_borrower::PossibleBorrowerMap;