[ty] Re-try argument matching for argument type expansion

Author: dhruvmanila

crates/ty_python_semantic/resources/mdtest/call/function.md

@@ -587,7 +587,7 @@ def _(t: tuple[int, str, int]) -> None:
     f(*t)
 
 def _(t: tuple[int, str] | tuple[int, str, int]) -> None:
-    # TODO: error: [invalid-argument-type]
+    # error: [no-matching-overload]
     f(*t)
 ```
 

crates/ty_python_semantic/resources/mdtest/call/overloads.md

@@ -889,6 +889,8 @@ def _(a: int | None):
     )
 ```
 
+### Retry from parameter matching
+
 ## Filtering based on `Any` / `Unknown`
 
 This is the step 5 of the overload call evaluation algorithm which specifies that:

crates/ty_python_semantic/src/types/call/arguments.rs

@@ -202,10 +202,22 @@ impl<'a, 'db> CallArguments<'a, 'db> {
                     for subtype in &expanded_types {
                         let mut new_expanded_types = pre_expanded_types.to_vec();
                         new_expanded_types[index] = Some(*subtype);
-                        expanded_arguments.push(CallArguments::new(
-                            self.arguments.clone(),
-                            new_expanded_types,
-                        ));
+
+                        // Update the arguments list to handle variadic argument expansion
+                        let mut new_arguments = self.arguments.clone();
+                        if let Argument::Variadic(_) = self.arguments[index] {
+                            // If the expanded type is a tuple, update the TupleLength
+                            if let Some(expanded_type) = new_expanded_types[index] {
+                                let length = expanded_type
+                                    .try_iterate(db)
+                                    .map(|tuple| tuple.len())
+                                    .unwrap_or(TupleLength::unknown());
+                                new_arguments[index] = Argument::Variadic(length);
+                            }
+                        }
+
+                        expanded_arguments
+                            .push(CallArguments::new(new_arguments, new_expanded_types));
                     }
                 }
 

crates/ty_python_semantic/src/types/call/bind.rs

@@ -131,7 +131,14 @@ impl<'db> Bindings<'db> {
         argument_types: &CallArguments<'_, 'db>,
     ) -> Result<Self, CallError<'db>> {
         for element in &mut self.elements {
-            element.check_types(db, argument_types);
+            if let Some((global_argument_forms, global_conflicting_forms)) =
+                element.check_types(db, argument_types)
+            {
+                // If this element returned global forms (indicating successful argument type expansion),
+                // update the Bindings with these forms
+                self.argument_forms = global_argument_forms.into();
+                self.conflicting_forms = global_conflicting_forms.into();
+            }
         }
 
         self.evaluate_known_cases(db);
@@ -1261,7 +1268,11 @@ impl<'db> CallableBinding<'db> {
         }
     }
 
-    fn check_types(&mut self, db: &'db dyn Db, argument_types: &CallArguments<'_, 'db>) {
+    fn check_types(
+        &mut self,
+        db: &'db dyn Db,
+        argument_types: &CallArguments<'_, 'db>,
+    ) -> Option<(Vec<Option<ParameterForm>>, Vec<bool>)> {
         // If this callable is a bound method, prepend the self instance onto the arguments list
         // before checking.
         let argument_types = argument_types.with_self(self.bound_type);
@@ -1275,14 +1286,14 @@ impl<'db> CallableBinding<'db> {
                 if let [overload] = self.overloads.as_mut_slice() {
                     overload.check_types(db, argument_types.as_ref());
                 }
-                return;
+                return None;
             }
             MatchingOverloadIndex::Single(index) => {
                 // If only one candidate overload remains, it is the winning match. Evaluate it as
                 // a regular (non-overloaded) call.
                 self.matching_overload_index = Some(index);
                 self.overloads[index].check_types(db, argument_types.as_ref());
-                return;
+                return None;
             }
             MatchingOverloadIndex::Multiple(indexes) => {
                 // If two or more candidate overloads remain, proceed to step 2.
@@ -1308,7 +1319,7 @@ impl<'db> CallableBinding<'db> {
             }
             MatchingOverloadIndex::Single(_) => {
                 // If only one overload evaluates without error, it is the winning match.
-                return;
+                return None;
             }
             MatchingOverloadIndex::Multiple(indexes) => {
                 // If two or more candidate overloads remain, proceed to step 4.
@@ -1318,7 +1329,7 @@ impl<'db> CallableBinding<'db> {
                 self.filter_overloads_using_any_or_unknown(db, argument_types.as_ref(), &indexes);
 
                 // We're returning here because this shouldn't lead to argument type expansion.
-                return;
+                return None;
             }
         }
 
@@ -1328,7 +1339,7 @@ impl<'db> CallableBinding<'db> {
 
         if expansions.peek().is_none() {
             // Return early if there are no argument types to expand.
-            return;
+            return None;
         }
 
         // State of the bindings _after_ evaluating (type checking) the matching overloads using
@@ -1377,18 +1388,24 @@ impl<'db> CallableBinding<'db> {
                     argument_type.display(db)
                 );
                 snapshotter.restore(self, post_evaluation_snapshot);
-                return;
+                return None;
             }
         }
 
+        tracing::debug!("Performing argument type expansion");
+
+        // Global argument_forms and conflicting_forms that will be merged across all expansions
+        let mut global_argument_forms: Vec<Option<ParameterForm>> = Vec::new();
+        let mut global_conflicting_forms: Vec<bool> = Vec::new();
+
         for expansion in expansions {
             let expanded_argument_lists = match expansion {
                 Expansion::LimitReached(index) => {
                     snapshotter.restore(self, post_evaluation_snapshot);
                     self.overload_call_return_type = Some(
                         OverloadCallReturnType::ArgumentTypeExpansionLimitReached(index),
                     );
-                    return;
+                    return Some((global_argument_forms, global_conflicting_forms));
                 }
                 Expansion::Expanded(argument_lists) => argument_lists,
             };
@@ -1400,11 +1417,41 @@ impl<'db> CallableBinding<'db> {
 
             let mut return_types = Vec::new();
 
-            for expanded_argument_types in &expanded_argument_lists {
+            for expanded_arguments in &expanded_argument_lists {
                 let pre_evaluation_snapshot = snapshotter.take(self);
 
+                // Clear the state of all overloads before re-evaluating from step 1
+                for overload in &mut self.overloads {
+                    overload.reset();
+                }
+
+                let mut argument_forms = vec![None; expanded_arguments.len()];
+                let mut conflicting_forms = vec![false; expanded_arguments.len()];
+
+                // The spec mentions that each expanded argument list should re-evaluate from step
+                // 2 which is the type checking step but we're re-evaluating from step 1. The tldr
+                // is that it allows ty to match the correct overload in case a variadic argument
+                // would expand into different number of arguments with each expansion. Refer to
+                // https://github.com/astral-sh/ty/issues/735 for more details.
+                for overload in &mut self.overloads {
+                    overload.match_parameters(
+                        db,
+                        expanded_arguments,
+                        &mut argument_forms,
+                        &mut conflicting_forms,
+                    );
+                }
+
+                // Merge argument_forms and conflicting_forms into global ones
+                Self::merge_argument_forms(
+                    &mut global_argument_forms,
+                    &mut global_conflicting_forms,
+                    &argument_forms,
+                    &conflicting_forms,
+                );
+
                 for (_, overload) in self.matching_overloads_mut() {
-                    overload.check_types(db, expanded_argument_types);
+                    overload.check_types(db, expanded_arguments);
                 }
 
                 let return_type = match self.matching_overload_index() {
@@ -1417,14 +1464,26 @@ impl<'db> CallableBinding<'db> {
 
                         self.filter_overloads_using_any_or_unknown(
                             db,
-                            expanded_argument_types,
+                            expanded_arguments,
                             &matching_overload_indexes,
                         );
 
                         Some(self.return_type())
                     }
                 };
 
+                tracing::debug!(
+                    "Return type after evaluating expanded arguments (`{}`): {}",
+                    expanded_arguments
+                        .iter_types()
+                        .map(|arg| arg.display(db).to_string())
+                        .collect::<Vec<_>>()
+                        .join(", "),
+                    return_type
+                        .map(|ty| ty.display(db).to_string())
+                        .unwrap_or_else(|| "None".to_string())
+                );
+
                 // This split between initializing and updating the merged evaluation state is
                 // required because otherwise it's difficult to differentiate between the
                 // following:
@@ -1468,7 +1527,7 @@ impl<'db> CallableBinding<'db> {
                         UnionType::from_elements(db, return_types),
                     ));
 
-                return;
+                return Some((global_argument_forms, global_conflicting_forms));
             }
         }
 
@@ -1477,6 +1536,44 @@ impl<'db> CallableBinding<'db> {
         // argument types. This is necessary because we restore the state to the pre-evaluation
         // snapshot when processing the expanded argument lists.
         snapshotter.restore(self, post_evaluation_snapshot);
+        None
+    }
+
+    /// Merge argument forms and conflicting forms into global ones.
+    fn merge_argument_forms(
+        global_argument_forms: &mut Vec<Option<ParameterForm>>,
+        global_conflicting_forms: &mut Vec<bool>,
+        local_argument_forms: &[Option<ParameterForm>],
+        local_conflicting_forms: &[bool],
+    ) {
+        // Resize global lists to match local length if needed
+        if global_argument_forms.len() < local_argument_forms.len() {
+            global_argument_forms.resize(local_argument_forms.len(), None);
+            global_conflicting_forms.resize(local_conflicting_forms.len(), false);
+        }
+
+        // Merge argument forms and conflicting forms
+        for (i, (local_form, local_conflict)) in local_argument_forms
+            .iter()
+            .zip(local_conflicting_forms.iter())
+            .enumerate()
+        {
+            // Update global argument form
+            if let Some(global_form) = &mut global_argument_forms[i] {
+                if let Some(local_form) = local_form {
+                    if *global_form != *local_form {
+                        // Different parameter forms, mark as conflicting
+                        global_conflicting_forms[i] = true;
+                        *global_form = *local_form; // Use the new form
+                    }
+                }
+            } else {
+                global_argument_forms[i] = *local_form;
+            }
+
+            // Update global conflicting form (true takes precedence)
+            global_conflicting_forms[i] |= local_conflict;
+        }
     }
 
     /// Filter overloads based on [`Any`] or [`Unknown`] argument types.
@@ -2048,6 +2145,14 @@ impl<'a, 'db> ArgumentMatcher<'a, 'db> {
         argument_type: Option<Type<'db>>,
         length: TupleLength,
     ) -> Result<(), ()> {
+        tracing::debug!("tuple length: {length:?}");
+        tracing::debug!(
+            "argument type: {}",
+            argument_type
+                .map(|ty| ty.display(db).to_string())
+                .unwrap_or_else(|| "<None>".to_string())
+        );
+
         let tuple = argument_type.map(|ty| ty.iterate(db));
         let mut argument_types = match tuple.as_ref() {
             Some(tuple) => Either::Left(tuple.all_elements().copied()),
@@ -2597,6 +2702,16 @@ impl<'db> Binding<'db> {
     pub(crate) fn errors(&self) -> &[BindingError<'db>] {
         &self.errors
     }
+
+    /// Resets the state of this binding to its initial state.
+    fn reset(&mut self) {
+        self.return_ty = Type::unknown();
+        self.specialization = None;
+        self.inherited_specialization = None;
+        self.argument_matches = Box::from([]);
+        self.parameter_tys = Box::from([]);
+        self.errors.clear();
+    }
 }
 
 #[derive(Clone, Debug)]