[ty] Garbage-collect reachability constraints

Cargo.toml

@@ -57,6 +57,9 @@ assert_fs = { version = "1.1.0" }
 argfile = { version = "0.2.0" }
 bincode = { version = "2.0.0" }
 bitflags = { version = "2.5.0" }
+bitvec = { version = "1.0.1", default-features = false, features = [
+    "alloc",
+] }
 bstr = { version = "1.9.1" }
 cachedir = { version = "0.3.1" }
 camino = { version = "1.1.7" }

crates/ty_python_semantic/Cargo.toml

@@ -26,6 +26,7 @@ ty_static = { workspace = true }
 
 anyhow = { workspace = true }
 bitflags = { workspace = true }
+bitvec = { workspace = true }
 camino = { workspace = true }
 colored = { workspace = true }
 compact_str = { workspace = true }

crates/ty_python_semantic/src/lib.rs

@@ -34,6 +34,7 @@ mod node_key;
 pub(crate) mod place;
 mod program;
 mod python_platform;
+mod rank;
 pub mod semantic_index;
 mod semantic_model;
 pub(crate) mod site_packages;

crates/ty_python_semantic/src/rank.rs

@@ -0,0 +1,83 @@
+//! A boxed bit slice that supports a constant-time `rank` operation.
+
+use bitvec::prelude::{BitBox, Msb0};
+use get_size2::GetSize;
+
+/// A boxed bit slice that supports a constant-time `rank` operation.
+///
+/// This can be used to "shrink" a large vector, where you only need to keep certain elements, and
+/// you want to continue to use the index in the large vector to identify each element.
+///
+/// First you create a new smaller vector, keeping only the elements of the large vector that you
+/// care about. Now you need a way to translate an index into the large vector (which no longer
+/// exists) into the corresponding index into the smaller vector. To do that, you create a bit
+/// slice, containing a bit for every element of the original large vector. Each bit in the bit
+/// slice indicates whether that element of the large vector was kept in the smaller vector. And
+/// the `rank` of the bit gives us the index of the element in the smaller vector.
+///
+/// However, the naive implementation of `rank` is O(n) in the size of the bit slice. To address
+/// that, we use a standard trick: we divide the bit slice into 64-bit chunks, and when
+/// constructing the bit slice, precalculate the rank of the first bit in each chunk. Then, to
+/// calculate the rank of an arbitrary bit, we first grab the precalculated rank of the chunk that
+/// bit belongs to, and add the rank of the bit within its (fixed-sized) chunk.
+///
+/// This trick adds O(1.5) bits of overhead per large vector element on 64-bit platforms, and O(2)
+/// bits of overhead on 32-bit platforms.
+#[derive(Clone, Debug, Eq, PartialEq, GetSize)]
+pub(crate) struct RankBitBox {
+    #[get_size(size_fn = bit_box_size)]
+    bits: BitBox<Chunk, Msb0>,
+    chunk_ranks: Box<[u32]>,
+}
+
+fn bit_box_size(bits: &BitBox<Chunk, Msb0>) -> usize {
+    bits.as_raw_slice().get_heap_size()
+}
+
+// bitvec does not support `u64` as a Store type on 32-bit platforms
+#[cfg(target_pointer_width = "64")]
+type Chunk = u64;
+#[cfg(not(target_pointer_width = "64"))]
+type Chunk = u32;
+
+const CHUNK_SIZE: usize = Chunk::BITS as usize;
+
+impl RankBitBox {
+    pub(crate) fn from_bits(iter: impl Iterator<Item = bool>) -> Self {
+        let bits: BitBox<Chunk, Msb0> = iter.collect();
+        let chunk_ranks = bits
+            .as_raw_slice()
+            .iter()
+            .scan(0u32, |rank, chunk| {
+                let result = *rank;
+                *rank += chunk.count_ones();
+                Some(result)
+            })
+            .collect();
+        Self { bits, chunk_ranks }
+    }
+
+    #[inline]
+    pub(crate) fn get_bit(&self, index: usize) -> Option<bool> {
+        self.bits.get(index).map(|bit| *bit)
+    }
+
+    /// Returns the number of bits _before_ (and not including) the given index that are set.
+    #[inline]
+    pub(crate) fn rank(&self, index: usize) -> u32 {
+        let chunk_index = index / CHUNK_SIZE;
+        let index_within_chunk = index % CHUNK_SIZE;
+        let chunk_rank = self.chunk_ranks[chunk_index];
+        if index_within_chunk == 0 {
+            return chunk_rank;
+        }
+
+        // To calculate the rank within the bit's chunk, we zero out the requested bit and every
+        // bit to the right, then count the number of 1s remaining (i.e., to the left of the
+        // requested bit).
+        let chunk = self.bits.as_raw_slice()[chunk_index];
+        let chunk_mask = Chunk::MAX << (CHUNK_SIZE - index_within_chunk);
+        let rank_within_chunk = (chunk & chunk_mask).count_ones();
+        chunk_rank + rank_within_chunk
+    }
+}

crates/ty_python_semantic/src/semantic_index/builder.rs

@@ -1021,6 +1021,14 @@ impl<'db, 'ast> SemanticIndexBuilder<'db, 'ast> {
 
         assert_eq!(&self.current_assignments, &[]);
 
+        for scope in &self.scopes {
+            if let Some(parent) = scope.parent() {
+                self.use_def_maps[parent]
+                    .reachability_constraints
+                    .mark_used(scope.reachability());
+            }
+        }
+
         let mut place_tables: IndexVec<_, _> = self
             .place_tables
             .into_iter()

crates/ty_python_semantic/src/semantic_index/reachability_constraints.rs

@@ -201,6 +201,7 @@ use rustc_hash::FxHashMap;
 use crate::Db;
 use crate::dunder_all::dunder_all_names;
 use crate::place::{RequiresExplicitReExport, imported_symbol};
+use crate::rank::RankBitBox;
 use crate::semantic_index::expression::Expression;
 use crate::semantic_index::place_table;
 use crate::semantic_index::predicate::{
@@ -283,6 +284,10 @@ impl ScopedReachabilityConstraintId {
     fn is_terminal(self) -> bool {
         self.0 >= SMALLEST_TERMINAL.0
     }
+
+    fn as_u32(self) -> u32 {
+        self.0
+    }
 }
 
 impl Idx for ScopedReachabilityConstraintId {
@@ -309,12 +314,18 @@ const SMALLEST_TERMINAL: ScopedReachabilityConstraintId = ALWAYS_FALSE;
 /// A collection of reachability constraints for a given scope.
 #[derive(Debug, PartialEq, Eq, salsa::Update, get_size2::GetSize)]
 pub(crate) struct ReachabilityConstraints {
-    interiors: IndexVec<ScopedReachabilityConstraintId, InteriorNode>,
+    /// The interior TDD nodes that were marked as used when being built.
+    used_interiors: Box<[InteriorNode]>,
+    /// A bit vector indicating which interior TDD nodes were marked as used. This is indexed by
+    /// the node's [`ScopedReachabilityConstraintId`]. The rank of the corresponding bit gives the
+    /// index of that node in the `used_interiors` vector.
+    used_indices: RankBitBox,
 }
 
 #[derive(Debug, Default, PartialEq, Eq)]
 pub(crate) struct ReachabilityConstraintsBuilder {
     interiors: IndexVec<ScopedReachabilityConstraintId, InteriorNode>,
+    interior_used: IndexVec<ScopedReachabilityConstraintId, bool>,
     interior_cache: FxHashMap<InteriorNode, ScopedReachabilityConstraintId>,
     not_cache: FxHashMap<ScopedReachabilityConstraintId, ScopedReachabilityConstraintId>,
     and_cache: FxHashMap<
@@ -334,11 +345,28 @@ pub(crate) struct ReachabilityConstraintsBuilder {
 }
 
 impl ReachabilityConstraintsBuilder {
-    pub(crate) fn build(mut self) -> ReachabilityConstraints {
-        self.interiors.shrink_to_fit();
-
+    pub(crate) fn build(self) -> ReachabilityConstraints {
+        let used_indices = RankBitBox::from_bits(self.interior_used.iter().copied());
+        let used_interiors = (self.interiors.into_iter())
+            .zip(self.interior_used)
+            .filter_map(|(interior, used)| used.then_some(interior))
+            .collect();
         ReachabilityConstraints {
-            interiors: self.interiors,
+            used_interiors,
+            used_indices,
+        }
+    }
+
+    /// Marks that a particular TDD node is used. This lets us throw away interior nodes that were
+    /// only calculated for intermediate values, and which don't need to be included in the final
+    /// built result.
+    pub(crate) fn mark_used(&mut self, node: ScopedReachabilityConstraintId) {
+        if !node.is_terminal() && !self.interior_used[node] {
+            self.interior_used[node] = true;
+            let node = self.interiors[node];
+            self.mark_used(node.if_true);
+            self.mark_used(node.if_ambiguous);
+            self.mark_used(node.if_false);
         }
     }
 
@@ -370,10 +398,10 @@ impl ReachabilityConstraintsBuilder {
             return node.if_true;
         }
 
-        *self
-            .interior_cache
-            .entry(node)
-            .or_insert_with(|| self.interiors.push(node))
+        *self.interior_cache.entry(node).or_insert_with(|| {
+            self.interior_used.push(false);
+            self.interiors.push(node)
+        })
     }
 
     /// Adds a new reachability constraint that checks a single [`Predicate`].
@@ -581,7 +609,21 @@ impl ReachabilityConstraints {
                 ALWAYS_TRUE => return Truthiness::AlwaysTrue,
                 AMBIGUOUS => return Truthiness::Ambiguous,
                 ALWAYS_FALSE => return Truthiness::AlwaysFalse,
-                _ => self.interiors[id],
+                _ => {
+                    // `id` gives us the index of this node in the IndexVec that we used when
+                    // constructing this BDD. When finalizing the builder, we threw away any
+                    // interior nodes that weren't marked as used. The `used_indices` bit vector
+                    // lets us verify that this node was marked as used, and the rank of that bit
+                    // in the bit vector tells us where this node lives in the "condensed"
+                    // `used_interiors` vector.
+                    let raw_index = id.as_u32() as usize;
+                    debug_assert!(
+                        self.used_indices.get_bit(raw_index).unwrap_or(false),
+                        "all used reachability constraints should have been marked as used",
+                    );
+                    let index = self.used_indices.rank(raw_index) as usize;
+                    self.used_interiors[index]
+                }
             };
             let predicate = &predicates[node.atom];
             match Self::analyze_single(db, predicate) {

crates/ty_python_semantic/src/semantic_index/use_def.rs

@@ -1118,7 +1118,41 @@ impl<'db> UseDefMapBuilder<'db> {
             .add_or_constraint(self.reachability, snapshot.reachability);
     }
 
+    fn mark_reachability_constraints(&mut self) {
+        // We only walk the fields that are copied through to the UseDefMap when we finish building
+        // it.
+        for bindings in &mut self.bindings_by_use {
+            bindings.finish(&mut self.reachability_constraints);
+        }
+        for constraint in self.node_reachability.values() {
+            self.reachability_constraints.mark_used(*constraint);
+        }
+        for place_state in &mut self.place_states {
+            place_state.finish(&mut self.reachability_constraints);
+        }
+        for reachable_definition in &mut self.reachable_definitions {
+            reachable_definition
+                .bindings
+                .finish(&mut self.reachability_constraints);
+            reachable_definition
+                .declarations
+                .finish(&mut self.reachability_constraints);
+        }
+        for declarations in self.declarations_by_binding.values_mut() {
+            declarations.finish(&mut self.reachability_constraints);
+        }
+        for bindings in self.bindings_by_definition.values_mut() {
+            bindings.finish(&mut self.reachability_constraints);
+        }
+        for eager_snapshot in &mut self.eager_snapshots {
+            eager_snapshot.finish(&mut self.reachability_constraints);
+        }
+        self.reachability_constraints.mark_used(self.reachability);
+    }
+
     pub(super) fn finish(mut self) -> UseDefMap<'db> {
+        self.mark_reachability_constraints();
+
         self.all_definitions.shrink_to_fit();
         self.place_states.shrink_to_fit();
         self.reachable_definitions.shrink_to_fit();