Add debug_assert to invoke_exact_args, lazy func_version reassignment

- Add debug_assert preconditions in invoke_exact_args - Add get_version_for_current_state() for lazy version reassignment after func_version invalidation - Document NEXT_TYPE_VERSION overflow policy
RustPython · youknowone · Mar 3, 2026 · Mar 2, 2026 · Mar 1, 2026 · Mar 1, 2026
commit 687e99f02b119b16d57b44fa470a732e44fe1c4e
diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs
@@ -603,6 +603,22 @@ impl Py<PyFunction> {
         self.func_version.load(Relaxed)
     }
 
+    /// _PyFunction_GetVersionForCurrentState
+    /// Returns the current version, assigning a fresh one if previously invalidated.
+    /// Returns 0 if the version counter has overflowed.
+    pub fn get_version_for_current_state(&self) -> u32 {
+        let v = self.func_version.load(Relaxed);
+        if v != 0 {
+            return v;
+        }
+        let new_v = FUNC_VERSION_COUNTER.fetch_add(1, Relaxed);
+        if new_v == 0 {
+            return 0; // Counter overflow
+        }
+        self.func_version.store(new_v, Relaxed);
+        new_v
+    }
+
     /// Check if this function is eligible for exact-args call specialization.
     /// Returns true if: no VARARGS, no VARKEYWORDS, no kwonly args, not generator/coroutine,
     /// and effective_nargs matches co_argcount.
@@ -627,6 +643,16 @@ impl Py<PyFunction> {
     pub fn invoke_exact_args(&self, args: &[PyObjectRef], vm: &VirtualMachine) -> PyResult {
         let code: PyRef<PyCode> = (*self.code).to_owned();
 
+        debug_assert_eq!(args.len(), code.arg_count as usize);
+        debug_assert!(code.flags.contains(bytecode::CodeFlags::NEWLOCALS));
+        debug_assert!(!code.flags.intersects(
+            bytecode::CodeFlags::VARARGS
+                | bytecode::CodeFlags::VARKEYWORDS
+                | bytecode::CodeFlags::GENERATOR
+                | bytecode::CodeFlags::COROUTINE
+        ));
+        debug_assert_eq!(code.kwonlyarg_count, 0);
+
         let frame = Frame::new(
             code.clone(),
             Scope::new(None, self.globals.clone()),

diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs
@@ -55,6 +55,10 @@ pub struct PyType {
     pub tp_version_tag: AtomicU32,
 }
 
+/// Monotonic counter for type version tags. Once it reaches `u32::MAX`,
+/// `assign_version_tag()` returns 0 permanently, disabling new inline-cache
+/// entries but not invalidating correctness (cache misses fall back to the
+/// generic path).
 static NEXT_TYPE_VERSION: AtomicU32 = AtomicU32::new(1);
 
 unsafe impl crate::object::Traverse for PyType {
@@ -199,7 +203,8 @@ fn is_subtype_with_mro(a_mro: &[PyTypeRef], a: &Py<PyType>, b: &Py<PyType>) -> b
 }
 
 impl PyType {
-    /// Assign a fresh version tag. Returns 0 on overflow (all caches invalidated).
+    /// Assign a fresh version tag. Returns 0 if the version counter has been
+    /// exhausted, in which case no new cache entries can be created.
     pub fn assign_version_tag(&self) -> u32 {
         loop {
             let current = NEXT_TYPE_VERSION.load(Ordering::Relaxed);

diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -4765,7 +4765,7 @@ impl ExecutingFrame<'_> {
         let callable = self.nth_value(nargs + 1);
 
         if let Some(func) = callable.downcast_ref::<PyFunction>() {
-            let version = func.func_version();
+            let version = func.get_version_for_current_state();
             if version == 0 {
                 unsafe {
                     self.code