Stop deduplicating constraints

Author: dcreager

crates/red_knot_python_semantic/src/semantic_index/builder.rs

@@ -457,7 +457,7 @@ impl<'db> SemanticIndexBuilder<'db> {
         let constraint_id = self.current_use_def_map_mut().add_constraint(constraint);
         let id = self
             .current_visibility_constraints_mut()
-            .add_atom(constraint_id, 0);
+            .add_atom(constraint_id);
         self.record_visibility_constraint_id(id);
         id
     }
@@ -1187,13 +1187,14 @@ where
                 // We need multiple copies of the visibility constraint for the while condition,
                 // since we need to model situations where the first evaluation of the condition
                 // returns True, but a later evaluation returns False.
-                let constraint_id = self.current_use_def_map_mut().add_constraint(constraint);
+                let first_constraint_id = self.current_use_def_map_mut().add_constraint(constraint);
+                let later_constraint_id = self.current_use_def_map_mut().add_constraint(constraint);
                 let first_vis_constraint_id = self
                     .current_visibility_constraints_mut()
-                    .add_atom(constraint_id, 0);
+                    .add_atom(first_constraint_id);
                 let later_vis_constraint_id = self
                     .current_visibility_constraints_mut()
-                    .add_atom(constraint_id, 1);
+                    .add_atom(later_constraint_id);
 
                 // Save aside any break states from an outer loop
                 let saved_break_states = std::mem::take(&mut self.loop_break_states);
@@ -1780,7 +1781,7 @@ where
                         };
                         let visibility_constraint = self
                             .current_visibility_constraints_mut()
-                            .add_atom(constraint_id, 0);
+                            .add_atom(constraint_id);
 
                         let after_expr = self.flow_snapshot();
 

crates/red_knot_python_semantic/src/semantic_index/constraint.rs

@@ -1,7 +1,6 @@
 use ruff_db::files::File;
 use ruff_index::{newtype_index, IndexVec};
 use ruff_python_ast::Singleton;
-use rustc_hash::FxHashMap;
 
 use crate::db::Db;
 use crate::semantic_index::expression::Expression;
@@ -18,16 +17,12 @@ pub(crate) type Constraints<'db> = IndexVec<ScopedConstraintId, Constraint<'db>>
 #[derive(Debug, Default)]
 pub(crate) struct ConstraintsBuilder<'db> {
     constraints: IndexVec<ScopedConstraintId, Constraint<'db>>,
-    constraint_cache: FxHashMap<Constraint<'db>, ScopedConstraintId>,
 }
 
 impl<'db> ConstraintsBuilder<'db> {
     /// Adds a constraint, ensuring that we only store any particular constraint once.
     pub(crate) fn add_constraint(&mut self, constraint: Constraint<'db>) -> ScopedConstraintId {
-        *self
-            .constraint_cache
-            .entry(constraint)
-            .or_insert_with(|| self.constraints.push(constraint))
+        self.constraints.push(constraint)
     }
 
     pub(crate) fn build(mut self) -> Constraints<'db> {

crates/red_knot_python_semantic/src/visibility_constraints.rs

@@ -236,60 +236,15 @@ impl std::fmt::Debug for ScopedVisibilityConstraintId {
 
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
 struct InteriorNode {
-    atom: Atom,
+    /// A "variable" that is evaluated as part of a TDD ternary function. For visibility
+    /// constraints, this is a `Constraint` that represents some runtime property of the Python
+    /// code that we are evaluating.
+    atom: ScopedConstraintId,
     if_true: ScopedVisibilityConstraintId,
     if_ambiguous: ScopedVisibilityConstraintId,
     if_false: ScopedVisibilityConstraintId,
 }
 
-/// A "variable" that is evaluated as part of a TDD ternary function. For visibility constraints,
-/// this is a `Constraint` that represents some runtime property of the Python code that we are
-/// evaluating. We intern these constraints in an arena ([`VisibilityConstraints::constraints`]).
-/// An atom is then an index into this arena.
-///
-/// By using a 32-bit index, we would typically allow 4 billion distinct constraints within a
-/// scope. However, we sometimes have to model how a `Constraint` can have a different runtime
-/// value at different points in the execution of the program. To handle this, we reserve the top
-/// byte of an atom to represent a "copy number". This is just an opaque value that allows
-/// different `Atom`s to evaluate the same `Constraint`. This yields a maximum of 16 million
-/// distinct `Constraint`s in a scope, and 256 possible copies of each of those constraints.
-#[derive(Clone, Copy, Eq, Hash, Ord, PartialEq, PartialOrd)]
-struct Atom(u32);
-
-impl Atom {
-    /// Deconstruct an atom into a constraint index and a copy number.
-    #[inline]
-    fn into_index_and_copy(self) -> (ScopedConstraintId, u8) {
-        let copy = self.0 >> 24;
-        let index = self.0 & 0x00ff_ffff;
-        (index.into(), copy as u8)
-    }
-
-    #[inline]
-    fn copy_of(mut self, copy: u8) -> Self {
-        // Clear out the previous copy number
-        self.0 &= 0x00ff_ffff;
-        // OR in the new one
-        self.0 |= u32::from(copy) << 24;
-        self
-    }
-}
-
-impl From<ScopedConstraintId> for Atom {
-    fn from(constraint: ScopedConstraintId) -> Atom {
-        Atom(constraint.as_u32())
-    }
-}
-
-// A custom Debug implementation that prints out the constraint index and copy number as distinct
-// fields.
-impl std::fmt::Debug for Atom {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let (index, copy) = self.into_index_and_copy();
-        f.debug_tuple("Atom").field(&index).field(&copy).finish()
-    }
-}
-
 impl ScopedVisibilityConstraintId {
     /// A special ID that is used for an "always true" / "always visible" constraint.
     pub(crate) const ALWAYS_TRUE: ScopedVisibilityConstraintId =
@@ -377,11 +332,6 @@ impl VisibilityConstraintsBuilder {
         }
     }
 
-    /// Adds a constraint, ensuring that we only store any particular constraint once.
-    fn add_constraint(&mut self, constraint: ScopedConstraintId, copy: u8) -> Atom {
-        Atom::from(constraint).copy_of(copy)
-    }
-
     /// Adds an interior node, ensuring that we always use the same visibility constraint ID for
     /// equal nodes.
     fn add_interior(&mut self, node: InteriorNode) -> ScopedVisibilityConstraintId {
@@ -397,17 +347,23 @@ impl VisibilityConstraintsBuilder {
             .or_insert_with(|| self.interiors.push(node))
     }
 
-    /// Adds a new visibility constraint that checks a single [`Constraint`]. Provide different
-    /// values for `copy` if you need to model that the constraint can evaluate to different
-    /// results at different points in the execution of the program being modeled.
+    /// Adds a new visibility constraint that checks a single [`Constraint`].
+    ///
+    /// [`ScopedConstraintId`]s are the “variables” that are evaluated by a TDD. A TDD variable has
+    /// the same value no matter how many times it appears in the ternary formula that the TDD
+    /// represents.
+    ///
+    /// However, we sometimes have to model how a `Constraint` can have a different runtime
+    /// value at different points in the execution of the program. To handle this, you can take
+    /// advantage of the fact that the [`Constraints`] arena does not deduplicate `Constraint`s.
+    /// You can add a `Constraint` multiple times, yielding different `ScopeConstraintId`s, which
+    /// you can then create separate TDD atoms for.
     pub(crate) fn add_atom(
         &mut self,
         constraint: ScopedConstraintId,
-        copy: u8,
     ) -> ScopedVisibilityConstraintId {
-        let atom = self.add_constraint(constraint, copy);
         self.add_interior(InteriorNode {
-            atom,
+            atom: constraint,
             if_true: ALWAYS_TRUE,
             if_ambiguous: AMBIGUOUS,
             if_false: ALWAYS_FALSE,
@@ -592,8 +548,7 @@ impl VisibilityConstraints {
                 ALWAYS_FALSE => return Truthiness::AlwaysFalse,
                 _ => self.interiors[id],
             };
-            let (constraint_id, _) = node.atom.into_index_and_copy();
-            let constraint = &constraints[constraint_id];
+            let constraint = &constraints[node.atom];
             match Self::analyze_single(db, constraint) {
                 Truthiness::AlwaysTrue => id = node.if_true,
                 Truthiness::Ambiguous => id = node.if_ambiguous,

crates/ruff_macros/src/newtype_index.rs

@@ -32,7 +32,7 @@ pub(super) fn generate_newtype_index(item: ItemStruct) -> syn::Result<proc_macro
     let semi_token = semi_token.unwrap_or_default();
     let output = quote! {
         #(#attrs)*
-        #[derive(Copy, Clone, Eq, PartialEq, Hash)]
+        #[derive(Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd)]
         #vis #struct_token #ident(std::num::NonZeroU32)#semi_token
 
         impl #ident {