coroutine, generator, io, exceptions and CPython allocation/layout metadata; five jit fixes

extra_tests/test_itertools.py

@@ -13,3 +13,18 @@ def test_islice_maxint():
 def test_islice_largeint():
     slic = itertools.islice(itertools.count(), 1, 10, sys.maxsize - 20)
     assert len(list(slic)) == 1
+
+
+def test_generic_sequence_iterator_exhaustion_after_jit():
+    class Sequence:
+        def __init__(self, values):
+            self.values = values
+
+        def __getitem__(self, index):
+            return self.values[index]
+
+    total = 0
+    for _ in range(20):
+        for value in Sequence(range(1000)):
+            total += value
+    assert total == 9_990_000

majit/majit-backend-dynasm/src/aarch64/assembler.rs

@@ -2215,9 +2215,9 @@ impl<'a> AssemblerARM64<'a> {
         // opref_to_slot stores ABSOLUTE jitframe slots (user position +
         // JITFRAME_FIXED_SIZE) so slot_offset(slot) gives the correct byte
         // offset without further adjustment.
-        for iarg in inputargs {
+        for (position, iarg) in inputargs.iter().enumerate() {
             self.opref_to_slot
-                .insert(iarg.opref(), JITFRAME_FIXED_SIZE + iarg.index as usize);
+                .insert(iarg.opref(), JITFRAME_FIXED_SIZE + position);
         }
         // Also sync any frame allocations from regalloc's FrameManager.
         for (&opref, lifetime) in ra.longevity.lifetimes_iter() {

majit/majit-backend-dynasm/src/x86/assembler.rs

@@ -2848,9 +2848,9 @@ impl<'a> Assembler386<'a> {
         // opref_to_slot stores ABSOLUTE jitframe slots (user position +
         // JITFRAME_FIXED_SIZE) so slot_offset(slot) gives the correct byte
         // offset without further adjustment.
-        for iarg in inputargs {
+        for (position, iarg) in inputargs.iter().enumerate() {
             self.opref_to_slot
-                .insert(iarg.opref(), JITFRAME_FIXED_SIZE + iarg.index as usize);
+                .insert(iarg.opref(), JITFRAME_FIXED_SIZE + position);
         }
         // Also sync any frame allocations from regalloc's FrameManager.
         for (&opref, lifetime) in ra.longevity.lifetimes_iter() {

majit/majit-metainterp/src/blackhole.rs

@@ -1213,6 +1213,17 @@ impl BlackholeInterpreter {
         if opcode == self.op_catch_exception {
             return self.route_to_catch(position, exc_value);
         }
+        // A guard resume coordinate may name the successor block's entry
+        // `-live-`, before that block mirrors the virtualizable and reaches
+        // the raising operation's trailing `-live-`.  The flattener emits the
+        // `catch_exception` immediately after that trailing marker
+        // (`flatten.py:206-217`).  Walk forward to the first such marker and
+        // accept only its immediately-following catch; crossing any other
+        // operation after it means this exception belongs to no handler at
+        // the resumed call site.
+        if let Some(catch_pos) = self.find_catch_after_resume_live(resume_live_pos) {
+            return self.route_to_catch(catch_pos, exc_value);
+        }
         // Backward case (after-residual-call guard): pyre resumes the post-call
         // `GUARD_NO_EXCEPTION` at the next opcode's `-live-`
         // (`pc_map[fallthrough_pc]`, jitcode_dispatch.rs / capture_resumedata),
@@ -1365,6 +1376,48 @@ impl BlackholeInterpreter {
         None
     }
 
+    fn find_catch_after_resume_live(&self, resume_live_pos: usize) -> Option<usize> {
+        let code = &self.jitcode.code;
+        let startpoints = self.jitcode.startpoints.as_ref()?;
+        let mut points: Vec<usize> = startpoints
+            .iter()
+            .copied()
+            .filter(|&q| q > resume_live_pos)
+            .collect();
+        points.sort_unstable();
+        let mut crossed_trailing_live = false;
+        for q in points {
+            let op = code[q];
+            if op == self.op_catch_exception {
+                return crossed_trailing_live.then_some(q);
+            }
+            if crossed_trailing_live {
+                return None;
+            }
+            if op == self.op_live {
+                crossed_trailing_live = true;
+                continue;
+            }
+            if !matches!(
+                op,
+                majit_translate::insns::BC_SETFIELD_VABLE_I
+                    | majit_translate::insns::BC_SETFIELD_VABLE_R
+                    | majit_translate::insns::BC_SETFIELD_VABLE_F
+                    | majit_translate::insns::BC_SETARRAYITEM_VABLE_I
+                    | majit_translate::insns::BC_SETARRAYITEM_VABLE_R
+                    | majit_translate::insns::BC_SETARRAYITEM_VABLE_F
+            ) {
+                // Only the codewriter's successor-block virtualizable mirror
+                // stores may separate the guard resume coordinate from the
+                // raising operation's trailing live marker.  Crossing an
+                // arbitrary operation would attach its exception to the next
+                // operation's handler and silently swallow it.
+                return None;
+            }
+        }
+        None
+    }
+
     /// blackhole.py:424-439 handle_rvmprof_enter.
     pub fn handle_rvmprof_enter(&mut self) {
         let code = &self.jitcode.code;
@@ -4284,6 +4337,57 @@ mod tests {
             assert_eq!(bh.return_type, BhReturnType::Int);
         }
 
+        #[test]
+        fn test_guard_exception_resume_finds_catch_after_successor_sync() {
+            let mut asm = majit_translate::codewriter::assembler::Assembler::new();
+            let mut b = JitCodeBuilder::default();
+            b.load_const_r_value(0, 1);
+            b.load_const_i_value(0, 2);
+            let resume_pc = b.current_pos();
+            b.live(&mut asm, &[], &[], &[]);
+            // The normal-flow successor mirrors virtualizable state before
+            // the can-raise block's trailing live/catch pair.
+            b.vable_setfield_int_with_base(0, 0, 0);
+            b.live(&mut asm, &[], &[], &[]);
+            let handler_lbl = b.new_label();
+            b.catch_exception(handler_lbl);
+            b.load_const_i_value(2, 99);
+            b.int_return(2);
+            b.mark_label(handler_lbl);
+            let handler_pc = b.current_pos();
+            b.load_const_i_value(2, 42);
+            b.int_return(2);
+            let jitcode = b.finish();
+
+            let mut builder = super::build_inline_call_only_bh_builder();
+            let mut bh = builder.acquire_interp();
+            bh.setposition(std::sync::Arc::new(jitcode), resume_pc);
+
+            assert!(bh.handle_exception_in_frame(0xCAFE_F00D));
+            assert_eq!(bh.position, handler_pc);
+        }
+
+        #[test]
+        fn test_guard_exception_resume_does_not_cross_another_operation() {
+            let mut asm = majit_translate::codewriter::assembler::Assembler::new();
+            let mut b = JitCodeBuilder::default();
+            let resume_pc = b.current_pos();
+            b.live(&mut asm, &[], &[], &[]);
+            b.load_const_i_value(0, 1);
+            b.live(&mut asm, &[], &[], &[]);
+            let handler_lbl = b.new_label();
+            b.catch_exception(handler_lbl);
+            b.mark_label(handler_lbl);
+            b.int_return(0);
+            let jitcode = b.finish();
+
+            let mut builder = super::build_inline_call_only_bh_builder();
+            let mut bh = builder.acquire_interp();
+            bh.setposition(std::sync::Arc::new(jitcode), resume_pc);
+
+            assert!(!bh.handle_exception_in_frame(0xCAFE_F00D));
+        }
+
         thread_local! {
             /// Address of the interpreter's `exception_last_value` slot, read
             /// back by [`probe_exception_slot_at_record_time`].