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 getattr(SymInfo, 'IsStruct', False) or getattr(SymInfo, 'IsRenum', False): 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._alloca_entry(ctx_val.type, name="__with_ctx") Gen._store(ctx_val, ctx_alloca) Gen._unregister_temp_ptr(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._alloca_entry(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)