Preserve partial keyword identity through vectorcall (#9127)

* Preserve partial keyword identity through vectorcall

Retain keyword objects in CALL_FUNCTION_EX and snapshot partial state before Python callbacks. Match CPython 3.15 final partial keyword merging without changing Lib or tests.

Assisted-by: OpenAI Codex:model version unavailable

* Register partial vectorcall through pyslot

Use the existing slot macro instead of a class extension hook. Return TypeError for non-string incoming names in the non-vectorcall fallback rather than panicking.

Assisted-by: OpenAI Codex:model version unavailable

crates/vm/src/builtins/dict.rs

@@ -586,7 +586,10 @@ impl PyDict {
     /// `__eq__`) is raised by the operation itself and must propagate unchanged.
     /// Hashing up front also means a `__hash__` that fails only intermittently
     /// is still reported.
-    fn hash_or_unhashable<K: DictKey + ?Sized>(key: &K, vm: &VirtualMachine) -> PyResult<PyHash> {
+    pub(crate) fn hash_or_unhashable<K: DictKey + ?Sized>(
+        key: &K,
+        vm: &VirtualMachine,
+    ) -> PyResult<PyHash> {
         key_hash(key, "dict", vm)
     }
 

crates/vm/src/frame.rs

@@ -25,7 +25,7 @@ use crate::{
     convert::{ToPyObject, ToPyResult},
     coroutine::Coro,
     exceptions::ExceptionCtor,
-    function::{ArgMapping, Callee, Either, FuncArgs, KwArgs, PyMethodFlags},
+    function::{ArgMapping, Callee, Either, FuncArgs, PyMethodFlags},
     object::PyAtomicBorrow,
     object::{Traverse, TraverseFn},
     protocol::{PyIter, PyIterReturn},
@@ -4700,8 +4700,8 @@ impl ExecutingFrame<'_> {
             }
             Instruction::CallFunctionEx => {
                 // Stack: [callable, self_or_null, args_tuple, kwargs_or_null]
-                let args = self.collect_ex_args(vm)?;
-                self.execute_call(args, vm)
+                let (args, kwnames) = self.collect_ex_args(vm)?;
+                self.execute_call_ex(args, kwnames, false, vm)
             }
             Instruction::CallIntrinsic1 { func } => {
                 let value = self.pop_value();
@@ -8329,8 +8329,8 @@ impl ExecutingFrame<'_> {
                 self.execute_call_instrumented(args, vm)
             }
             Instruction::InstrumentedCallFunctionEx => {
-                let args = self.collect_ex_args(vm)?;
-                self.execute_call_instrumented(args, vm)
+                let (args, kwnames) = self.collect_ex_args(vm)?;
+                self.execute_call_ex(args, kwnames, true, vm)
             }
             Instruction::InstrumentedLoadSuperAttr => {
                 let oparg = bytecode::LoadSuperAttr::from(u32::from(arg));
@@ -9032,9 +9032,13 @@ impl ExecutingFrame<'_> {
         FuncArgs::with_kwargs_names(args, kwarg_names)
     }
 
-    fn collect_ex_args(&mut self, vm: &VirtualMachine) -> PyResult<FuncArgs> {
+    fn collect_ex_args(
+        &mut self,
+        vm: &VirtualMachine,
+    ) -> PyResult<(Vec<PyObjectRef>, Vec<PyObjectRef>)> {
         let kwargs_or_null = self.pop_value_opt();
-        let mut kwargs = KwArgs::default();
+        let mut kwnames = Vec::new();
+        let mut keyword_values = Vec::new();
         if let Some(kw_obj) = kwargs_or_null {
             // Stack: [callable, self_or_null, args_tuple]
             let callable = self.nth_value(2);
@@ -9043,16 +9047,17 @@ impl ExecutingFrame<'_> {
             Self::iterate_mapping_keys(vm, &kw_obj, &func_str, |key, value| {
                 // `PyStr`, not `PyUtf8Str`: CPython only checks that the key is a
                 // `str`, not that it is valid UTF-8, so surrogate keys are accepted.
-                let key_str = key
-                    .downcast_ref::<PyStr>()
+                key.downcast_ref::<PyStr>()
                     .ok_or_else(|| vm.new_type_error("keywords must be strings"))?;
-                kwargs.insert(key_str.as_wtf8().to_owned(), value);
+                // Preserve string-subclass identity until the callee binds or merges keywords.
+                kwnames.push(key);
+                keyword_values.push(value);
                 Ok(())
             })?
         };
 
         let args_obj = self.pop_value();
-        let args = if let Some(tuple) = args_obj.downcast_ref::<PyTuple>() {
+        let mut args = if let Some(tuple) = args_obj.downcast_ref::<PyTuple>() {
             tuple.as_slice().to_vec()
         } else {
             // Single *arg passed directly; convert to sequence at runtime.
@@ -9075,7 +9080,8 @@ impl ExecutingFrame<'_> {
                 }
             })?
         };
-        Ok(FuncArgs { args, kwargs })
+        args.extend(keyword_values);
+        Ok((args, kwnames))
     }
 
     /// Returns a display string for a callable object for use in error messages.
@@ -9267,22 +9273,34 @@ impl ExecutingFrame<'_> {
     }
 
     #[inline]
-    fn execute_call(&mut self, args: FuncArgs, vm: &VirtualMachine) -> FrameResult {
+    fn execute_call_ex(
+        &mut self,
+        mut args: Vec<PyObjectRef>,
+        kwnames: Vec<PyObjectRef>,
+        instrumented: bool,
+        vm: &VirtualMachine,
+    ) -> FrameResult {
         // Stack: [callable, self_or_null, ...]
         let self_or_null = self.pop_value_opt(); // Option<PyObjectRef>
         let callable = self.pop_value();
 
-        let final_args = if let Some(self_val) = self_or_null {
-            let mut args = args;
-            args.prepend_arg(self_val);
-            args
+        if let Some(self_val) = self_or_null {
+            args.insert(0, self_val);
+        }
+        let nargs = args.len() - kwnames.len();
+        let arg0 = if instrumented {
+            args[..nargs].first().cloned()
         } else {
-            args
+            None
         };
-
-        let value = callable.call(final_args, vm)?;
-        self.push_value(value);
-        Ok(None)
+        let invoke = || callable.vectorcall(args, nargs, Some(&kwnames), vm);
+        if instrumented {
+            self.execute_call_instrumented_inner(&callable, arg0, invoke, vm)
+        } else {
+            let value = invoke()?;
+            self.push_value(value);
+            Ok(None)
+        }
     }
 
     /// Instrumented version of execute_call: fires CALL, C_RETURN, and C_RAISE events.
@@ -9298,29 +9316,36 @@ impl ExecutingFrame<'_> {
             args
         };
 
+        let arg0 = final_args.args.first().cloned();
+        self.execute_call_instrumented_inner(&callable, arg0, || callable.call(final_args, vm), vm)
+    }
+
+    fn execute_call_instrumented_inner(
+        &mut self,
+        callable: &PyObject,
+        arg0: Option<PyObjectRef>,
+        invoke: impl FnOnce() -> PyResult,
+        vm: &VirtualMachine,
+    ) -> FrameResult {
         let is_python_call = callable.downcast_ref_if_exact::<PyFunction>(vm).is_some();
 
         // Fire CALL event
         let call_arg0 = if self.monitoring_mask & MonitoringEvent::Call.mask() != 0 {
-            let arg0 = final_args
-                .args
-                .first()
-                .cloned()
-                .unwrap_or_else(|| monitoring::get_missing(vm));
+            let arg0 = arg0.unwrap_or_else(|| monitoring::get_missing(vm));
             let offset = (self.lasti() - 1) * 2;
-            monitoring::fire_call(vm, self.code, offset, &callable, arg0.clone())?;
+            monitoring::fire_call(vm, self.code, offset, callable, arg0.clone())?;
             Some(arg0)
         } else {
             None
         };
 
-        match callable.call(final_args, vm) {
+        match invoke() {
             Ok(value) => {
                 if let Some(arg0) = call_arg0
                     && !is_python_call
                 {
                     let offset = (self.lasti() - 1) * 2;
-                    monitoring::fire_c_return(vm, self.code, offset, &callable, arg0)?;
+                    monitoring::fire_c_return(vm, self.code, offset, callable, arg0)?;
                 }
                 self.push_value(value);
                 Ok(None)
@@ -9330,7 +9355,7 @@ impl ExecutingFrame<'_> {
                     && !is_python_call
                 {
                     let offset = (self.lasti() - 1) * 2;
-                    match monitoring::fire_c_raise(vm, self.code, offset, &callable, arg0) {
+                    match monitoring::fire_c_raise(vm, self.code, offset, callable, arg0) {
                         Ok(()) => exc,
                         Err(monitor_exc) => monitor_exc,
                     }

crates/vm/src/stdlib/_functools.rs

@@ -5,7 +5,8 @@ mod _functools {
     use crate::{
         Context, Py, PyObject, PyObjectRef, PyPayload, PyResult, VirtualMachine,
         builtins::{
-            PyBoundMethod, PyDict, PyDictRef, PyGenericAlias, PyTuple, PyType, PyTypeRef, object,
+            PyBoundMethod, PyDict, PyDictRef, PyGenericAlias, PyStr, PyTuple, PyType, PyTypeRef,
+            object,
         },
         common::{hash::PyHash, lock::PyRwLock},
         function::{
@@ -277,6 +278,106 @@ mod _functools {
         flags(BASETYPE, HAS_DICT, HAS_WEAKREF)
     )]
     impl PyPartial {
+        #[pyslot]
+        fn vectorcall(
+            zelf: &PyObject,
+            args: Vec<PyObjectRef>,
+            nargs: usize,
+            kwnames: Option<&[PyObjectRef]>,
+            vm: &VirtualMachine,
+        ) -> PyResult {
+            let zelf = zelf.downcast_ref::<Self>().unwrap();
+            // Keep the call's state alive before keyword lookups can execute
+            // __hash__/__eq__ and replace it through __setstate__.
+            let func = zelf.func.load_owned();
+            let stored_args = zelf.args.load_owned();
+            let keywords = zelf.keywords.load_owned();
+            let phcount = zelf.phcount.load(Ordering::Relaxed);
+
+            // Check if we have enough args to fill placeholders
+            if nargs < phcount {
+                return Err(vm.new_type_error(format!(
+                    "missing positional arguments in 'partial' call; expected at least {phcount}, got {nargs}"
+                )));
+            }
+
+            // Build combined args, replacing placeholders
+            let mut combined_args = Vec::with_capacity(stored_args.as_slice().len() + args.len());
+            let mut new_args_iter = args[..nargs].iter();
+
+            for stored_arg in stored_args.as_slice() {
+                if is_placeholder(stored_arg) {
+                    // Replace placeholder with next new arg
+                    if let Some(new_arg) = new_args_iter.next() {
+                        combined_args.push(new_arg.clone());
+                    } else {
+                        // This shouldn't happen if phcount check passed
+                        combined_args.push(stored_arg.clone());
+                    }
+                } else {
+                    combined_args.push(stored_arg.clone());
+                }
+            }
+            // Append remaining new args
+            combined_args.extend(new_args_iter.cloned());
+            let combined_nargs = combined_args.len();
+
+            let kwnames = kwnames.unwrap_or_default();
+            if func.class().slots.vectorcall.load().is_none() {
+                // FuncArgs callables retain their existing keyword binding path.
+                // In particular, do not introduce vectorcall's hash callbacks
+                // when the wrapped callable does not support vectorcall.
+                let mut kwargs = KwArgs::default();
+                for (key, value) in keywords.items_vec() {
+                    let key = key
+                        .downcast_ref::<PyStr>()
+                        .ok_or_else(|| vm.new_type_error("keywords must be strings"))?;
+                    kwargs.insert(key.as_wtf8().to_owned(), value);
+                }
+                for (key, value) in kwnames.iter().zip(&args[nargs..]) {
+                    let key = key
+                        .downcast_ref::<PyStr>()
+                        .ok_or_else(|| vm.new_type_error("keywords must be strings"))?;
+                    kwargs.insert(key.as_wtf8().to_owned(), value.clone());
+                }
+                return func.call(FuncArgs::new(combined_args, kwargs), vm);
+            }
+            if keywords.is_empty() {
+                combined_args.extend_from_slice(&args[nargs..]);
+                return func.vectorcall(combined_args, combined_nargs, Some(kwnames), vm);
+            }
+
+            // Like partial_vectorcall, copy stored keywords only when an incoming
+            // key replaces one. New keys remain a tail, preserving their identity
+            // and avoiding extra hash callbacks from an intermediate dictionary.
+            let mut merged_keywords = None;
+            let mut tail_names = Vec::new();
+            let mut tail_values = Vec::new();
+            for (key, value) in kwnames.iter().zip(&args[nargs..]) {
+                let hash = PyDict::hash_or_unhashable(&**key, vm)?;
+                if keywords.entries.contains(vm, &**key, hash)? {
+                    let merged = merged_keywords.get_or_insert_with(|| keywords.copy());
+                    merged.inner_setitem(&**key, value.clone(), vm)?;
+                } else {
+                    tail_names.push(key.clone());
+                    tail_values.push(value.clone());
+                }
+            }
+
+            let merged = merged_keywords.as_ref().unwrap_or(&keywords);
+            let mut combined_names = Vec::with_capacity(merged.__len__() + tail_names.len());
+            for (key, value) in merged.items_vec() {
+                key.downcast_ref::<PyStr>()
+                    .ok_or_else(|| vm.new_type_error("keywords must be strings"))?;
+                combined_names.push(key);
+                combined_args.push(value);
+            }
+            combined_names.extend(tail_names);
+            combined_args.extend(tail_values);
+
+            func.vectorcall(combined_args, combined_nargs, Some(&combined_names), vm)
+        }
+
         #[pygetset]
         fn __dict__(zelf: &Py<Self>, vm: &VirtualMachine) -> PyDictRef {
             zelf.as_object()
@@ -534,60 +635,14 @@ mod _functools {
         type Args = FuncArgs;
 
         fn call(zelf: &Py<Self>, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
-            // Clone and release lock before calling Python code to prevent deadlock
-            let func = zelf.func.load_owned();
-            let stored_args = zelf.args.load_owned();
-            let keywords = zelf.keywords.load_owned();
-            let phcount = zelf.phcount.load(Ordering::Relaxed);
-
-            // Check if we have enough args to fill placeholders
-            if phcount > 0 && args.args.len() < phcount {
-                return Err(vm.new_type_error(format!(
-                    "missing positional arguments in 'partial' call; expected at least {}, got {}",
-                    phcount,
-                    args.args.len()
-                )));
-            }
-
-            // Build combined args, replacing placeholders
-            let mut combined_args =
-                Vec::with_capacity(stored_args.as_slice().len() + args.args.len());
-            let mut new_args_iter = args.args.iter();
-
-            for stored_arg in stored_args.as_slice() {
-                if is_placeholder(stored_arg) {
-                    // Replace placeholder with next new arg
-                    if let Some(new_arg) = new_args_iter.next() {
-                        combined_args.push(new_arg.clone());
-                    } else {
-                        // This shouldn't happen if phcount check passed
-                        combined_args.push(stored_arg.clone());
-                    }
-                } else {
-                    combined_args.push(stored_arg.clone());
-                }
-            }
-            // Append remaining new args
-            combined_args.extend(new_args_iter.cloned());
-
-            // Merge keywords from self.keywords and args.kwargs
-            let mut final_kwargs = crate::function::KwArgsMap::default();
-
-            // Add keywords from self.keywords
-            for (key, value) in &*keywords {
-                // `expect_str()` would panic on surrogate keys; keep them as WTF-8.
-                let key_str = key
-                    .downcast_ref::<crate::builtins::PyStr>()
-                    .ok_or_else(|| vm.new_type_error("keywords must be strings"))?;
-                final_kwargs.insert(key_str.as_wtf8().to_owned(), value);
-            }
-
-            // Add keywords from args.kwargs (these override self.keywords)
+            let nargs = args.args.len();
+            let mut values = args.args;
+            let mut kwnames = Vec::with_capacity(args.kwargs.len());
             for (key, value) in args.kwargs {
-                final_kwargs.insert(key, value);
+                kwnames.push(vm.ctx.new_str(key).into());
+                values.push(value);
             }
-
-            func.call(FuncArgs::new(combined_args, KwArgs::new(final_kwargs)), vm)
+            Self::vectorcall(zelf.as_object(), values, nargs, Some(&kwnames), vm)
         }
     }