fix: handle small try blocks (#145)

* fix: handle small try blocks

We make sure that we can handle small (and potentially empty) try
blocks. These are blocks where the end offset is smaller or at most
equal to the start offset. These blocks can occur when a single
instruction that requires EXTENDED_ARG opcodes is wrapped in a try
block. Original code objects seem to wrap the whole instruction,
including the EXTENDED_ARG opcodes. The fix in this PR wraps just
the instruction, leaving the extra EXTENDED_ARG opcodes outside the
try block. This might not be a problem in general, but it is perhaps
not ideal.

* document changes

doc/changelog.rst

@@ -1,6 +1,15 @@
 ChangeLog
 =========
 
+2024-05-28: Version 0.15.2
+--------------------------
+
+Bugfixes:
+
+- Ensure that empty or small (one-instruction) try blocks are handled without
+  problems when compiling and de-compiling abstract code for CPython 3.11 and
+  later. PR #145
+
 2023-10-13: Version 0.15.1
 --------------------------
 

src/bytecode/concrete.py

@@ -989,14 +989,19 @@ def to_bytecode(
             # Handle TryBegin pseudo instructions
             if offset in ex_start:
                 entry = ex_start[offset]
-                tb_instr = TryBegin(
-                    Label(),
-                    entry.push_lasti,
-                    entry.stack_depth if conserve_exception_block_stackdepth else UNSET,
-                )
-                # Per entry store the pseudo instruction associated
-                tb_instrs[entry] = tb_instr
-                instructions.append(tb_instr)
+                # Check if the try begin was already created by an entry
+                # with a end offset less or equal to the start offset.
+                if entry not in tb_instrs:
+                    tb_instr = TryBegin(
+                        Label(),
+                        entry.push_lasti,
+                        entry.stack_depth
+                        if conserve_exception_block_stackdepth
+                        else UNSET,
+                    )
+                    # Per entry store the pseudo instruction associated
+                    tb_instrs[entry] = tb_instr
+                    instructions.append(tb_instr)
 
             jump_target = c_instr.get_jump_target(offset)
             size = c_instr.size
@@ -1063,7 +1068,22 @@ def to_bytecode(
             if current_instr_offset in ex_end:
                 entries = ex_end[current_instr_offset]
                 for entry in reversed(entries):
-                    instructions.append(TryEnd(tb_instrs[entry]))
+                    try:
+                        instructions.append(TryEnd(tb_instrs[entry]))
+                    except KeyError:
+                        # The end offset is behind the start offset, so we
+                        # need to create
+                        tb_instr = TryBegin(
+                            Label(),
+                            entry.push_lasti,
+                            entry.stack_depth
+                            if conserve_exception_block_stackdepth
+                            else UNSET,
+                        )
+                        # Per entry store the pseudo instruction associated
+                        tb_instrs[entry] = tb_instr
+                        instructions.append(tb_instr)
+                        instructions.append(TryEnd(tb_instr))
 
         # Replace jump targets with labels
         for index, jump_target in jumps:

tests/test_bytecode.py

@@ -712,6 +712,72 @@ def recompile_code_and_inner(code):
                 while callable(f := f()):
                     pass
 
+    def test_empty_try_block(self):
+        if sys.version_info < (3, 11):
+            self.skipTest("Exception tables were introduced in 3.11")
+
+        import bytecode as b
+
+        def foo():
+            return 42
+
+        code = Bytecode.from_code(foo.__code__)
+
+        try_begin = b.TryBegin(Label(), push_lasti=True)
+        code[1:1] = [try_begin, b.TryEnd(try_begin), try_begin.target]
+
+        foo.__code__ = code.to_code()
+
+        # Test that the function is still good
+        self.assertEqual(foo(), 42)
+
+        # Test that we can re-decompile the code
+        code = Bytecode.from_code(foo.__code__)
+        foo.__code__ = code.to_code()
+
+        # Test that the function is still good
+        self.assertEqual(foo(), 42)
+
+        # Do another round trip
+        Bytecode.from_code(foo.__code__).to_code()
+
+    def test_try_block_around_extended_arg(self):
+        """Test that we can handle small try blocks around opcodes that require
+        extended arguments.
+
+        We wrap a jump instruction between a TryBegin and TryEnd, and ensure
+        that the jump target is further away as to require an extended argument
+        for the branching instruction. We then test that we can compile and
+        de-compile the code object without issues.
+        """
+        if sys.version_info < (3, 11):
+            self.skipTest("Exception tables were introduced in 3.11")
+
+        import bytecode as b
+
+        def foo():
+            return 42
+
+        bc = Bytecode.from_code(foo.__code__)
+
+        try_begin = b.TryBegin(Label(), push_lasti=True)
+        bc[1:1] = [
+            try_begin,
+            Instr("JUMP_FORWARD", try_begin.target),
+            b.TryEnd(try_begin),
+            *(Instr("NOP") for _ in range(400)),
+            try_begin.target,
+        ]
+
+        foo.__code__ = bc.to_code()
+
+        self.assertEqual(foo(), 42)
+
+        # Do another round trip
+        foo.__code__ = Bytecode.from_code(foo.__code__).to_code()
+
+        self.assertEqual(foo(), 42)
+
 
 if __name__ == "__main__":
     unittest.main()  # pragma: no cover