138 lines
8.6 KiB
Python
138 lines
8.6 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
|
|
import ast
|
|
import llvmlite.ir as ir
|
|
|
|
|
|
class WithHandle(BaseHandle):
|
|
def _HandleWithLlvm(self, Node):
|
|
Gen = self.Trans.LlvmGen
|
|
for item in Node.items:
|
|
context_expr = item.context_expr
|
|
asname = None
|
|
target = getattr(item, 'optional_vars', getattr(item, 'target', getattr(item, 'asname', None)))
|
|
if target and isinstance(target, ast.Name):
|
|
asname = target.id
|
|
ClassName = None
|
|
if isinstance(context_expr, ast.Call):
|
|
if isinstance(context_expr.func, ast.Name):
|
|
ClassName = context_expr.func.id
|
|
elif isinstance(context_expr.func, ast.Attribute):
|
|
ClassName = context_expr.func.attr
|
|
if ClassName and ClassName not in Gen.structs:
|
|
if ClassName in self.Trans.SymbolTable:
|
|
SymInfo = self.Trans.SymbolTable[ClassName]
|
|
if SymInfo.IsStruct or SymInfo.IsRenum:
|
|
pass
|
|
else:
|
|
ClassName = None
|
|
else:
|
|
ClassName = None
|
|
elif isinstance(context_expr, ast.Name):
|
|
VarName = context_expr.id
|
|
if VarName in Gen.var_struct_class:
|
|
ClassName = Gen.var_struct_class[VarName]
|
|
ctx_val = self.HandleExprLlvm(context_expr)
|
|
if not ctx_val:
|
|
continue
|
|
original_ctx_val = ctx_val
|
|
if ClassName and isinstance(ctx_val.type, ir.PointerType):
|
|
pointee = ctx_val.type.pointee
|
|
if isinstance(pointee, (ir.LiteralStructType, ir.IdentifiedStructType)):
|
|
found = Gen.find_struct_by_pointee(pointee)
|
|
if found:
|
|
CN, ST = found
|
|
ClassName = CN
|
|
StructPtrType = ir.PointerType(Gen.structs[ClassName]) if ClassName and ClassName in Gen.structs else None
|
|
if StructPtrType and isinstance(ctx_val.type, ir.PointerType) and ctx_val.type.pointee != Gen.structs.get(ClassName):
|
|
ctx_val = Gen.builder.bitcast(ctx_val, StructPtrType, name="cast_with_ctx")
|
|
ctx_alloca = Gen._allocaEntry(ctx_val.type, name="__with_ctx")
|
|
Gen._store(ctx_val, ctx_alloca)
|
|
Gen._UnregisterTempPtr(ctx_val)
|
|
if ClassName:
|
|
Gen.var_struct_class['__with_ctx'] = ClassName
|
|
Gen._var_to_heap_ptr['__with_ctx'] = ctx_val
|
|
original_asname_heap_ptr = Gen._var_to_heap_ptr.get(asname) if asname else None
|
|
ctx_is_var_ref = isinstance(context_expr, ast.Name)
|
|
enter_result = None
|
|
if ClassName:
|
|
EnterFuncName = f'{ClassName}.__enter__'
|
|
if EnterFuncName in Gen.functions:
|
|
ctx_Loaded = Gen._load(ctx_alloca, name="__with_ctx")
|
|
enter_call = Gen.builder.call(Gen.functions[EnterFuncName], [ctx_Loaded], name="call_enter")
|
|
enter_result = enter_call
|
|
if isinstance(enter_call.type, ir.PointerType):
|
|
pointee = enter_call.type.pointee
|
|
if isinstance(pointee, ir.PointerType) and isinstance(pointee.pointee, (ir.LiteralStructType, ir.IdentifiedStructType)):
|
|
enter_result = Gen._load(enter_call, name="enter_deref")
|
|
elif self._is_char_pointer(enter_call) and StructPtrType:
|
|
enter_result = Gen.builder.bitcast(enter_call, StructPtrType, name="cast_enter")
|
|
if asname and enter_result:
|
|
if isinstance(enter_result.type, ir.PointerType) and isinstance(enter_result.type.pointee, (ir.LiteralStructType, ir.IdentifiedStructType)):
|
|
found = Gen.find_struct_by_pointee(enter_result.type.pointee)
|
|
if found:
|
|
CN, ST = found
|
|
Gen.var_struct_class[asname] = CN
|
|
Gen._var_to_heap_ptr[asname] = enter_result
|
|
Gen._register_local_heap_ptr(enter_result, VarName=asname)
|
|
if asname in Gen._reg_values:
|
|
OldVal = Gen._reg_values[asname]
|
|
var = Gen._allocaEntry(OldVal.type, name=asname)
|
|
Gen._store(OldVal, var)
|
|
Gen.variables[asname] = var
|
|
del Gen._reg_values[asname]
|
|
if asname in Gen.variables and Gen.variables[asname] is not None:
|
|
VarPtr = Gen.variables[asname]
|
|
if VarPtr.type.pointee == enter_result.type:
|
|
Gen._store(enter_result, VarPtr)
|
|
elif isinstance(VarPtr.type.pointee, ir.PointerType) and isinstance(enter_result.type, ir.PointerType):
|
|
CastedVal = Gen.builder.bitcast(enter_result, VarPtr.type.pointee, name=f"cast_{asname}")
|
|
Gen._store(CastedVal, VarPtr)
|
|
else:
|
|
NewVar = Gen._alloca(enter_result.type, name=asname)
|
|
Gen._store(enter_result, NewVar)
|
|
Gen.variables[asname] = NewVar
|
|
else:
|
|
Gen._reg_values[asname] = enter_result
|
|
Gen.variables[asname] = None
|
|
self.Trans.VarScopes.append({})
|
|
if asname:
|
|
self.Trans.VarScopes[-1][asname] = True
|
|
self.HandleBodyLlvm(Node.body)
|
|
if self.Trans.VarScopes:
|
|
self.Trans.VarScopes.pop()
|
|
if ClassName:
|
|
ExitFuncName = f'{ClassName}.__exit__'
|
|
if ExitFuncName in Gen.functions:
|
|
ctx_Loaded = Gen._load(ctx_alloca, name="__with_ctx")
|
|
Gen.builder.call(Gen.functions[ExitFuncName], [ctx_Loaded], name="call_exit")
|
|
Gen._unregister_local_heap_ptr(original_ctx_val)
|
|
if ctx_val is not original_ctx_val:
|
|
Gen._unregister_local_heap_ptr(ctx_val)
|
|
if ctx_is_var_ref and original_asname_heap_ptr and original_asname_heap_ptr is not original_ctx_val and original_asname_heap_ptr is not ctx_val:
|
|
Gen._unregister_local_heap_ptr(original_asname_heap_ptr)
|
|
if '__with_ctx' in Gen._var_to_heap_ptr:
|
|
del Gen._var_to_heap_ptr['__with_ctx']
|
|
if asname and enter_result is not None and isinstance(ctx_val.type, ir.PointerType) and isinstance(enter_result.type, ir.PointerType):
|
|
ctx_int = Gen.builder.ptrtoint(ctx_val, ir.IntType(64), name="ctx_int")
|
|
enter_int = Gen.builder.ptrtoint(enter_result, ir.IntType(64), name="enter_int")
|
|
same_ptr = Gen.builder.icmp_unsigned('==', ctx_int, enter_int, name="same_ptr_check")
|
|
with_free_bb = Gen.func.append_basic_block("with_free_ctx")
|
|
with_skip_bb = Gen.func.append_basic_block("with_skip_free")
|
|
Gen.builder.cbranch(same_ptr, with_skip_bb, with_free_bb)
|
|
Gen.builder.position_at_start(with_free_bb)
|
|
raw = Gen.builder.bitcast(ctx_val, ir.PointerType(ir.IntType(8)), name="with_free_cast")
|
|
free_func = Gen.get_or_declare_c_func('free', ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.IntType(8))]))
|
|
Gen.builder.call(free_func, [raw])
|
|
Gen.builder.branch(with_skip_bb)
|
|
Gen.builder.position_at_start(with_skip_bb)
|
|
Gen._unregister_local_heap_ptr(enter_result)
|
|
Gen._unregister_local_heap_ptr(original_ctx_val)
|
|
elif isinstance(ctx_val.type, ir.PointerType):
|
|
raw = Gen.builder.bitcast(ctx_val, ir.PointerType(ir.IntType(8)), name="with_free_cast")
|
|
free_func = Gen.get_or_declare_c_func('free', ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.IntType(8))]))
|
|
Gen.builder.call(free_func, [raw])
|
|
Gen._unregister_local_heap_ptr(original_ctx_val) |