Skip to content

Commit f22790c

Browse files
stroxlermeta-codesync[bot]
authored andcommitted
Implement JoinStyle::map, pass join style to TypeInfo::join
Summary: Just plumbing. Reviewed By: yangdanny97 Differential Revision: D84381068 fbshipit-source-id: f5afc8980e22e13eb373a75862e17b7a7b1e999a
1 parent 839d13d commit f22790c

File tree

3 files changed

+37
-11
lines changed

3 files changed

+37
-11
lines changed

crates/pyrefly_types/src/type_info.rs

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,16 @@ pub enum JoinStyle<T> {
4141
NarrowOf(T),
4242
}
4343

44+
impl<T> JoinStyle<T> {
45+
pub fn map<S>(&self, f: impl FnOnce(&T) -> S) -> JoinStyle<S> {
46+
match self {
47+
JoinStyle::SimpleMerge => JoinStyle::SimpleMerge,
48+
JoinStyle::ReassignmentOf(x) => JoinStyle::ReassignmentOf(f(x)),
49+
JoinStyle::NarrowOf(x) => JoinStyle::NarrowOf(f(x)),
50+
}
51+
}
52+
}
53+
4454
/// The `TypeInfo` datatype represents type information associated with a
4555
/// name or expression in a control flow context.
4656
///
@@ -112,7 +122,11 @@ impl TypeInfo {
112122
/// - Drop narrowing for facet chains where at least one branch does not narrow
113123
///
114124
/// In the case where there are no branches, we get `Never` with no narrows.
115-
pub fn join(mut branches: Vec<Self>, union_types: &impl Fn(Vec<Type>) -> Type) -> Self {
125+
pub fn join(
126+
mut branches: Vec<Self>,
127+
union_types: &impl Fn(Vec<Type>) -> Type,
128+
_join_style: JoinStyle<Arc<TypeInfo>>,
129+
) -> Self {
116130
match branches.len() {
117131
0 => Self::of_ty(Type::never()),
118132
1 => branches.pop().unwrap(),
@@ -535,6 +549,7 @@ mod tests {
535549
use crate::class::ClassType;
536550
use crate::display::tests::fake_class;
537551
use crate::facet::FacetKind;
552+
use crate::type_info::JoinStyle;
538553
use crate::type_info::TypeInfo;
539554
use crate::types::TArgs;
540555
use crate::types::Type;
@@ -632,13 +647,17 @@ mod tests {
632647

633648
#[test]
634649
fn test_type_info_empty_join() {
635-
let type_info = TypeInfo::join(Vec::new(), &|ts| {
636-
if ts.is_empty() {
637-
fake_class_type("Never")
638-
} else {
639-
fake_class_type("FakeUnionType")
640-
}
641-
});
650+
let type_info = TypeInfo::join(
651+
Vec::new(),
652+
&|ts| {
653+
if ts.is_empty() {
654+
fake_class_type("Never")
655+
} else {
656+
fake_class_type("FakeUnionType")
657+
}
658+
},
659+
JoinStyle::SimpleMerge,
660+
);
642661
assert_eq!(type_info.to_string(), "Never");
643662
}
644663

pyrefly/lib/alt/narrow.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
use num_traits::ToPrimitive;
99
use pyrefly_python::ast::Ast;
10+
use pyrefly_types::type_info::JoinStyle;
1011
use pyrefly_util::prelude::SliceExt;
1112
use ruff_python_ast::Arguments;
1213
use ruff_python_ast::AtomicNodeIndex;
@@ -921,6 +922,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
921922
NarrowOp::Or(ops) => TypeInfo::join(
922923
ops.map(|op| self.narrow(type_info, op, range, errors)),
923924
&|tys| self.unions(tys),
925+
JoinStyle::SimpleMerge,
924926
),
925927
}
926928
}

pyrefly/lib/alt/solve.rs

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ use pyrefly_python::ast::Ast;
1414
use pyrefly_python::dunder;
1515
use pyrefly_python::short_identifier::ShortIdentifier;
1616
use pyrefly_types::facet::FacetKind;
17+
use pyrefly_types::type_info::JoinStyle;
1718
use pyrefly_types::typed_dict::ExtraItems;
1819
use pyrefly_types::typed_dict::TypedDict;
1920
use pyrefly_util::prelude::SliceExt;
@@ -1940,7 +1941,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
19401941
Binding::Narrow(k, op, range) => {
19411942
self.narrow(self.get_idx(*k).as_ref(), op, *range, errors)
19421943
}
1943-
Binding::Phi(_, ks) => {
1944+
Binding::Phi(join_style, ks) => {
19441945
if ks.len() == 1 {
19451946
self.get_idx(*ks.first().unwrap()).arc_clone()
19461947
} else {
@@ -1957,7 +1958,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
19571958
}
19581959
})
19591960
.collect::<Vec<_>>();
1960-
TypeInfo::join(type_infos, &|ts| self.unions(ts))
1961+
TypeInfo::join(
1962+
type_infos,
1963+
&|ts| self.unions(ts),
1964+
join_style.map(|idx| self.get_idx(*idx)),
1965+
)
19611966
}
19621967
}
19631968
Binding::LoopPhi(default, ks) => {
@@ -1980,7 +1985,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
19801985
}
19811986
})
19821987
.collect::<Vec<_>>();
1983-
TypeInfo::join(type_infos, &|ts| self.unions(ts))
1988+
TypeInfo::join(type_infos, &|ts| self.unions(ts), JoinStyle::SimpleMerge)
19841989
}
19851990
}
19861991
Binding::AssignToAttribute(attr, got) => {

0 commit comments

Comments
 (0)