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 _ResolveInheritedMethod(self, ClassName: str, MethodName: str, Gen: "Translator.LlvmGen") -> tuple[str | None, "ir.Function | None"]: """沿继承链查找方法，返回 (定义方法的类名, 函数对象) 当子类未覆写 __enter__/__exit__ 等方法时，沿 Gen.class_parent 链查找父类实现。 _generate_inherited_method_wrappers 只为 class_methods[ClassName] 中的方法生成包装， 不包括父类独有方法，因此 with 语句需要单独处理继承查找。 """ # 先检查子类自身 FullName: str = f'{ClassName}.{MethodName}' func = Gen._find_function(FullName) if hasattr(Gen, '_find_function') else Gen.functions.get(FullName) if func: return (ClassName, func) # 沿继承链查找 p: str | None = Gen.class_parent.get(ClassName) if hasattr(Gen, 'class_parent') else None while p: parent_full: str = f'{p}.{MethodName}' parent_func = Gen._find_function(parent_full) if hasattr(Gen, '_find_function') else Gen.functions.get(parent_full) if parent_func: return (p, parent_func) p = Gen.class_parent.get(p) return (None, None) def _HandleWithLlvm(self, Node: ast.With) -> None: Gen: "Translator.LlvmGen" = self.Trans.LlvmGen for item in Node.items: context_expr: ast.expr = item.context_expr asname: str | None = None target: ast.expr | None = getattr(item, 'optional_vars', getattr(item, 'target', getattr(item, 'asname', None))) if target and isinstance(target, ast.Name): asname = target.id ClassName: str | None = 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: SymInfo: "SymbolTable.SymbolInfo" = self.Trans.SymbolTable.lookup(ClassName) if SymInfo: if SymInfo.IsStruct or SymInfo.IsRenum: pass else: ClassName = None else: ClassName = None elif isinstance(context_expr, ast.Name): VarName: str = context_expr.id if VarName in Gen.var_struct_class: ClassName = Gen.var_struct_class[VarName] ctx_val: ir.Value | None = self.HandleExprLlvm(context_expr) if not ctx_val: continue original_ctx_val: ir.Value = ctx_val if ClassName and isinstance(ctx_val.type, ir.PointerType): pointee: ir.Type = ctx_val.type.pointee if isinstance(pointee, (ir.LiteralStructType, ir.IdentifiedStructType)): found: tuple[str, ir.Type] | None = Gen.find_struct_by_pointee(pointee) if found: CN: str = found[0] ST: ir.Type = found[1] ClassName = CN StructPtrType: ir.PointerType | None = 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: ir.AllocaInstr = 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: ir.Value | None = Gen._var_to_heap_ptr.get(asname) if asname else None ctx_is_var_ref: bool = isinstance(context_expr, ast.Name) enter_result: ir.Value | None = None if ClassName: # 沿继承链查找 __enter__（子类可能未覆写，使用父类实现） enter_class, enter_func = self._ResolveInheritedMethod(ClassName, '__enter__', Gen) if enter_func is not None: ctx_Loaded: ir.Value = Gen._load(ctx_alloca, name="__with_ctx") # 父类方法需要 self 为父类指针类型，bitcast 子类指针 if enter_class != ClassName and enter_class in Gen.structs: parent_ptr_type = ir.PointerType(Gen.structs[enter_class]) if ctx_Loaded.type != parent_ptr_type: ctx_Loaded = Gen.builder.bitcast(ctx_Loaded, parent_ptr_type, name="cast_enter_self") enter_call: ir.CallInstr = Gen.builder.call(enter_func, [ctx_Loaded], name="call_enter") enter_result = enter_call if isinstance(enter_call.type, ir.PointerType): pointee: ir.Type = 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: tuple[str, ir.Type] | None = Gen.find_struct_by_pointee(enter_result.type.pointee) if found: CN: str = found[0] ST: ir.Type = found[1] 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: ir.Value = Gen._reg_values[asname] var: ir.AllocaInstr = 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: ir.Value = 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: ir.Value = Gen.builder.bitcast(enter_result, VarPtr.type.pointee, name=f"cast_{asname}") Gen._store(CastedVal, VarPtr) else: NewVar: ir.AllocaInstr = 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 # 如果类有 __provides__，自动生成 push 代码 if ClassName and ClassName in Gen.class_provides: # 编译期检查: with provider 栈深度是否超过 _MAX_PROVIDER_DEPTH (32) # _with_provides_stack 只反映当前嵌套路径，顺序执行的 with 会在退出时弹栈 current_depth: int = sum(len(pl) for pl in Gen._with_provides_stack) new_fields: int = len(Gen.class_provides[ClassName]) if current_depth + new_fields > 32: raise SyntaxError( f"[TransPyC] with provider stack overflow at line {Node.lineno}: " f"current depth {current_depth} + new {new_fields} fields > max 32. " f"Class '{ClassName}' provides too many fields or with nesting too deep." ) # 编译期压栈：记录当前 with 上下文提供的字段名，供 __require_must__ 静态检查 Gen._with_provides_stack.append(Gen.class_provides[ClassName]) ctx_Loaded: ir.Value = Gen._load(ctx_alloca, name="__with_ctx") ctx_i8: ir.Value = ctx_Loaded if isinstance(ctx_Loaded.type, ir.PointerType) and ctx_Loaded.type != ir.PointerType(ir.IntType(8)): ctx_i8 = Gen.builder.bitcast(ctx_Loaded, ir.PointerType(ir.IntType(8)), name="ctx_push_i8") for field_name in Gen.class_provides[ClassName]: field_str: ir.Value = Gen._create_string_global(field_name) push_func = Gen.functions.get('_push_provider') if push_func: Gen.builder.call(push_func, [field_str, ctx_i8], name="call_push_provider") self.Trans.VarScopes.append({}) if asname: self.Trans.VarScopes[-1][asname] = True self.HandleBodyLlvm(Node.body) if self.Trans.VarScopes: self.Trans.VarScopes.pop() # 如果类有 __provides__，自动生成 pop 代码 if ClassName and ClassName in Gen.class_provides: for _ in Gen.class_provides[ClassName]: pop_func = Gen.functions.get('_pop_provider') if pop_func: Gen.builder.call(pop_func, [], name="call_pop_provider") # 编译期弹栈 if Gen._with_provides_stack: Gen._with_provides_stack.pop() if ClassName: # 沿继承链查找 __exit__（子类可能未覆写，使用父类实现） exit_class, exit_func = self._ResolveInheritedMethod(ClassName, '__exit__', Gen) if exit_func is not None: ctx_Loaded: ir.Value = Gen._load(ctx_alloca, name="__with_ctx") # 父类方法需要 self 为父类指针类型，bitcast 子类指针 if exit_class != ClassName and exit_class in Gen.structs: parent_ptr_type = ir.PointerType(Gen.structs[exit_class]) if ctx_Loaded.type != parent_ptr_type: ctx_Loaded = Gen.builder.bitcast(ctx_Loaded, parent_ptr_type, name="cast_exit_self") Gen.builder.call(exit_func, [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: ir.Value = Gen.builder.ptrtoint(ctx_val, ir.IntType(64), name="ctx_int") enter_int: ir.Value = Gen.builder.ptrtoint(enter_result, ir.IntType(64), name="enter_int") same_ptr: ir.Value = Gen.builder.icmp_unsigned('==', ctx_int, enter_int, name="same_ptr_check") with_free_bb: ir.Block = Gen.func.append_basic_block("with_free_ctx") with_skip_bb: ir.Block = 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: ir.Value = Gen.builder.bitcast(ctx_val, ir.PointerType(ir.IntType(8)), name="with_free_cast") free_func: ir.Function = 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: ir.Value = Gen.builder.bitcast(ctx_val, ir.PointerType(ir.IntType(8)), name="with_free_cast") free_func: ir.Function = 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)