[red-knot] Reuse alist cells when possible

.github/workflows/daily_property_tests.yaml

@@ -47,6 +47,7 @@ jobs:
         run: |
           export QUICKCHECK_TESTS=100000
           for _ in {1..5}; do
+            cargo test --locked --release --package ruff_index -- --ignored list::property_tests
             cargo test --locked --release --package red_knot_python_semantic -- --ignored types::property_tests::stable
           done
 

Cargo.toml

@@ -117,6 +117,8 @@ pretty_assertions = "1.3.0"
 proc-macro2 = { version = "1.0.79" }
 pyproject-toml = { version = "0.13.4" }
 quick-junit = { version = "0.5.0" }
+quickcheck = { version = "1.0.3", default-features = false }
+quickcheck_macros = { version = "1.0.0" }
 quote = { version = "1.0.23" }
 rand = { version = "0.9.0" }
 rayon = { version = "1.10.0" }

crates/red_knot_python_semantic/Cargo.toml

@@ -55,8 +55,8 @@ anyhow = { workspace = true }
 dir-test = { workspace = true }
 insta = { workspace = true }
 tempfile = { workspace = true }
-quickcheck = { version = "1.0.3", default-features = false }
-quickcheck_macros = { version = "1.0.0" }
+quickcheck = { workspace = true }
+quickcheck_macros = { workspace = true }
 
 [features]
 serde = ["ruff_db/serde", "dep:serde", "ruff_python_ast/serde"]

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -22,6 +22,7 @@ use crate::semantic_index::definition::{
     WithItemDefinitionNodeRef,
 };
 use crate::semantic_index::expression::{Expression, ExpressionKind};
+use crate::semantic_index::narrowing_constraints::NarrowingConstraintsBuilder;
 use crate::semantic_index::predicate::{
     PatternPredicate, PatternPredicateKind, Predicate, PredicateNode, ScopedPredicateId,
 };
@@ -289,9 +290,9 @@ impl<'db> SemanticIndexBuilder<'db> {
         &mut self.use_def_maps[scope_id]
     }
 
-    fn current_use_def_map(&self) -> &UseDefMapBuilder<'db> {
+    fn current_narrowing_constraints_mut(&mut self) -> &mut NarrowingConstraintsBuilder {
         let scope_id = self.current_scope();
-        &self.use_def_maps[scope_id]
+        &mut self.use_def_maps[scope_id].narrowing_constraints
     }
 
     fn current_visibility_constraints_mut(&mut self) -> &mut VisibilityConstraintsBuilder {
@@ -304,14 +305,23 @@ impl<'db> SemanticIndexBuilder<'db> {
         &mut self.ast_ids[scope_id]
     }
 
-    fn flow_snapshot(&self) -> FlowSnapshot {
-        self.current_use_def_map().snapshot()
+    fn flow_clone(&mut self, state: &FlowSnapshot) -> FlowSnapshot {
+        state.clone(self.current_narrowing_constraints_mut())
+    }
+
+    fn flow_snapshot(&mut self) -> FlowSnapshot {
+        self.current_use_def_map_mut().snapshot()
     }
 
     fn flow_restore(&mut self, state: FlowSnapshot) {
         self.current_use_def_map_mut().restore(state);
     }
 
+    fn flow_restore_clone(&mut self, state: &FlowSnapshot) {
+        let state = self.flow_clone(state);
+        self.current_use_def_map_mut().restore(state);
+    }
+
     fn flow_merge(&mut self, state: FlowSnapshot) {
         self.current_use_def_map_mut().merge(state);
     }
@@ -1154,7 +1164,7 @@ where
                     post_clauses.push(self.flow_snapshot());
                     // we can only take an elif/else branch if none of the previous ones were
                     // taken
-                    self.flow_restore(no_branch_taken.clone());
+                    self.flow_restore_clone(&no_branch_taken);
                     self.record_negated_narrowing_constraint(last_predicate);
 
                     let elif_predicate = if let Some(elif_test) = clause_test {
@@ -1241,7 +1251,7 @@ where
                 // To model this correctly, we need two copies of the while condition constraint,
                 // since the first and later evaluations might produce different results.
                 let post_body = self.flow_snapshot();
-                self.flow_restore(pre_loop.clone());
+                self.flow_restore_clone(&pre_loop);
                 self.record_negated_visibility_constraint(first_vis_constraint_id);
                 self.flow_merge(post_body);
                 self.record_negated_narrowing_constraint(predicate);
@@ -1387,7 +1397,7 @@ where
                 for (i, case) in cases.iter().enumerate() {
                     if i != 0 {
                         post_case_snapshots.push(self.flow_snapshot());
-                        self.flow_restore(after_subject.clone());
+                        self.flow_restore_clone(&after_subject);
                     }
 
                     self.current_match_case = Some(CurrentMatchCase::new(&case.pattern));
@@ -1416,7 +1426,7 @@ where
                     .is_some_and(|case| case.guard.is_none() && case.pattern.is_wildcard())
                 {
                     post_case_snapshots.push(self.flow_snapshot());
-                    self.flow_restore(after_subject.clone());
+                    self.flow_restore_clone(&after_subject);
 
                     for id in &vis_constraints {
                         self.record_negated_visibility_constraint(*id);
@@ -1512,7 +1522,7 @@ where
                         // as we'll immediately call `self.flow_restore()` to a different state
                         // as soon as this loop over the handlers terminates.
                         if i < (num_handlers - 1) {
-                            self.flow_restore(pre_except_state.clone());
+                            self.flow_restore_clone(&pre_except_state);
                         }
                     }
 
@@ -1548,7 +1558,8 @@ where
 
             ast::Stmt::Break(_) => {
                 if self.loop_state().is_inside() {
-                    self.loop_break_states.push(self.flow_snapshot());
+                    let snapshot = self.flow_snapshot();
+                    self.loop_break_states.push(snapshot);
                 }
                 // Everything in the current block after a terminal statement is unreachable.
                 self.mark_unreachable();
@@ -1708,7 +1719,7 @@ where
                 self.visit_expr(body);
                 let visibility_constraint = self.record_visibility_constraint(predicate);
                 let post_body = self.flow_snapshot();
-                self.flow_restore(pre_if.clone());
+                self.flow_restore_clone(&pre_if);
 
                 self.record_negated_narrowing_constraint(predicate);
                 self.visit_expr(orelse);

crates/red_knot_python_semantic/src/semantic_index/builder/except_handlers.rs

@@ -41,7 +41,7 @@ impl TryNodeContextStackManager {
 
     /// Retrieve the stack that is at the top of our stack of stacks.
     /// For each `try` block on that stack, push the snapshot onto the `try` block
-    pub(super) fn record_definition(&mut self, builder: &SemanticIndexBuilder) {
+    pub(super) fn record_definition(&mut self, builder: &mut SemanticIndexBuilder) {
         self.current_try_context_stack().record_definition(builder);
     }
 
@@ -77,7 +77,7 @@ impl TryNodeContextStack {
     }
 
     /// For each `try` block on the stack, push the snapshot onto the `try` block
-    fn record_definition(&mut self, builder: &SemanticIndexBuilder) {
+    fn record_definition(&mut self, builder: &mut SemanticIndexBuilder) {
         for context in &mut self.0 {
             context.record_definition(builder.flow_snapshot());
         }

crates/red_knot_python_semantic/src/semantic_index/narrowing_constraints.rs

@@ -28,21 +28,16 @@
 //!
 //! [`Predicate`]: crate::semantic_index::predicate::Predicate
 
-use ruff_index::list::{ListBuilder, ListSetReverseIterator, ListStorage};
-use ruff_index::newtype_index;
+use ruff_index::list::{List, ListBuilder, ListSetReverseIterator, ListStorage};
 
 use crate::semantic_index::predicate::ScopedPredicateId;
 
 /// A narrowing constraint associated with a live binding.
 ///
 /// A constraint is a list of [`Predicate`]s that each constrain the type of the binding's symbol.
 ///
-/// An instance of this type represents a _non-empty_ narrowing constraint. You will often wrap
-/// this in `Option` and use `None` to represent an empty narrowing constraint.
-///
 /// [`Predicate`]: crate::semantic_index::predicate::Predicate
-#[newtype_index]
-pub(crate) struct ScopedNarrowingConstraintId;
+pub(crate) type ScopedNarrowingConstraint = List<ScopedNarrowingConstraintPredicate>;
 
 /// One of the [`Predicate`]s in a narrowing constraint, which constraints the type of the
 /// binding's symbol.
@@ -71,7 +66,7 @@ impl From<ScopedPredicateId> for ScopedNarrowingConstraintPredicate {
 /// A collection of narrowing constraints for a given scope.
 #[derive(Debug, Eq, PartialEq)]
 pub(crate) struct NarrowingConstraints {
-    lists: ListStorage<ScopedNarrowingConstraintId, ScopedNarrowingConstraintPredicate>,
+    lists: ListStorage<ScopedNarrowingConstraintPredicate>,
 }
 
 // Building constraints
@@ -80,7 +75,7 @@ pub(crate) struct NarrowingConstraints {
 /// A builder for creating narrowing constraints.
 #[derive(Debug, Default, Eq, PartialEq)]
 pub(crate) struct NarrowingConstraintsBuilder {
-    lists: ListBuilder<ScopedNarrowingConstraintId, ScopedNarrowingConstraintPredicate>,
+    lists: ListBuilder<ScopedNarrowingConstraintPredicate>,
 }
 
 impl NarrowingConstraintsBuilder {
@@ -90,38 +85,45 @@ impl NarrowingConstraintsBuilder {
         }
     }
 
+    /// Clones a narrowing constraint.
+    pub(crate) fn clone_constraint(
+        &mut self,
+        constraint: &ScopedNarrowingConstraint,
+    ) -> ScopedNarrowingConstraint {
+        self.lists.clone_list(constraint)
+    }
+
     /// Adds a predicate to an existing narrowing constraint.
     pub(crate) fn add_predicate_to_constraint(
         &mut self,
-        constraint: Option<ScopedNarrowingConstraintId>,
+        constraint: &mut ScopedNarrowingConstraint,
         predicate: ScopedNarrowingConstraintPredicate,
-    ) -> Option<ScopedNarrowingConstraintId> {
-        self.lists.insert(constraint, predicate)
+    ) {
+        *constraint = self.lists.insert(std::mem::take(constraint), predicate);
     }
 
     /// Returns the intersection of two narrowing constraints. The result contains the predicates
     /// that appear in both inputs.
     pub(crate) fn intersect_constraints(
         &mut self,
-        a: Option<ScopedNarrowingConstraintId>,
-        b: Option<ScopedNarrowingConstraintId>,
-    ) -> Option<ScopedNarrowingConstraintId> {
+        a: ScopedNarrowingConstraint,
+        b: ScopedNarrowingConstraint,
+    ) -> ScopedNarrowingConstraint {
         self.lists.intersect(a, b)
     }
 }
 
 // Iteration
 // ---------
 
-pub(crate) type NarrowingConstraintsIterator<'a> = std::iter::Copied<
-    ListSetReverseIterator<'a, ScopedNarrowingConstraintId, ScopedNarrowingConstraintPredicate>,
->;
+pub(crate) type NarrowingConstraintsIterator<'a> =
+    std::iter::Copied<ListSetReverseIterator<'a, ScopedNarrowingConstraintPredicate>>;
 
 impl NarrowingConstraints {
     /// Iterates over the predicates in a narrowing constraint.
     pub(crate) fn iter_predicates(
         &self,
-        set: Option<ScopedNarrowingConstraintId>,
+        set: &ScopedNarrowingConstraint,
     ) -> NarrowingConstraintsIterator<'_> {
         self.lists.iter_set_reverse(set).copied()
     }
@@ -143,7 +145,7 @@ mod tests {
     impl NarrowingConstraintsBuilder {
         pub(crate) fn iter_predicates(
             &self,
-            set: Option<ScopedNarrowingConstraintId>,
+            set: &ScopedNarrowingConstraint,
         ) -> NarrowingConstraintsIterator<'_> {
             self.lists.iter_set_reverse(set).copied()
         }

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -441,7 +441,7 @@ impl<'map, 'db> Iterator for BindingWithConstraintsIterator<'map, 'db> {
                 narrowing_constraint: ConstraintsIterator {
                     predicates,
                     constraint_ids: narrowing_constraints
-                        .iter_predicates(live_binding.narrowing_constraint),
+                        .iter_predicates(&live_binding.narrowing_constraint),
                 },
                 visibility_constraint: live_binding.visibility_constraint,
             })
@@ -506,12 +506,25 @@ impl<'db> Iterator for DeclarationsIterator<'_, 'db> {
 impl std::iter::FusedIterator for DeclarationsIterator<'_, '_> {}
 
 /// A snapshot of the definitions and constraints state at a particular point in control flow.
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 pub(super) struct FlowSnapshot {
     symbol_states: IndexVec<ScopedSymbolId, SymbolState>,
     scope_start_visibility: ScopedVisibilityConstraintId,
 }
 
+impl FlowSnapshot {
+    pub(super) fn clone(&self, narrowing_constraints: &mut NarrowingConstraintsBuilder) -> Self {
+        Self {
+            symbol_states: self
+                .symbol_states
+                .iter()
+                .map(|state| state.clone(narrowing_constraints))
+                .collect(),
+            scope_start_visibility: self.scope_start_visibility,
+        }
+    }
+}
+
 #[derive(Debug)]
 pub(super) struct UseDefMapBuilder<'db> {
     /// Append-only array of [`Definition`].
@@ -655,8 +668,12 @@ impl<'db> UseDefMapBuilder<'db> {
     ) {
         let def_id = self.all_definitions.push(Some(declaration));
         let symbol_state = &mut self.symbol_states[symbol];
-        self.bindings_by_declaration
-            .insert(declaration, symbol_state.bindings().clone());
+        self.bindings_by_declaration.insert(
+            declaration,
+            symbol_state
+                .bindings()
+                .clone(&mut self.narrowing_constraints),
+        );
         symbol_state.record_declaration(def_id);
     }
 
@@ -676,24 +693,33 @@ impl<'db> UseDefMapBuilder<'db> {
     pub(super) fn record_use(&mut self, symbol: ScopedSymbolId, use_id: ScopedUseId) {
         // We have a use of a symbol; clone the current bindings for that symbol, and record them
         // as the live bindings for this use.
-        let new_use = self
-            .bindings_by_use
-            .push(self.symbol_states[symbol].bindings().clone());
+        let new_use = self.bindings_by_use.push(
+            self.symbol_states[symbol]
+                .bindings()
+                .clone(&mut self.narrowing_constraints),
+        );
         debug_assert_eq!(use_id, new_use);
     }
 
     pub(super) fn snapshot_eager_bindings(
         &mut self,
         enclosing_symbol: ScopedSymbolId,
     ) -> ScopedEagerBindingsId {
-        self.eager_bindings
-            .push(self.symbol_states[enclosing_symbol].bindings().clone())
+        self.eager_bindings.push(
+            self.symbol_states[enclosing_symbol]
+                .bindings()
+                .clone(&mut self.narrowing_constraints),
+        )
     }
 
     /// Take a snapshot of the current visible-symbols state.
-    pub(super) fn snapshot(&self) -> FlowSnapshot {
+    pub(super) fn snapshot(&mut self) -> FlowSnapshot {
         FlowSnapshot {
-            symbol_states: self.symbol_states.clone(),
+            symbol_states: self
+                .symbol_states
+                .iter()
+                .map(|state| state.clone(&mut self.narrowing_constraints))
+                .collect(),
             scope_start_visibility: self.scope_start_visibility,
         }
     }
@@ -713,10 +739,9 @@ impl<'db> UseDefMapBuilder<'db> {
         // If the snapshot we are restoring is missing some symbols we've recorded since, we need
         // to fill them in so the symbol IDs continue to line up. Since they don't exist in the
         // snapshot, the correct state to fill them in with is "undefined".
-        self.symbol_states.resize(
-            num_symbols,
-            SymbolState::undefined(self.scope_start_visibility),
-        );
+        self.symbol_states.resize_with(num_symbols, || {
+            SymbolState::undefined(self.scope_start_visibility)
+        });
     }
 
     /// Merge the given snapshot into the current state, reflecting that we might have taken either