from __future__ import annotations from typing import TYPE_CHECKING if TYPE_CHECKING: from lib.core.translator import Translator from lib.core.Handles.HandlesBase import BaseHandle, EXCEPTION_CODE_MAP import ast import llvmlite.ir as ir class TryHandle(BaseHandle): def _get_exception_type_code(self, type_node: ast.expr) -> int: if isinstance(type_node, ast.Name): name: str = type_node.id if name in EXCEPTION_CODE_MAP: return EXCEPTION_CODE_MAP[name] if name in self.Trans.exception_registry: return self.Trans.exception_registry[name] return 1 if isinstance(type_node, ast.Attribute): name: str = type_node.attr if name in EXCEPTION_CODE_MAP: return EXCEPTION_CODE_MAP[name] if name in self.Trans.exception_registry: return self.Trans.exception_registry[name] return 1 if isinstance(type_node, ast.Tuple): for elt in type_node.elts: code: int = self._get_exception_type_code(elt) if code != 1: return code return 1 return 1 def _get_exception_name(self, type_node: ast.expr) -> str | None: if isinstance(type_node, ast.Name): return type_node.id if isinstance(type_node, ast.Attribute): return type_node.attr return None def _get_all_match_codes(self, type_node: ast.expr) -> list[int]: name: str | None = self._get_exception_name(type_node) if name is None: return [self._get_exception_type_code(type_node)] if name not in self.Trans.exception_registry: return [self._get_exception_type_code(type_node)] codes: list[int] = [self.Trans.exception_registry[name]] self._collect_descendant_codes(name, codes) return codes def _collect_descendant_codes(self, parent_name: str, codes: list[int]) -> None: for child_name, parent in self.Trans.exception_parents.items(): if parent == parent_name: child_code: int | None = self.Trans.exception_registry.get(child_name) if child_code is not None and child_code not in codes: codes.append(child_code) self._collect_descendant_codes(child_name, codes) def _HandleTryLlvm(self, Node: ast.Try) -> None: Gen: "Translator.LlvmGen" = self.Trans.LlvmGen if not Gen.func: return exception_code: ir.AllocaInstr = Gen._allocaEntry(ir.IntType(32), name="eh_code") eh_message: ir.AllocaInstr = Gen._allocaEntry(ir.PointerType(ir.IntType(8)), name="eh_message") Gen._store(ir.Constant(ir.IntType(32), 0), exception_code) Gen._store(ir.Constant(ir.PointerType(ir.IntType(8)), None), eh_message) HasFinally: bool = bool(Node.finalbody) if HasFinally: Gen.eh_finally_stack.append((None, None, None)) else: Gen.eh_finally_stack.append(None) DispatchBB: ir.Block = Gen.func.append_basic_block(name="try.dispatch") AfterBB: ir.Block = Gen.func.append_basic_block(name="try.after") except_blocks: list[tuple[ir.Block, list[int], ast.ExceptHandler]] = [] for i, handler in enumerate(Node.handlers): ExcTypeCodes: list[int] = [1] if handler.type: ExcTypeCodes = self._get_all_match_codes(handler.type) ExceptBB: ir.Block = Gen.func.append_basic_block(name=f"try.except_{i}") Gen.eh_except_block_stack.append((DispatchBB, None, exception_code, eh_message)) except_blocks.append((ExceptBB, ExcTypeCodes, handler)) TryBB: ir.Block = Gen.func.append_basic_block(name="try.body") ElseBB: ir.Block | None = None if Node.orelse: ElseBB = Gen.func.append_basic_block(name="try.else") Gen.builder.branch(TryBB) Gen.builder.position_at_start(TryBB) self.HandleBodyLlvm(Node.body) if not Gen.builder.block.is_terminated: Gen._store(ir.Constant(ir.IntType(32), 0), exception_code) if ElseBB: Gen.builder.branch(ElseBB) else: Gen.builder.branch(AfterBB) Gen.builder.position_at_start(DispatchBB) eh_code_val: ir.Value = Gen._load(exception_code, name="Load_eh_code") next_check_bb: ir.Block | None = None handler_var_map: dict[int, tuple[str, ir.AllocaInstr, ir.Value | None]] = {} UnhandledBB: ir.Block = Gen.func.append_basic_block(name="try.unhandled") for i, (ExceptBB, ExcTypeCodes, handler) in enumerate(except_blocks): MatchBB: ir.Block = Gen.func.append_basic_block(name=f"try.match_{i}") if i > 0 and next_check_bb is not None: Gen.builder.position_at_start(next_check_bb) eh_code_val = Gen._load(exception_code, name=f"Load_eh_code_{i}") if handler.type is None: Gen.builder.branch(MatchBB) else: match_cond: ir.Value | None = None for code in ExcTypeCodes: is_code_match: ir.Value = Gen.builder.icmp_signed('==', eh_code_val, ir.Constant(ir.IntType(32), code), name=f"eh_match_{i}_c{code}") if match_cond is None: match_cond = is_code_match else: match_cond = Gen.builder.or_(match_cond, is_code_match, name=f"eh_match_any_{i}") if i < len(except_blocks) - 1: next_check_bb = Gen.func.append_basic_block(name=f"try.check_{i}") else: next_check_bb = UnhandledBB Gen.builder.cbranch(match_cond, MatchBB, next_check_bb) Gen.builder.position_at_start(MatchBB) if handler.name: var_alloca: ir.AllocaInstr = Gen._allocaEntry(ir.IntType(8).as_pointer(), name=handler.name) eh_msg_val: ir.Value = Gen._load(eh_message, name=f"Load_eh_msg_{handler.name}") Gen._store(eh_msg_val, var_alloca) handler_var_map[i] = (handler.name, var_alloca, Gen.variables.get(handler.name)) Gen.builder.branch(ExceptBB) Gen.builder.position_at_start(UnhandledBB) eh_msg_out_arg: ir.Argument | None = None eh_code_out_arg: ir.Argument | None = None if Gen.func and len(Gen.func.args) > 0: for arg in Gen.func.args: if arg.name == '__eh_msg_out__': eh_msg_out_arg = arg elif arg.name == '__eh_code_out__': eh_code_out_arg = arg if eh_msg_out_arg is not None and eh_code_out_arg is not None: eh_msg_val: ir.Value = Gen._load(eh_message, name="Load_eh_msg_propagate") eh_code_val_prop: ir.Value = Gen._load(exception_code, name="Load_eh_code_propagate") Gen._store(eh_msg_val, eh_msg_out_arg) Gen._store(eh_code_val_prop, eh_code_out_arg) ret_type: ir.Type | None = getattr(Gen.func, 'ftype', None) and Gen.func.ftype.return_type if getattr(Gen, 'func', None) else None if isinstance(ret_type, ir.IdentifiedStructType): if ret_type.elements: zero_val: ir.Constant = ir.Constant(ret_type, [ ir.Constant(et, None) if isinstance(et, (ir.PointerType, ir.IdentifiedStructType, ir.LiteralStructType, ir.ArrayType)) else ir.Constant(et, 0) if isinstance(et, (ir.IntType, ir.FloatType, ir.DoubleType)) else ir.Constant(et, ir.Undefined) for et in ret_type.elements]) else: zero_val = ir.Constant(ret_type, None) Gen.builder.ret(zero_val) elif isinstance(ret_type, ir.PointerType): Gen.builder.ret(ir.Constant(ret_type, None)) elif isinstance(ret_type, (ir.FloatType, ir.DoubleType)): Gen.builder.ret(ir.Constant(ret_type, 0.0)) elif isinstance(ret_type, ir.IntType): Gen.builder.ret(ir.Constant(ret_type, 1)) elif isinstance(ret_type, ir.VoidType): Gen.builder.ret_void() else: Gen.builder.ret(ir.Constant(ir.IntType(32), 1)) else: Gen._store(ir.Constant(ir.IntType(32), 0), exception_code) Gen.builder.branch(AfterBB) for idx, (ExceptBB, ExcTypeCodes, handler) in enumerate(except_blocks): Gen.builder.position_at_start(ExceptBB) old_var: ir.Value | None = None if idx in handler_var_map: hname: str = handler_var_map[idx][0] var_alloca: ir.AllocaInstr = handler_var_map[idx][1] old_var = handler_var_map[idx][2] Gen.variables[hname] = var_alloca self.HandleBodyLlvm(handler.body) if idx in handler_var_map: hname = handler_var_map[idx][0] old_var = handler_var_map[idx][2] if old_var is not None: Gen.variables[hname] = old_var elif hname in Gen.variables: del Gen.variables[hname] if not Gen.builder.block.is_terminated: Gen.builder.branch(AfterBB) Gen.eh_except_block_stack.pop() if ElseBB: Gen.builder.position_at_start(ElseBB) self.HandleBodyLlvm(Node.orelse) if not Gen.builder.block.is_terminated: Gen.builder.branch(AfterBB) Gen.builder.position_at_start(AfterBB) if HasFinally: Gen.eh_finally_stack.pop() if Node.finalbody: FinallyBB: ir.Block = Gen.func.append_basic_block(name="try.finally") pending_ret_flag: ir.AllocaInstr = Gen._allocaEntry(ir.IntType(32), name="pending_ret_flag") pending_ret_val: ir.AllocaInstr | None = Gen._allocaEntry(ir.IntType(32), name="pending_ret_val") if Gen.func.type.pointee.return_type == ir.IntType(32) else None Gen.eh_finally_stack.append((FinallyBB, pending_ret_flag, pending_ret_val)) Gen.builder.branch(FinallyBB) Gen.builder.position_at_start(FinallyBB) Gen._store(ir.Constant(ir.IntType(32), 0), pending_ret_flag) self.HandleBodyLlvm(Node.finalbody) if Gen.eh_finally_stack: _, pending_ret_flag, pending_ret_val = Gen.eh_finally_stack[-1] Gen.eh_finally_stack.pop() ret_flag_val: ir.Value = Gen._load(pending_ret_flag, name="Load_ret_flag") ret_flag_is_set: ir.Value = Gen.builder.icmp_signed('!=', ret_flag_val, ir.Constant(ir.IntType(32), 0), name="check_ret_flag") RetBB: ir.Block = Gen.func.append_basic_block(name="finally.ret") ContinueBB: ir.Block = Gen.func.append_basic_block(name="try.continue") Gen.builder.cbranch(ret_flag_is_set, RetBB, ContinueBB) Gen.builder.position_at_start(RetBB) if pending_ret_val: ret_val: ir.Value = Gen._load(pending_ret_val, name="Load_ret_val") Gen.builder.ret(ret_val) else: Gen.builder.ret_void() Gen.builder.position_at_start(ContinueBB)