Also check globals is exact dict for LOAD_GLOBAL fast path

youknowone · youknowone · commit 29c0cf2db933 · 2026-03-02T01:04:51.000+09:00
get_chain_exact bypasses __missing__ on dict subclasses.
Move get_chain_exact to PyExact&lt;PyDict&gt; impl with debug_assert,
and have get_chain delegate to it. Store builtins_dict as
Option&lt;&amp;PyExact&lt;PyDict&gt;&gt; to enforce exact type at compile time.

Use PyRangeIterator::next_fast() instead of pub(crate) fields.
Fix comment style issues.
diff --git a/crates/vm/src/builtins/dict.rs b/crates/vm/src/builtins/dict.rs
@@ -5,7 +5,7 @@ use super::{
 use crate::common::lock::LazyLock;
 use crate::object::{Traverse, TraverseFn};
 use crate::{
-    AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
+    AsObject, Context, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyRefExact, PyResult,
     TryFromObject, atomic_func,
     builtins::{
         PyTuple,
@@ -681,22 +681,29 @@ impl Py<PyDict> {
         let self_exact = self.exact_dict(vm);
         let other_exact = other.exact_dict(vm);
         if self_exact && other_exact {
-            self.entries.get_chain(&other.entries, vm, key)
+            // SAFETY: exact_dict checks passed
+            let self_exact = unsafe { PyExact::ref_unchecked(self) };
+            let other_exact = unsafe { PyExact::ref_unchecked(other) };
+            self_exact.get_chain_exact(other_exact, key, vm)
         } else if let Some(value) = self.get_item_opt(key, vm)? {
             Ok(Some(value))
         } else {
             other.get_item_opt(key, vm)
         }
     }
+}
 
-    /// Like `get_chain` but skips the exact_dict type checks. Use when both
-    /// dicts are known to be exact dict types (e.g. globals + builtins).
-    pub fn get_chain_exact<K: DictKey + ?Sized>(
+impl PyExact<PyDict> {
+    /// Look up `key` in `self`, falling back to `other`.
+    /// Both dicts must be exact `dict` types (enforced by `PyExact`).
+    pub(crate) fn get_chain_exact<K: DictKey + ?Sized>(
         &self,
         other: &Self,
         key: &K,
         vm: &VirtualMachine,
     ) -> PyResult<Option<PyObjectRef>> {
+        debug_assert!(self.class().is(vm.ctx.types.dict_type));
+        debug_assert!(other.class().is(vm.ctx.types.dict_type));
         self.entries.get_chain(&other.entries, vm, key)
     }
 }
diff --git a/crates/vm/src/builtins/range.rs b/crates/vm/src/builtins/range.rs
@@ -607,10 +607,23 @@ impl IterNext for PyLongRangeIterator {
 #[pyclass(module = false, name = "range_iterator")]
 #[derive(Debug)]
 pub struct PyRangeIterator {
-    pub(crate) index: AtomicCell<usize>,
-    pub(crate) start: isize,
-    pub(crate) step: isize,
-    pub(crate) length: usize,
+    index: AtomicCell<usize>,
+    start: isize,
+    step: isize,
+    length: usize,
+}
+
+impl PyRangeIterator {
+    /// Advance and return next value without going through the iterator protocol.
+    #[inline]
+    pub(crate) fn next_fast(&self) -> Option<isize> {
+        let index = self.index.fetch_add(1);
+        if index < self.length {
+            Some(self.start + (index as isize) * self.step)
+        } else {
+            None
+        }
+    }
 }
 
 impl PyPayload for PyRangeIterator {
diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs
@@ -1,8 +1,8 @@
 #[cfg(feature = "flame")]
 use crate::bytecode::InstructionMetadata;
 use crate::{
-    AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, PyStackRef, TryFromObject,
-    VirtualMachine,
+    AsObject, Py, PyExact, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, PyStackRef,
+    TryFromObject, VirtualMachine,
     builtins::{
         PyBaseException, PyBaseExceptionRef, PyCode, PyCoroutine, PyDict, PyDictRef, PyGenerator,
         PyInterpolation, PyList, PySet, PySlice, PyStr, PyStrInterned, PyTemplate, PyTraceback,
@@ -447,7 +447,14 @@ impl Py<Frame> {
             locals: &self.locals,
             globals: &self.globals,
             builtins: &self.builtins,
-            builtins_dict: self.builtins.downcast_ref_if_exact::<PyDict>(vm),
+            builtins_dict: if self.globals.class().is(vm.ctx.types.dict_type) {
+                self.builtins
+                    .downcast_ref_if_exact::<PyDict>(vm)
+                    // SAFETY: downcast_ref_if_exact already verified exact type
+                    .map(|d| unsafe { PyExact::ref_unchecked(d) })
+            } else {
+                None
+            },
             lasti: &self.lasti,
             object: self,
             state: &mut state,
@@ -530,9 +537,11 @@ struct ExecutingFrame<'a> {
     locals: &'a ArgMapping,
     globals: &'a PyDictRef,
     builtins: &'a PyObjectRef,
-    /// Cached downcast of builtins to PyDict. builtins never changes during
-    /// frame execution, so we avoid repeating the downcast on every LOAD_GLOBAL.
-    builtins_dict: Option<&'a Py<PyDict>>,
+    /// Cached downcast of builtins to PyDict for fast LOAD_GLOBAL.
+    /// Only set when both globals and builtins are exact dict types (not
+    /// subclasses), so that `__missing__` / `__getitem__` overrides are
+    /// not bypassed.
+    builtins_dict: Option<&'a PyExact<PyDict>>,
     object: &'a Py<Frame>,
     lasti: &'a PyAtomic<u32>,
     state: &'a mut FrameState,
@@ -3028,8 +3037,10 @@ impl ExecutingFrame<'_> {
     #[inline]
     fn load_global_or_builtin(&self, name: &Py<PyStr>, vm: &VirtualMachine) -> PyResult {
         if let Some(builtins_dict) = self.builtins_dict {
-            // Fast path: both globals (PyDictRef) and builtins are exact dicts
-            self.globals
+            // Fast path: both globals and builtins are exact dicts
+            // SAFETY: builtins_dict is only set when globals is also exact dict
+            let globals_exact = unsafe { PyExact::ref_unchecked(self.globals.as_ref()) };
+            globals_exact
                 .get_chain_exact(builtins_dict, name, vm)?
                 .ok_or_else(|| {
                     vm.new_name_error(format!("name '{name}' is not defined"), name.to_owned())
@@ -3720,13 +3731,10 @@ impl ExecutingFrame<'_> {
 
         // FOR_ITER_RANGE: bypass generic iterator protocol for range iterators
         if let Some(range_iter) = top.downcast_ref_if_exact::<PyRangeIterator>(vm) {
-            let index = range_iter.index.fetch_add(1);
-            if index < range_iter.length {
-                let value = range_iter.start + (index as isize) * range_iter.step;
+            if let Some(value) = range_iter.next_fast() {
                 self.push_value(vm.ctx.new_int(value).into());
                 return Ok(true);
             }
-            // Exhausted
             if vm.use_tracing.get() && !vm.is_none(&self.object.trace.lock()) {
                 let stop_exc = vm.new_stop_iteration(None);
                 self.fire_exception_trace(&stop_exc, vm)?;
@@ -4085,17 +4093,17 @@ impl ExecutingFrame<'_> {
             self.push_value(vm.ctx.new_bool(result).into());
             return Ok(None);
         }
-        // COMPARE_OP_FLOAT
+        // COMPARE_OP_FLOAT: leaf type, cannot recurse — skip rich_compare dispatch.
+        // Falls through on NaN (partial_cmp returns None) for correct != semantics.
         if let (Some(a_f), Some(b_f)) = (
             a.downcast_ref_if_exact::<PyFloat>(vm),
             b.downcast_ref_if_exact::<PyFloat>(vm),
         ) {
-            let result = a_f
-                .to_f64()
-                .partial_cmp(&b_f.to_f64())
-                .is_some_and(|ord| cmp_op.eval_ord(ord));
-            self.push_value(vm.ctx.new_bool(result).into());
-            return Ok(None);
+            if let Some(ord) = a_f.to_f64().partial_cmp(&b_f.to_f64()) {
+                let result = cmp_op.eval_ord(ord);
+                self.push_value(vm.ctx.new_bool(result).into());
+                return Ok(None);
+            }
         }
 
         let value = a.rich_compare(b, cmp_op, vm)?;
diff --git a/crates/vm/src/protocol/object.rs b/crates/vm/src/protocol/object.rs
@@ -277,16 +277,16 @@ impl PyObject {
 
     // Perform a comparison, raising TypeError when the requested comparison
     // operator is not supported.
-    // see: CPython PyObject_RichCompare / do_richcompare
+    // see: PyObject_RichCompare / do_richcompare
     #[inline] // called by ExecutingFrame::execute_compare with const op
     fn _cmp(
         &self,
         other: &Self,
         op: PyComparisonOp,
         vm: &VirtualMachine,
     ) -> PyResult<Either<PyObjectRef, bool>> {
-        // Single recursion guard for the entire comparison (matching CPython's
-        // Py_EnterRecursiveCallTstate placement in do_richcompare).
+        // Single recursion guard for the entire comparison
+        // (do_richcompare in Objects/object.c).
         vm.with_recursion("in comparison", || self._cmp_inner(other, op, vm))
     }
 
