60 lines
3.3 KiB
Python
60 lines
3.3 KiB
Python
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 AssertHandle(BaseHandle):
|
|
def _HandleAssertLlvm(self, Node: ast.Assert) -> None:
|
|
Gen: "Translator.LlvmGen" = self.Trans.LlvmGen
|
|
cond_val: ir.Value | None = self.HandleExprLlvm(Node.test)
|
|
if not cond_val:
|
|
return
|
|
if not isinstance(cond_val.type, ir.IntType) or cond_val.type.width != 1:
|
|
zero: ir.Constant = ir.Constant(cond_val.type, 0)
|
|
cond_val = Gen.builder.icmp_signed('!=', cond_val, zero, name="assert_cond")
|
|
AssertFailBB: ir.Block = Gen.func.append_basic_block(name="assert.fail")
|
|
AssertOkBB: ir.Block = Gen.func.append_basic_block(name="assert.ok")
|
|
Gen.builder.cbranch(cond_val, AssertOkBB, AssertFailBB)
|
|
Gen.builder.position_at_start(AssertFailBB)
|
|
exc_code: int = 9
|
|
exc_msg: ir.Value | None = None
|
|
if Node.msg:
|
|
if isinstance(Node.msg, ast.Constant) and isinstance(Node.msg.value, str):
|
|
exc_msg = self.HandleExprLlvm(Node.msg)
|
|
elif isinstance(Node.msg, ast.Call) and isinstance(Node.msg.func, ast.Name):
|
|
exc_code = EXCEPTION_CODE_MAP.get(Node.msg.func.id, 99)
|
|
if Node.msg.args:
|
|
first_arg: ast.expr = Node.msg.args[0]
|
|
if isinstance(first_arg, ast.Constant) and isinstance(first_arg.value, str):
|
|
exc_msg = self.HandleExprLlvm(first_arg)
|
|
elif isinstance(Node.msg, ast.Name):
|
|
exc_code = EXCEPTION_CODE_MAP.get(Node.msg.id, 99)
|
|
else:
|
|
exc_msg = self.HandleExprLlvm(Node.msg)
|
|
if Gen.eh_except_block_stack:
|
|
ExceptBB: ir.Block = Gen.eh_except_block_stack[-1][0]
|
|
EndBB: ir.Block | None = Gen.eh_except_block_stack[-1][1]
|
|
exception_code: ir.AllocaInstr = Gen.eh_except_block_stack[-1][2]
|
|
eh_message: ir.AllocaInstr = Gen.eh_except_block_stack[-1][3]
|
|
Gen._store(ir.Constant(ir.IntType(32), exc_code), exception_code)
|
|
if exc_msg:
|
|
Gen._store(exc_msg, eh_message)
|
|
else:
|
|
null_ptr: ir.Constant = ir.Constant(ir.PointerType(ir.IntType(8)), None)
|
|
Gen._store(null_ptr, eh_message)
|
|
Gen.builder.branch(ExceptBB)
|
|
else:
|
|
printf_func: ir.Function = Gen.get_or_declare_c_func('printf', ir.FunctionType(ir.IntType(32), [ir.PointerType(ir.IntType(8))], var_arg=True))
|
|
if exc_msg:
|
|
fmt: ir.Value = Gen.emit_constant("AssertionError: %s\n", 'string')
|
|
Gen.builder.call(printf_func, [fmt, exc_msg], name="assert_fail_print")
|
|
else:
|
|
fmt: ir.Value = Gen.emit_constant("AssertionError\n", 'string')
|
|
Gen.builder.call(printf_func, [fmt], name="assert_fail_print")
|
|
exit_func: ir.Function = Gen.get_or_declare_c_func('exit', ir.FunctionType(ir.VoidType(), [ir.IntType(32)]))
|
|
Gen.builder.call(exit_func, [ir.Constant(ir.IntType(32), 1)], name="call_exit")
|
|
Gen.builder.unreachable()
|
|
Gen.builder.position_at_start(AssertOkBB) |