@@ -45,6 +45,87 @@ use crate::types::{
4545use crate :: { Db , FxOrderMap , FxOrderSet } ;
4646use smallvec:: SmallVec ;
4747
48+ /// Extract `(core, guard)` from truthiness-guarded intersections.
49+ ///
50+ /// e.g.
51+ /// - `A & ~AlwaysTruthy` -> `Some((A, ~AlwaysTruthy))`
52+ /// - `A & ~AlwaysFalsy` -> `Some((A, ~AlwaysFalsy))`
53+ /// - `A` -> `None`
54+ /// - `A & ~AlwaysTruthy & ~AlwaysFalsy` -> `None` (not a single-guard shape)
55+ ///
56+ /// This only recognizes the "single truthiness guard" forms used by truthiness narrowing.
57+ fn split_truthiness_guarded_intersection < ' db > (
58+ db : & ' db dyn Db ,
59+ ty : Type < ' db > ,
60+ ) -> Option < ( Type < ' db > , Type < ' db > ) > {
61+ let Type :: Intersection ( intersection) = ty else {
62+ return None ;
63+ } ;
64+ let falsy = Type :: AlwaysTruthy . negate ( db) ;
65+ let truthy = Type :: AlwaysFalsy . negate ( db) ;
66+
67+ let has_not_truthy = intersection. negative ( db) . contains ( & Type :: AlwaysTruthy ) ;
68+ let has_not_falsy = intersection. negative ( db) . contains ( & Type :: AlwaysFalsy ) ;
69+ let guard = match ( has_not_truthy, has_not_falsy) {
70+ ( true , false ) => falsy,
71+ ( false , true ) => truthy,
72+ _ => return None ,
73+ } ;
74+
75+ let mut core = IntersectionBuilder :: new ( db) ;
76+ for positive in intersection. positive ( db) {
77+ core = core. add_positive ( * positive) ;
78+ }
79+ for negative in intersection. negative ( db) {
80+ if ( guard == falsy && * negative == Type :: AlwaysTruthy )
81+ || ( guard == truthy && * negative == Type :: AlwaysFalsy )
82+ {
83+ continue ;
84+ }
85+ core = core. add_negative ( * negative) ;
86+ }
87+ Some ( ( core. build ( ) , guard) )
88+ }
89+
90+ /// Try to merge a complementary guarded pair into an unguarded core.
91+ ///
92+ /// e.g.
93+ /// - `(A & ~AlwaysTruthy, A & ~AlwaysFalsy)` -> `Some(A)`
94+ /// - `(A & ~AlwaysTruthy, B & ~AlwaysFalsy)` -> `Some(A | B)` if reconstruction is exact
95+ /// - `(A & ~AlwaysTruthy, C)` -> `None`
96+ ///
97+ /// Safety rule:
98+ /// The candidate merge is accepted only if adding each original guard back reconstructs
99+ /// exactly the original operands (`left` and `right`).
100+ ///
101+ /// TODO: This processing is specialized for `AlwaysTruthy/AlwaysFalsy`.
102+ /// It would be nice to generalize this in the future.
103+ /// Discussion: <https://github.com/astral-sh/ty/issues/224>
104+ fn merge_truthiness_guarded_pair < ' db > (
105+ db : & ' db dyn Db ,
106+ left : Type < ' db > ,
107+ right : Type < ' db > ,
108+ ) -> Option < Type < ' db > > {
109+ let ( left_core, left_guard) = split_truthiness_guarded_intersection ( db, left) ?;
110+ let ( right_core, right_guard) = split_truthiness_guarded_intersection ( db, right) ?;
111+ if left_guard == right_guard {
112+ return None ;
113+ }
114+
115+ if left_core. is_equivalent_to ( db, right_core) {
116+ return Some ( left_core) ;
117+ }
118+
119+ let candidate = UnionType :: from_elements ( db, [ left_core, right_core] ) ;
120+ let left_reconstructed = IntersectionType :: from_two_elements ( db, candidate, left_guard) ;
121+ let right_reconstructed = IntersectionType :: from_two_elements ( db, candidate, right_guard) ;
122+ if left_reconstructed == left && right_reconstructed == right {
123+ Some ( candidate)
124+ } else {
125+ None
126+ }
127+ }
128+
48129#[ derive( Debug , Clone , Copy , PartialEq , Eq ) ]
49130enum LiteralKind < ' db > {
50131 Int ,
@@ -650,10 +731,13 @@ impl<'db> UnionBuilder<'db> {
650731 }
651732
652733 fn push_type ( & mut self , ty : Type < ' db > , seen_aliases : & mut Vec < Type < ' db > > ) {
653- let bool_pair = if let Some ( LiteralValueTypeKind :: Bool ( b) ) = ty. as_literal_value_kind ( ) {
654- Some ( LiteralValueTypeKind :: Bool ( !b) )
655- } else {
656- None
734+ let mut ty = ty;
735+ let bool_pair = |ty : Type < ' db > | {
736+ if let Some ( LiteralValueTypeKind :: Bool ( b) ) = ty. as_literal_value_kind ( ) {
737+ Some ( LiteralValueTypeKind :: Bool ( !b) )
738+ } else {
739+ None
740+ }
657741 } ;
658742
659743 // If an alias gets here, it means we aren't unpacking aliases, and we also
@@ -686,9 +770,16 @@ impl<'db> UnionBuilder<'db> {
686770 return ;
687771 }
688772
773+ // Fold `(T & ~AlwaysTruthy) | (T & ~AlwaysFalsy)` to `T`.
774+ if let Some ( merged_type) = merge_truthiness_guarded_pair ( self . db , ty, element_type) {
775+ to_remove. push ( i) ;
776+ ty = merged_type;
777+ continue ;
778+ }
779+
689780 if element_type
690781 . as_literal_value_kind ( )
691- . zip ( bool_pair)
782+ . zip ( bool_pair ( ty ) )
692783 . is_some_and ( |( element, pair) | element == pair)
693784 {
694785 self . add_in_place_impl ( KnownClass :: Bool . to_instance ( self . db ) , seen_aliases) ;
0 commit comments