some optimization and simplification

This commit is contained in:
Eric Mark Martin 2025-12-20 14:02:20 -05:00
parent 7b3eea38ef
commit 34390c8b9c

View file

@ -19,7 +19,7 @@ use crate::types::{
use ruff_db::parsed::{ParsedModuleRef, parsed_module};
use ruff_python_stdlib::identifiers::is_identifier;
use itertools::Itertools;
use itertools::{Either, Itertools};
use ruff_python_ast as ast;
use ruff_python_ast::{BoolOp, ExprBoolOp};
use rustc_hash::FxHashMap;
@ -366,28 +366,30 @@ impl<'db> NarrowingConstraint<'db> {
pub(crate) fn merge_constraint_and(&self, other: &Self, db: &'db dyn Db) -> Self {
// Distribute AND over OR: (A1 | A2 | ...) AND (B1 | B2 | ...)
// becomes (A1 & B1) | (A1 & B2) | ... | (A2 & B1) | ...
let new_disjuncts = self
let new_disjuncts = other
.disjuncts
.iter()
.cartesian_product(other.disjuncts.iter())
.map(|(left_conj, right_conj)| {
.flat_map(|right_conj| {
// We iterate the RHS first because if it has a typeguard then we don't need to consider the LHS
if right_conj.typeguard.is_some() {
// If the right conjunct has a TypeGuard, it "wins" the conjunction
*right_conj
Either::Left(std::iter::once(*right_conj))
} else {
// Intersect the regular constraints
let new_regular = IntersectionBuilder::new(db)
.add_positive(left_conj.constraint)
.add_positive(right_conj.constraint)
.build();
// Otherwise, we need to consider all LHS disjuncts
Either::Right(self.disjuncts.iter().map(|left_conj| {
let new_regular = IntersectionBuilder::new(db)
.add_positive(left_conj.constraint)
.add_positive(right_conj.constraint)
.build();
Conjunction {
constraint: new_regular,
typeguard: left_conj.typeguard,
}
Conjunction {
constraint: new_regular,
typeguard: left_conj.typeguard,
}
}))
}
})
.collect::<SmallVec<_>>();
.collect();
NarrowingConstraint {
disjuncts: new_disjuncts,
@ -450,10 +452,13 @@ fn merge_constraints_and<'db>(
///
/// However, if a place appears in only one branch of the OR, we need to widen it
/// to `object` in the overall result (because the other branch doesn't constrain it).
///
/// When none of the disjuncts have `TypeGuard`, we simplify the constraint types
/// via `UnionBuilder` to enable simplifications like `~AlwaysFalsy | ~AlwaysTruthy -> object`.
fn merge_constraints_or<'db>(
into: &mut NarrowingConstraints<'db>,
from: &NarrowingConstraints<'db>,
_db: &'db dyn Db,
db: &'db dyn Db,
) {
// For places that appear in `into` but not in `from`, widen to object
into.retain(|key, _| from.contains_key(key));
@ -461,11 +466,35 @@ fn merge_constraints_or<'db>(
for (key, from_constraint) in from {
match into.entry(*key) {
Entry::Occupied(mut entry) => {
// Simply concatenate the disjuncts
entry
.get_mut()
let into_constraint = entry.get_mut();
// Concatenate disjuncts
into_constraint
.disjuncts
.extend(from_constraint.disjuncts.clone());
// If none of the disjuncts have TypeGuard, we can simplify the constraint types
// via UnionBuilder. This enables simplifications like:
// `~AlwaysFalsy | ~AlwaysTruthy -> object`
let all_regular = into_constraint
.disjuncts
.iter()
.all(|conj| conj.typeguard.is_none());
if all_regular {
// Simplify via UnionBuilder
let simplified = UnionType::from_elements(
db,
into_constraint.disjuncts.iter().map(|conj| conj.constraint),
);
// If simplified to object, we can drop the constraint entirely
if simplified.is_object() {
// Remove this entry since it provides no constraint
entry.remove();
} else {
// Replace with simplified constraint
into_constraint.disjuncts = smallvec![Conjunction::regular(simplified)];
}
}
}
Entry::Vacant(_) => {
// Place only appears in `from`, not in `into`. No constraint needed.