Add CallKwPy, CallKwBoundMethod, CallKwNonPy, BinaryOpSubscrListSlice specialization

youknowone · youknowone · commit 5ac7f954cef1 · 2026-03-02T16:40:55.000+09:00
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -1293,8 +1293,21 @@ impl ExecutingFrame<'_> {
                 self.execute_call(args, vm)
             }
             Instruction::CallKw { nargs } => {
+                let nargs = nargs.get(arg);
+                let instr_idx = self.lasti() as usize - 1;
+                let cache_base = instr_idx + 1;
+                let counter = self.code.instructions.read_cache_u16(cache_base);
+                if counter > 0 {
+                    unsafe {
+                        self.code
+                            .instructions
+                            .write_cache_u16(cache_base, counter - 1);
+                    }
+                } else {
+                    self.specialize_call_kw(vm, nargs, instr_idx, cache_base);
+                }
                 // Stack: [callable, self_or_null, arg1, ..., argN, kwarg_names]
-                let args = self.collect_keyword_args(nargs.get(arg));
+                let args = self.collect_keyword_args(nargs);
                 self.execute_call(args, vm)
             }
             Instruction::CallFunctionEx => {
@@ -4086,6 +4099,77 @@ impl ExecutingFrame<'_> {
                 let args = self.collect_positional_args(nargs);
                 self.execute_call(args, vm)
             }
+            Instruction::CallKwPy => {
+                let instr_idx = self.lasti() as usize - 1;
+                let cache_base = instr_idx + 1;
+                let cached_version = self.code.instructions.read_cache_u32(cache_base + 1);
+                let nargs: u32 = arg.into();
+                // Stack: [callable, self_or_null, arg1, ..., argN, kwarg_names]
+                // callable is at position nargs + 2 from top (nargs args + kwarg_names + self_or_null)
+                let callable = self.nth_value(nargs + 2);
+                if let Some(func) = callable.downcast_ref::<PyFunction>()
+                    && func.func_version() == cached_version
+                    && cached_version != 0
+                {
+                    let args = self.collect_keyword_args(nargs);
+                    let self_or_null = self.pop_value_opt();
+                    let callable = self.pop_value();
+                    let func = callable.downcast_ref::<PyFunction>().unwrap();
+                    let final_args = if let Some(self_val) = self_or_null {
+                        let mut args = args;
+                        args.prepend_arg(self_val);
+                        args
+                    } else {
+                        args
+                    };
+                    let result = func.invoke(final_args, vm)?;
+                    self.push_value(result);
+                    return Ok(None);
+                }
+                self.deoptimize_call_kw();
+                let args = self.collect_keyword_args(nargs);
+                self.execute_call(args, vm)
+            }
+            Instruction::CallKwBoundMethod => {
+                let instr_idx = self.lasti() as usize - 1;
+                let cache_base = instr_idx + 1;
+                let cached_version = self.code.instructions.read_cache_u32(cache_base + 1);
+                let nargs: u32 = arg.into();
+                // Stack: [callable, self_or_null(=self), arg1, ..., argN, kwarg_names]
+                let callable = self.nth_value(nargs + 2);
+                if let Some(func) = callable.downcast_ref::<PyFunction>()
+                    && func.func_version() == cached_version
+                    && cached_version != 0
+                {
+                    let args = self.collect_keyword_args(nargs);
+                    let self_val = self.pop_value(); // self_or_null is always Some here
+                    let callable = self.pop_value();
+                    let func = callable.downcast_ref::<PyFunction>().unwrap();
+                    let mut final_args = args;
+                    final_args.prepend_arg(self_val);
+                    let result = func.invoke(final_args, vm)?;
+                    self.push_value(result);
+                    return Ok(None);
+                }
+                self.deoptimize_call_kw();
+                let args = self.collect_keyword_args(nargs);
+                self.execute_call(args, vm)
+            }
+            Instruction::CallKwNonPy => {
+                let instr_idx = self.lasti() as usize - 1;
+                let cache_base = instr_idx + 1;
+                let cached_tag = self.code.instructions.read_cache_u32(cache_base + 1);
+                let nargs: u32 = arg.into();
+                let callable = self.nth_value(nargs + 2);
+                let callable_tag = callable as *const PyObject as u32;
+                if cached_tag == callable_tag {
+                    let args = self.collect_keyword_args(nargs);
+                    return self.execute_call(args, vm);
+                }
+                self.deoptimize_call_kw();
+                let args = self.collect_keyword_args(nargs);
+                self.execute_call(args, vm)
+            }
             Instruction::LoadSuperAttrAttr => {
                 let oparg = u32::from(arg);
                 let attr_name = self.code.names[(oparg >> 2) as usize];
@@ -6681,6 +6765,57 @@ impl ExecutingFrame<'_> {
         }
     }
 
+    fn specialize_call_kw(
+        &mut self,
+        _vm: &VirtualMachine,
+        nargs: u32,
+        instr_idx: usize,
+        cache_base: usize,
+    ) {
+        // Stack: [callable, self_or_null, arg1, ..., argN, kwarg_names]
+        // callable is at position nargs + 2 from top
+        let stack = &self.state.stack;
+        let stack_len = stack.len();
+        let self_or_null_is_some = stack[stack_len - nargs as usize - 2].is_some();
+        let callable = self.nth_value(nargs + 2);
+
+        if let Some(func) = callable.downcast_ref::<PyFunction>() {
+            let version = func.func_version();
+            if version == 0 {
+                unsafe {
+                    self.code
+                        .instructions
+                        .write_adaptive_counter(cache_base, ADAPTIVE_BACKOFF_VALUE);
+                }
+                return;
+            }
+
+            let new_op = if self_or_null_is_some {
+                Instruction::CallKwBoundMethod
+            } else {
+                Instruction::CallKwPy
+            };
+            unsafe {
+                self.code.instructions.replace_op(instr_idx, new_op);
+                self.code
+                    .instructions
+                    .write_cache_u32(cache_base + 1, version);
+            }
+            return;
+        }
+
+        // General fallback
+        let callable_tag = callable as *const PyObject as u32;
+        unsafe {
+            self.code
+                .instructions
+                .replace_op(instr_idx, Instruction::CallKwNonPy);
+            self.code
+                .instructions
+                .write_cache_u32(cache_base + 1, callable_tag);
+        }
+    }
+
     fn specialize_send(&mut self, instr_idx: usize, cache_base: usize) {
         // Stack: [receiver, val] — receiver is at position 1
         let receiver = self.nth_value(1);
@@ -6885,6 +7020,22 @@ impl ExecutingFrame<'_> {
         }
     }
 
+    fn deoptimize_call_kw(&mut self) {
+        let instr_idx = self.lasti() as usize - 1;
+        let cache_base = instr_idx + 1;
+        unsafe {
+            self.code.instructions.replace_op(
+                instr_idx,
+                Instruction::CallKw {
+                    nargs: Arg::marker(),
+                },
+            );
+            self.code
+                .instructions
+                .write_adaptive_counter(cache_base, ADAPTIVE_BACKOFF_VALUE);
+        }
+    }
+
     fn deoptimize_for_iter(&mut self) {
         let instr_idx = self.lasti() as usize - 1;
         let cache_base = instr_idx + 1;
