From d021b3cedbbff02172bffca8efb48aa8f44c2dc8 Mon Sep 17 00:00:00 2001
From: Andre Masella <andre@masella.name>
Date: Wed, 19 Jul 2023 16:34:33 -0400
Subject: [PATCH] Add unwind handle for Python flow

For nested control flows
---
 dispyatcher/cpython.py | 41 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 39 insertions(+), 2 deletions(-)

diff --git a/dispyatcher/cpython.py b/dispyatcher/cpython.py
index dcf6bea..b55b725 100644
--- a/dispyatcher/cpython.py
+++ b/dispyatcher/cpython.py
@@ -6,8 +6,8 @@
 
 import dispyatcher.general
 import llvmlite.ir.types
-from dispyatcher import ArgumentManagement, BaseTransferUnaryHandle, ControlFlow, DerefPointer, F, FlowState, Handle,\
-    Identity, Type, ReturnManagement
+from dispyatcher import ArgumentManagement, BaseTransferUnaryHandle, ControlFlow, DerefPointer, F, FlowState, Handle, \
+    Identity, Type, ReturnManagement, TemporaryValue
 from dispyatcher.accessors import GetElementPointer
 from dispyatcher.general import BaseIndirectFunction, CurrentProcessFunction, MachineType
 from dispyatcher.permute import implode_args
@@ -217,6 +217,43 @@ def handle_return(self) -> Tuple[Type, ReturnManagement]:
         return self.__ty, self.__transfer
 
 
+class CheckAndUnwind(Handle[PythonControlFlow]):
+    """
+    Wraps another handle, usually a function, and checks if a Python exception has been raised.
+
+    This serves as an adapter between arbitrary function call handles that don't return useful status codes and the flow
+    control mechanism.
+    """
+    __handle: Handle[PythonControlFlow]
+
+    def __init__(self, handle: Handle[PythonControlFlow]):
+        super().__init__()
+        self.__handle = handle
+
+    def __str__(self) -> str:
+        return f"CheckAndUnwind({self.__handle})"
+
+    def generate_handle_ir(self, flow: PythonControlFlow, args: Sequence[IRValue]) -> Union[TemporaryValue, IRValue]:
+        result = self.__handle.generate_handle_ir(flow, args)
+        occurred_fn = flow.use_native_function("PyErr_Occurred", PY_OBJECT_TYPE.machine_type(), ())
+        comparison = flow.builder.icmp_unsigned('==',
+                                                flow.builder.call(occurred_fn, ()),
+                                                PY_OBJECT_TYPE.machine_type()(None))
+        fail_block = flow.builder.append_basic_block('check_failed')
+        success_block = flow.builder.append_basic_block('check_ok')
+        flow.builder.cbranch(comparison, fail_block, success_block)
+        flow.builder.position_at_start(fail_block)
+        flow.unwind()
+        flow.builder.position_at_start(success_block)
+        return result
+
+    def handle_arguments(self) -> Sequence[Tuple[Type, ArgumentManagement]]:
+        return self.__handle.handle_arguments()
+
+    def handle_return(self) -> Tuple[Type, ReturnManagement]:
+        return self.__handle.handle_return()
+
+
 class CFunctionHandle(BaseIndirectFunction):
     """
     Wraps a ``ctypes`` generated function in a handle