[ty] Use OrderedSet/Map in more places

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>

crates/ty_python_semantic/resources/mdtest/typed_dict.md

@@ -2124,20 +2124,26 @@ shows up in a subset of the union members) is present, but that isn't generally
 field, it could be *assigned to* with another `TypedDict` that does:
 
 ```py
+from typing_extensions import Literal
+
 class Foo(TypedDict):
     foo: int
 
 class Bar(TypedDict):
     bar: int
 
-def disappointment(u: Foo | Bar):
+def disappointment(u: Foo | Bar, v: Literal["foo"]):
     if "foo" in u:
         # We can't narrow the union here...
         reveal_type(u)  # revealed: Foo | Bar
     else:
         # ...(even though we *can* narrow it here)...
-        # TODO: This should narrow to `Bar`, because "foo" is required in `Foo`.
+        reveal_type(u)  # revealed: Bar
+
+    if v in u:
         reveal_type(u)  # revealed: Foo | Bar
+    else:
+        reveal_type(u)  # revealed: Bar
 
 # ...because `u` could turn out to be one of these.
 class FooBar(TypedDict):
@@ -2148,6 +2154,39 @@ static_assert(is_assignable_to(FooBar, Foo))
 static_assert(is_assignable_to(FooBar, Bar))
 ```
 
+`not in` works in the opposite way to `in`: we can narrow in the positive case, but we cannot narrow
+in the negative case. The following snippet also tests our narrowing behaviour for intersections
+that contain `TypedDict`s, and unions that contain intersections that contain `TypedDict`s:
+
+```py
+from typing_extensions import Literal, Any
+from ty_extensions import Intersection, is_assignable_to, static_assert
+
+def _(t: Bar, u: Foo | Intersection[Bar, Any], v: Intersection[Bar, Any], w: Literal["bar"]):
+    reveal_type(u)  # revealed: Foo | (Bar & Any)
+    reveal_type(v)  # revealed: Bar & Any
+
+    if "bar" not in t:
+        reveal_type(t)  # revealed: Never
+    else:
+        reveal_type(t)  # revealed: Bar
+
+    if "bar" not in u:
+        reveal_type(u)  # revealed: Foo
+    else:
+        reveal_type(u)  # revealed: Foo | (Bar & Any)
+
+    if "bar" not in v:
+        reveal_type(v)  # revealed: Never
+    else:
+        reveal_type(v)  # revealed: Bar & Any
+
+    if w not in u:
+        reveal_type(u)  # revealed: Foo
+    else:
+        reveal_type(u)  # revealed: Foo | (Bar & Any)
+```
+
 TODO: The narrowing that we didn't do above will become possible when we add support for
 `closed=True`. This is [one of the main use cases][closed] that motivated the `closed` feature.
 

crates/ty_python_semantic/src/types.rs

@@ -7267,10 +7267,7 @@ impl<'db> Type<'db> {
             }
 
             (Some(Place::Defined(new_method, ..)), Place::Defined(init_method, ..)) => {
-                let callable = UnionBuilder::new(db)
-                    .add(*new_method)
-                    .add(*init_method)
-                    .build();
+                let callable = UnionType::from_elements(db, [new_method, init_method]);
 
                 let new_method_bindings = new_method
                     .bindings(db)
@@ -10758,11 +10755,7 @@ fn walk_type_var_constraints<'db, V: visitor::TypeVisitor<'db> + ?Sized>(
 
 impl<'db> TypeVarConstraints<'db> {
     fn as_type(self, db: &'db dyn Db) -> Type<'db> {
-        let mut builder = UnionBuilder::new(db);
-        for ty in self.elements(db) {
-            builder = builder.add(*ty);
-        }
-        builder.build()
+        UnionType::from_elements(db, self.elements(db))
     }
 
     fn to_instance(self, db: &'db dyn Db) -> Option<TypeVarConstraints<'db>> {
@@ -14107,21 +14100,19 @@ impl<'db> UnionType<'db> {
         self.try_map(db, |element| element.to_instance(db))
     }
 
-    pub(crate) fn filter(self, db: &'db dyn Db, f: impl FnMut(&Type<'db>) -> bool) -> Type<'db> {
-        let current = self.elements(db);
-        let new: Vec<Type<'db>> = current.iter().copied().filter(f).collect();
-        match new.len() {
-            0 => Type::Never,
-            1 => new[0],
-            len if len == current.len() => Type::Union(self),
-            _ => new
-                .iter()
-                .fold(UnionBuilder::new(db), |builder, element| {
-                    builder.add(*element)
-                })
-                .recursively_defined(self.recursively_defined(db))
-                .build(),
-        }
+    pub(crate) fn filter(
+        self,
+        db: &'db dyn Db,
+        mut f: impl FnMut(&Type<'db>) -> bool,
+    ) -> Type<'db> {
+        self.elements(db)
+            .iter()
+            .filter(|ty| f(ty))
+            .fold(UnionBuilder::new(db), |builder, element| {
+                builder.add(*element)
+            })
+            .recursively_defined(self.recursively_defined(db))
+            .build()
     }
 
     pub(crate) fn map_with_boundness(

crates/ty_python_semantic/src/types/class.rs

@@ -8,6 +8,7 @@ use super::{
     SubclassOfType, Truthiness, Type, TypeQualifiers, class_base::ClassBase,
     function::FunctionType,
 };
+use crate::FxOrderMap;
 use crate::place::TypeOrigin;
 use crate::semantic_index::definition::{Definition, DefinitionState};
 use crate::semantic_index::scope::{NodeWithScopeKind, Scope, ScopeKind};
@@ -161,8 +162,8 @@ fn fields_cycle_initial<'db>(
     _self: ClassLiteral<'db>,
     _specialization: Option<Specialization<'db>>,
     _field_policy: CodeGeneratorKind<'db>,
-) -> FxIndexMap<Name, Field<'db>> {
-    FxIndexMap::default()
+) -> FxOrderMap<Name, Field<'db>> {
+    FxOrderMap::default()
 }
 
 /// A category of classes with code generation capabilities (with synthesized methods).
@@ -3147,7 +3148,7 @@ impl<'db> ClassLiteral<'db> {
         db: &'db dyn Db,
         specialization: Option<Specialization<'db>>,
         field_policy: CodeGeneratorKind<'db>,
-    ) -> FxIndexMap<Name, Field<'db>> {
+    ) -> FxOrderMap<Name, Field<'db>> {
         if field_policy == CodeGeneratorKind::NamedTuple {
             // NamedTuples do not allow multiple inheritance, so it is sufficient to enumerate the
             // fields of this class only.
@@ -3195,8 +3196,8 @@ impl<'db> ClassLiteral<'db> {
         db: &'db dyn Db,
         specialization: Option<Specialization<'db>>,
         field_policy: CodeGeneratorKind,
-    ) -> FxIndexMap<Name, Field<'db>> {
-        let mut attributes = FxIndexMap::default();
+    ) -> FxOrderMap<Name, Field<'db>> {
+        let mut attributes = FxOrderMap::default();
 
         let class_body_scope = self.body_scope(db);
         let table = place_table(db, class_body_scope);

crates/ty_python_semantic/src/types/diagnostic.rs

@@ -31,7 +31,7 @@ use crate::types::{
     protocol_class::ProtocolClass,
 };
 use crate::types::{DataclassFlags, KnownInstanceType, MemberLookupPolicy, TypeVarInstance};
-use crate::{Db, DisplaySettings, FxIndexMap, Program, declare_lint};
+use crate::{Db, DisplaySettings, FxOrderMap, Program, declare_lint};
 use itertools::Itertools;
 use ruff_db::{
     diagnostic::{Annotation, Diagnostic, Span, SubDiagnostic, SubDiagnosticSeverity},
@@ -3001,7 +3001,7 @@ pub(crate) fn report_instance_layout_conflict(
 /// The inner data is an `IndexMap` to ensure that diagnostics regarding conflicting disjoint bases
 /// are reported in a stable order.
 #[derive(Debug, Default)]
-pub(super) struct IncompatibleBases<'db>(FxIndexMap<DisjointBase<'db>, IncompatibleBaseInfo<'db>>);
+pub(super) struct IncompatibleBases<'db>(FxOrderMap<DisjointBase<'db>, IncompatibleBaseInfo<'db>>);
 
 impl<'db> IncompatibleBases<'db> {
     pub(super) fn insert(

crates/ty_python_semantic/src/types/enums.rs

@@ -2,7 +2,7 @@ use ruff_python_ast::name::Name;
 use rustc_hash::FxHashMap;
 
 use crate::{
-    Db, FxIndexMap,
+    Db, FxOrderMap,
     place::{Place, PlaceAndQualifiers, place_from_bindings, place_from_declarations},
     semantic_index::{place_table, use_def_map},
     types::{
@@ -13,7 +13,7 @@ use crate::{
 
 #[derive(Debug, PartialEq, Eq, salsa::Update)]
 pub(crate) struct EnumMetadata<'db> {
-    pub(crate) members: FxIndexMap<Name, Type<'db>>,
+    pub(crate) members: FxOrderMap<Name, Type<'db>>,
     pub(crate) aliases: FxHashMap<Name, Name>,
 }
 
@@ -22,7 +22,7 @@ impl get_size2::GetSize for EnumMetadata<'_> {}
 impl EnumMetadata<'_> {
     fn empty() -> Self {
         EnumMetadata {
-            members: FxIndexMap::default(),
+            members: FxOrderMap::default(),
             aliases: FxHashMap::default(),
         }
     }
@@ -253,7 +253,7 @@ pub(crate) fn enum_metadata<'db>(
 
             Some((name.clone(), value_ty))
         })
-        .collect::<FxIndexMap<_, _>>();
+        .collect::<FxOrderMap<_, _>>();
 
     if members.is_empty() {
         // Enum subclasses without members are not considered enums.

crates/ty_python_semantic/src/types/infer/builder/type_expression.rs

@@ -1083,13 +1083,9 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
                     &mut self.inner_expression_inference_state,
                     InnerExpressionInferenceState::Get,
                 );
-                let union = union
-                    .elements(self.db())
-                    .iter()
-                    .fold(UnionBuilder::new(self.db()), |builder, elem| {
-                        builder.add(self.infer_subscript_type_expression(subscript, *elem))
-                    })
-                    .build();
+                let union = union.map(self.db(), |element| {
+                    self.infer_subscript_type_expression(subscript, *element)
+                });
                 self.inner_expression_inference_state = previous_slice_inference_state;
                 union
             }

crates/ty_python_semantic/src/types/narrow.rs

@@ -12,7 +12,7 @@ use crate::types::enums::{enum_member_literals, enum_metadata};
 use crate::types::function::KnownFunction;
 use crate::types::infer::{ExpressionInference, infer_same_file_expression_type};
 use crate::types::typed_dict::{
-    SynthesizedTypedDictType, TypedDictFieldBuilder, TypedDictSchema, TypedDictType,
+    SynthesizedTypedDictType, TypedDictField, TypedDictFieldBuilder, TypedDictSchema, TypedDictType,
 };
 use crate::types::{
     CallableType, ClassLiteral, ClassType, IntersectionBuilder, IntersectionType, KnownClass,
@@ -926,10 +926,7 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
                 .build();
 
             // Keep order: first literal complement, then broader arms.
-            let result = UnionBuilder::new(self.db)
-                .add(narrowed_single)
-                .add(rest_union)
-                .build();
+            let result = UnionType::from_elements(self.db, [narrowed_single, rest_union]);
             Some(result)
         } else {
             None
@@ -1027,23 +1024,31 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
             && rhs_ty.is_singleton(self.db)
         {
             let is_positive_check = is_positive == (ops[0] == ast::CmpOp::Is);
-            let filtered = union.filter(self.db, |elem| {
-                elem.as_nominal_instance()
-                    .and_then(|inst| inst.tuple_spec(self.db))
-                    .and_then(|spec| spec.py_index(self.db, index).ok())
-                    .is_none_or(|el_ty| {
-                        if is_positive_check {
-                            // `is X` context: keep tuples where element could be X
-                            !el_ty.is_disjoint_from(self.db, rhs_ty)
-                        } else {
-                            // `is not X` context: keep tuples where element is not always X
-                            !el_ty.is_subtype_of(self.db, rhs_ty)
-                        }
-                    })
-            });
-            if filtered != Type::Union(union) {
+            let filtered: Vec<_> = union
+                .elements(self.db)
+                .iter()
+                .filter(|elem| {
+                    elem.as_nominal_instance()
+                        .and_then(|inst| inst.tuple_spec(self.db))
+                        .and_then(|spec| spec.py_index(self.db, index).ok())
+                        .is_none_or(|el_ty| {
+                            if is_positive_check {
+                                // `is X` context: keep tuples where element could be X
+                                !el_ty.is_disjoint_from(self.db, rhs_ty)
+                            } else {
+                                // `is not X` context: keep tuples where element is not always X
+                                !el_ty.is_subtype_of(self.db, rhs_ty)
+                            }
+                        })
+                })
+                .copied()
+                .collect();
+            if filtered.len() < union.elements(self.db).len() {
                 let place = self.expect_place(&subscript_place_expr);
-                constraints.insert(place, NarrowingConstraint::typeguard(filtered));
+                constraints.insert(
+                    place,
+                    NarrowingConstraint::regular(UnionType::from_elements(self.db, filtered)),
+                );
             }
         }
 
@@ -1091,6 +1096,75 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
             }
         }
 
+        // Narrow unions and intersections of `TypedDict` in cases where required keys are
+        // excluded:
+        //
+        // class Foo(TypedDict):
+        //     foo: int
+        // class Bar(TypedDict):
+        //     bar: int
+        //
+        // def _(u: Foo | Bar):
+        //     if "foo" not in u:
+        //         reveal_type(u)  # revealed: Bar
+        if matches!(&**ops, [ast::CmpOp::In | ast::CmpOp::NotIn])
+            && let Type::StringLiteral(key) = inference.expression_type(&**left)
+            && let Some(rhs_place_expr) = place_expr(&comparators[0])
+            && let rhs_type = inference.expression_type(&comparators[0])
+            && is_typeddict_or_union_with_typeddicts(self.db, rhs_type)
+        {
+            let is_negative_check = is_positive == (ops[0] == ast::CmpOp::NotIn);
+            if is_negative_check {
+                let requires_key = |td: TypedDictType<'db>| -> bool {
+                    td.items(self.db)
+                        .get(key.value(self.db))
+                        .is_some_and(TypedDictField::is_required)
+                };
+
+                let narrowed = match rhs_type {
+                    Type::TypedDict(td) => {
+                        if requires_key(td) {
+                            Type::Never
+                        } else {
+                            rhs_type
+                        }
+                    }
+                    Type::Intersection(intersection) => {
+                        if intersection
+                            .positive(self.db)
+                            .iter()
+                            .copied()
+                            .filter_map(Type::as_typed_dict)
+                            .any(requires_key)
+                        {
+                            Type::Never
+                        } else {
+                            rhs_type
+                        }
+                    }
+                    Type::Union(union) => {
+                        // remove all members of the union that would require the key
+                        union.filter(self.db, |ty| match ty {
+                            Type::TypedDict(td) => !requires_key(*td),
+                            Type::Intersection(intersection) => !intersection
+                                .positive(self.db)
+                                .iter()
+                                .copied()
+                                .filter_map(Type::as_typed_dict)
+                                .any(requires_key),
+                            _ => true,
+                        })
+                    }
+                    _ => rhs_type,
+                };
+
+                if narrowed != rhs_type {
+                    let place = self.expect_place(&rhs_place_expr);
+                    constraints.insert(place, NarrowingConstraint::typeguard(narrowed));
+                }
+            }
+        }
+
         let mut last_rhs_ty: Option<Type> = None;
 
         for (op, (left, right)) in std::iter::zip(&**ops, comparator_tuples) {
@@ -1631,25 +1705,33 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
         }
 
         // Filter the union based on whether each tuple element at the index could match the rhs.
-        let filtered = union.filter(self.db, |elem| {
-            elem.as_nominal_instance()
-                .and_then(|inst| inst.tuple_spec(self.db))
-                .and_then(|spec| spec.py_index(self.db, index).ok())
-                .is_none_or(|el_ty| {
-                    if constrain_with_equality {
-                        // Keep tuples where element could be equal to rhs.
-                        !el_ty.is_disjoint_from(self.db, rhs_type)
-                    } else {
-                        // Keep tuples where element is not always equal to rhs.
-                        !el_ty.is_subtype_of(self.db, rhs_type)
-                    }
-                })
-        });
+        let filtered: Vec<_> = union
+            .elements(self.db)
+            .iter()
+            .filter(|elem| {
+                elem.as_nominal_instance()
+                    .and_then(|inst| inst.tuple_spec(self.db))
+                    .and_then(|spec| spec.py_index(self.db, index).ok())
+                    .is_none_or(|el_ty| {
+                        if constrain_with_equality {
+                            // Keep tuples where element could be equal to rhs.
+                            !el_ty.is_disjoint_from(self.db, rhs_type)
+                        } else {
+                            // Keep tuples where element is not always equal to rhs.
+                            !el_ty.is_subtype_of(self.db, rhs_type)
+                        }
+                    })
+            })
+            .copied()
+            .collect();
 
         // Only create a constraint if we actually narrowed something.
-        if filtered != rhs_type {
+        if filtered.len() < union.elements(self.db).len() {
             let place = self.expect_place(&subscript_place_expr);
-            Some((place, NarrowingConstraint::typeguard(filtered)))
+            Some((
+                place,
+                NarrowingConstraint::regular(UnionType::from_elements(self.db, filtered)),
+            ))
         } else {
             None
         }
@@ -1661,18 +1743,13 @@ impl<'db, 'ast> NarrowingConstraintsBuilder<'db, 'ast> {
 fn is_typeddict_or_union_with_typeddicts<'db>(db: &'db dyn Db, ty: Type<'db>) -> bool {
     match ty {
         Type::TypedDict(_) => true,
-        Type::Union(union) => {
-            union
-                .elements(db)
-                .iter()
-                .any(|union_member_ty| match union_member_ty {
-                    Type::TypedDict(_) => true,
-                    Type::Intersection(intersection) => {
-                        intersection.positive(db).iter().any(Type::is_typed_dict)
-                    }
-                    _ => false,
-                })
+        Type::Intersection(intersection) => {
+            intersection.positive(db).iter().any(Type::is_typed_dict)
         }
+        Type::Union(union) => union
+            .elements(db)
+            .iter()
+            .any(|union_member_ty| is_typeddict_or_union_with_typeddicts(db, *union_member_ty)),
         _ => false,
     }
 }