[ty] Avoid storing redundant reachability indexes (#25453)

charliermarsh · web-flow · commit 5a6c811505e5 · 2026-05-29T23:01:39.000+02:00
## Summary

When all reachability nodes for a scope are retained, their IDs already
match their positions in the stored node array. In that case, this PR
avoids storing an extra bit vector for mapping node IDs back to the
array. We only allocate that mapping when unused intermediate nodes were
removed.

This reduces retained memory for reachability constraints while
preserving the existing compaction for scopes that need it. It also
moves node lookup into `ReachabilityConstraints::get_interior_node`, so
both paths use the same lookup logic.
diff --git a/crates/ty_python_core/src/reachability_constraints.rs b/crates/ty_python_core/src/reachability_constraints.rs
@@ -144,29 +144,31 @@ pub struct ReachabilityConstraints {
     /// 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,
+    ///
+    /// If all interior nodes were retained, the original ID can be used directly instead.
+    used_indices: Option<Box<RankBitBox>>,
 }
 
 impl ReachabilityConstraints {
     /// Look up an interior node by its constraint ID.
     pub fn get_interior_node(&self, id: ScopedReachabilityConstraintId) -> InteriorNode {
         debug_assert!(!id.is_terminal());
         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]
+        if let Some(used_indices) = &self.used_indices {
+            debug_assert!(
+                used_indices.get_bit(raw_index).unwrap_or(false),
+                "all used reachability constraints should have been marked as used",
+            );
+            let index = used_indices.rank(raw_index) as usize;
+            self.used_interiors[index]
+        } else {
+            self.used_interiors[raw_index]
+        }
     }
 
     pub fn used_interiors(&self) -> &[InteriorNode] {
         &self.used_interiors
     }
-
-    pub fn used_indices(&self) -> &RankBitBox {
-        &self.used_indices
-    }
 }
 
 #[derive(Debug, Default, PartialEq, Eq)]
@@ -193,14 +195,21 @@ pub struct ReachabilityConstraintsBuilder {
 
 impl ReachabilityConstraintsBuilder {
     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 {
-            used_interiors,
-            used_indices,
+        if self.interior_used.iter().all(|used| *used) {
+            ReachabilityConstraints {
+                used_interiors: self.interiors.raw.into_boxed_slice(),
+                used_indices: None,
+            }
+        } else {
+            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 {
+                used_interiors,
+                used_indices: Some(Box::new(used_indices)),
+            }
         }
     }
 
diff --git a/crates/ty_python_semantic/src/reachability.rs b/crates/ty_python_semantic/src/reachability.rs
@@ -629,21 +629,7 @@ impl<'db> ReachabilityConstraintsExtension<'db> for ReachabilityConstraints {
                 Id::ALWAYS_TRUE => return Truthiness::AlwaysTrue,
                 Id::AMBIGUOUS => return Truthiness::Ambiguous,
                 Id::ALWAYS_FALSE => return Truthiness::AlwaysFalse,
-                _ => {
-                    // `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]
-                }
+                _ => self.get_interior_node(id),
             };
             let predicate = &predicates[node.atom()];
             match analyze_single(db, predicate) {
