Implement JoinStyle::map, pass join style to TypeInfo::join

Summary: Just plumbing.

Reviewed By: yangdanny97

Differential Revision: D84381068

fbshipit-source-id: f5afc8980e22e13eb373a75862e17b7a7b1e999a

Author: stroxler

crates/pyrefly_types/src/type_info.rs

@@ -41,6 +41,16 @@ pub enum JoinStyle<T> {
     NarrowOf(T),
 }
 
+impl<T> JoinStyle<T> {
+    pub fn map<S>(&self, f: impl FnOnce(&T) -> S) -> JoinStyle<S> {
+        match self {
+            JoinStyle::SimpleMerge => JoinStyle::SimpleMerge,
+            JoinStyle::ReassignmentOf(x) => JoinStyle::ReassignmentOf(f(x)),
+            JoinStyle::NarrowOf(x) => JoinStyle::NarrowOf(f(x)),
+        }
+    }
+}
+
 /// The `TypeInfo` datatype represents type information associated with a
 /// name or expression in a control flow context.
 ///
@@ -112,7 +122,11 @@ impl TypeInfo {
     /// - Drop narrowing for facet chains where at least one branch does not narrow
     ///
     /// In the case where there are no branches, we get `Never` with no narrows.
-    pub fn join(mut branches: Vec<Self>, union_types: &impl Fn(Vec<Type>) -> Type) -> Self {
+    pub fn join(
+        mut branches: Vec<Self>,
+        union_types: &impl Fn(Vec<Type>) -> Type,
+        _join_style: JoinStyle<Arc<TypeInfo>>,
+    ) -> Self {
         match branches.len() {
             0 => Self::of_ty(Type::never()),
             1 => branches.pop().unwrap(),
@@ -535,6 +549,7 @@ mod tests {
     use crate::class::ClassType;
     use crate::display::tests::fake_class;
     use crate::facet::FacetKind;
+    use crate::type_info::JoinStyle;
     use crate::type_info::TypeInfo;
     use crate::types::TArgs;
     use crate::types::Type;
@@ -632,13 +647,17 @@ mod tests {
 
     #[test]
     fn test_type_info_empty_join() {
-        let type_info = TypeInfo::join(Vec::new(), &|ts| {
-            if ts.is_empty() {
-                fake_class_type("Never")
-            } else {
-                fake_class_type("FakeUnionType")
-            }
-        });
+        let type_info = TypeInfo::join(
+            Vec::new(),
+            &|ts| {
+                if ts.is_empty() {
+                    fake_class_type("Never")
+                } else {
+                    fake_class_type("FakeUnionType")
+                }
+            },
+            JoinStyle::SimpleMerge,
+        );
         assert_eq!(type_info.to_string(), "Never");
     }
 

pyrefly/lib/alt/narrow.rs

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

pyrefly/lib/alt/solve.rs

@@ -14,6 +14,7 @@ use pyrefly_python::ast::Ast;
 use pyrefly_python::dunder;
 use pyrefly_python::short_identifier::ShortIdentifier;
 use pyrefly_types::facet::FacetKind;
+use pyrefly_types::type_info::JoinStyle;
 use pyrefly_types::typed_dict::ExtraItems;
 use pyrefly_types::typed_dict::TypedDict;
 use pyrefly_util::prelude::SliceExt;
@@ -1940,7 +1941,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
             Binding::Narrow(k, op, range) => {
                 self.narrow(self.get_idx(*k).as_ref(), op, *range, errors)
             }
-            Binding::Phi(_, ks) => {
+            Binding::Phi(join_style, ks) => {
                 if ks.len() == 1 {
                     self.get_idx(*ks.first().unwrap()).arc_clone()
                 } else {
@@ -1957,7 +1958,11 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                             }
                         })
                         .collect::<Vec<_>>();
-                    TypeInfo::join(type_infos, &|ts| self.unions(ts))
+                    TypeInfo::join(
+                        type_infos,
+                        &|ts| self.unions(ts),
+                        join_style.map(|idx| self.get_idx(*idx)),
+                    )
                 }
             }
             Binding::LoopPhi(default, ks) => {
@@ -1980,7 +1985,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
                             }
                         })
                         .collect::<Vec<_>>();
-                    TypeInfo::join(type_infos, &|ts| self.unions(ts))
+                    TypeInfo::join(type_infos, &|ts| self.unions(ts), JoinStyle::SimpleMerge)
                 }
             }
             Binding::AssignToAttribute(attr, got) => {