from __future__ import annotations from typing import TYPE_CHECKING import ast import sys import llvmlite.ir as ir if TYPE_CHECKING: from lib.core.translator import Translator from lib.core.VLogger import get_logger as _vlog from lib.core.Handles.HandlesBase import BaseHandle class AugAssignHandle(BaseHandle): _AUGOP_MAP: dict[type, str] = { ast.Add: '__iadd__', ast.Sub: '__isub__', ast.Mult: '__imul__', ast.Div: '__itruediv__', ast.FloorDiv: '__ifloordiv__', ast.Mod: '__imod__', ast.Pow: '__ipow__', } _AUGOP_FALLBACK: dict[type, str] = { ast.Add: '__add__', ast.Sub: '__sub__', ast.Mult: '__mul__', ast.Div: '__truediv__', ast.FloorDiv: '__floordiv__', ast.Mod: '__mod__', ast.Pow: '__pow__', } def _is_aug_unsigned(self, Node: ast.AugAssign, OldVal: ir.Value | None, Value: ir.Value) -> bool: """判断增量赋值操作是否应使用无符号指令""" Gen: "Translator.LlvmGen" = self.Trans.LlvmGen if isinstance(Node.target, ast.Name): if Gen._check_node_unsigned(Node.target): return True elif isinstance(Node.target, ast.Attribute): if Gen._check_node_unsigned(Node.target): return True elif isinstance(Node.target, ast.Subscript): if Gen._check_node_unsigned(Node.target): return True if Node.value is not None and Gen._check_node_unsigned(Node.value): return True return False def _HandleAugAssignLlvm(self, Node: ast.AugAssign) -> None: Gen: "Translator.LlvmGen" = self.Trans.LlvmGen if isinstance(Node.target, ast.Name): VarName: str = Node.target.id if (VarName not in Gen.variables and VarName not in Gen.global_vars and VarName not in Gen._reg_values and VarName not in Gen._direct_values): src_info: str = Gen._get_node_info() _vlog().error(f"变量 '{VarName}' 未声明，不能用于增量赋值（需用 global 声明或先声明局部变量）{src_info}") sys.exit(1) if VarName in Gen.var_const_flags: src_info: str = Gen._get_node_info() _vlog().error(f"不能对 const 变量 '{VarName}' 增量赋值{src_info}") sys.exit(1) ClassName: str | None = Gen.var_struct_class.get(VarName) iop_name: str | None = self._AUGOP_MAP.get(type(Node.op)) if ClassName and iop_name: iFuncName: str = f'{ClassName}.{iop_name}' FuncName: str = f'{ClassName}.{self._AUGOP_FALLBACK.get(type(Node.op), "")}' if iFuncName in Gen.functions or FuncName in Gen.functions: OldVal: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(ast.Name(id=VarName, ctx=ast.Load())) Value: ir.Value | None = self.HandleExprLlvm(Node.value) if OldVal and Value: result: ir.Value | None if iFuncName in Gen.functions: result = self.Trans.ExprHandler._try_operator_overLoad(ClassName, iop_name, OldVal, Gen, Value) else: result = self.Trans.ExprHandler._try_operator_overLoad(ClassName, self._AUGOP_FALLBACK[type(Node.op)], OldVal, Gen, Value) if result is not None: if VarName in Gen.variables and Gen.variables[VarName] is not None: VarPtr: ir.Value = Gen.variables[VarName] if VarPtr.type.pointee == result.type: Gen._store(result, VarPtr) elif isinstance(VarPtr.type.pointee, ir.PointerType) and isinstance(result.type, ir.PointerType): CastedVal: ir.Value = Gen.builder.bitcast(result, VarPtr.type.pointee, name=f"cast_{VarName}") Gen._store(CastedVal, VarPtr) else: NewVar: ir.AllocaInstr = Gen._allocaEntry(result.type, name=VarName) Gen._store(result, NewVar) Gen.variables[VarName] = NewVar else: Gen._reg_values[VarName] = result Gen.variables[VarName] = None return Value: ir.Value | None = self.HandleExprLlvm(Node.value) if not Value: return is_unsigned: bool = self._is_aug_unsigned(Node, None, Value) if isinstance(Node.target, ast.Name): VarName: str = Node.target.id Op: str = self.Trans.ExprHandler.GetOpSymbol(Node.op) if VarName in Gen.global_vars and VarName in Gen.module.globals: GVar: ir.GlobalVariable = Gen.module.globals[VarName] Gen.variables[VarName] = GVar OldVal: ir.Value = Gen._load(GVar, name=VarName) NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) Gen._store(NewVal, GVar) return if VarName in Gen._reg_values: OldVal: ir.Value = Gen._reg_values[VarName] saved_block: ir.Block = Gen.builder.block entry_block: ir.Block = Gen.func.entry_basic_block Gen.builder.position_at_start(entry_block) var: ir.AllocaInstr = Gen.builder.alloca(OldVal.type, name=VarName) Gen._store(OldVal, var) Gen.builder.position_at_end(saved_block) Gen.variables[VarName] = var del Gen._reg_values[VarName] if VarName in Gen._direct_values: OldVal: ir.Value = Gen._direct_values.pop(VarName) saved_block: ir.Block = Gen.builder.block entry_block: ir.Block = Gen.func.entry_basic_block Gen.builder.position_at_start(entry_block) var: ir.AllocaInstr = Gen.builder.alloca(OldVal.type, name=VarName) Gen._store(OldVal, var) Gen.builder.position_at_end(saved_block) Gen.variables[VarName] = var if Gen.builder.block.is_terminated: return if VarName in Gen.variables and Gen.variables[VarName] is not None: OldVal: ir.Value = Gen._load(Gen.variables[VarName], name=VarName) store_ptr: ir.Value = Gen.variables[VarName] if isinstance(OldVal.type, ir.PointerType) and isinstance(Value.type, ir.IntType): if Value.type.width != 64: Value = Gen.builder.zext(Value, ir.IntType(64), name="zext_ptr_offset") NewVal: ir.Value if Op == '-' or Op == '+': if Op == '-': NegValue: ir.Value = Gen.builder.neg(Value, name="neg_ptr_offset") NewVal = Gen.builder.gep(OldVal, [NegValue], name="ptr_sub") else: NewVal = Gen.builder.gep(OldVal, [Value], name="ptr_add") else: ptr_val: ir.Value = OldVal OldVal = Gen._load(OldVal, name="deref_ptr_for_op") orig_type: ir.Type = OldVal.type if isinstance(OldVal.type, ir.IntType) and isinstance(Value.type, ir.IntType): if OldVal.type.width < Value.type.width: OldVal = Gen.builder.zext(OldVal, Value.type, name="zext_left") elif OldVal.type.width > Value.type.width: Value = Gen.builder.zext(Value, OldVal.type, name="zext_right") NewVal = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) if isinstance(NewVal.type, ir.IntType) and isinstance(orig_type, ir.IntType): if NewVal.type.width > orig_type.width: NewVal = Gen.builder.trunc(NewVal, orig_type, name="trunc_result") store_ptr = ptr_val else: if isinstance(OldVal.type, ir.IntType) and isinstance(Value.type, ir.IntType): if OldVal.type.width < Value.type.width: OldVal = Gen.builder.zext(OldVal, Value.type, name="zext_left") elif OldVal.type.width > Value.type.width: Value = Gen.builder.zext(Value, OldVal.type, name="zext_right") elif isinstance(OldVal.type, ir.PointerType) and isinstance(Value.type, ir.IntType): if Op in ('+', '-'): if Value.type.width < 64: Value = Gen.builder.zext(Value, ir.IntType(64), name="zext_ptr_offset_aug") NewVal: ir.Value if Op == '-': NegValue: ir.Value = Gen.builder.neg(Value, name="neg_ptr_offset_aug") NewVal = Gen.builder.gep(OldVal, [NegValue], name="ptr_sub_aug") else: NewVal = Gen.builder.gep(OldVal, [Value], name="ptr_add_aug") Gen._store(NewVal, store_ptr) return NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) if isinstance(store_ptr.type, ir.PointerType): TargetType: ir.Type = store_ptr.type.pointee if NewVal.type != TargetType: Coerced: ir.Value | None = Gen._coerce_value(NewVal, TargetType) if Coerced is not None: NewVal = Coerced Gen._store(NewVal, store_ptr) elif VarName == 'pos': var: ir.AllocaInstr = Gen._allocaEntry(Value.type, name=VarName) Gen.variables[VarName] = var OldVal: ir.Value = Gen._load(var, name=VarName) NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) Gen._store(NewVal, var) elif isinstance(Node.target, ast.Attribute): AttrName: str = Node.target.attr OldVal: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(Node.target) if OldVal: def _deref_if_ptr(val: ir.Value) -> ir.Value: if isinstance(val.type, ir.PointerType): pointee: ir.Type = val.type.pointee if isinstance(pointee, (ir.IntType, ir.FloatType, ir.DoubleType)): return Gen._load(val, name="deref_aug") return val OldVal = _deref_if_ptr(OldVal) Op: str = self.Trans.ExprHandler.GetOpSymbol(Node.op) if isinstance(OldVal.type, ir.IntType) and isinstance(Value.type, ir.IntType): if OldVal.type.width < Value.type.width: OldVal = Gen.builder.zext(OldVal, Value.type, name="zext_left_attr") elif OldVal.type.width > Value.type.width: Value = Gen.builder.zext(Value, OldVal.type, name="zext_right_attr") NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) self.Trans.AssignHandler._HandleAttributeStoreLlvm(Node.target, NewVal) elif isinstance(Node.target, ast.Subscript): SubTarget: ast.Subscript = Node.target ElemPtr: ir.Value | None = self.Trans.ExprAttrHandle.HandleSubscriptPtrLlvm(SubTarget) if ElemPtr: if isinstance(ElemPtr.type, ir.PointerType): OldVal: ir.Value = Gen._load(ElemPtr, name="old_subscript_val") Op: str = self.Trans.ExprHandler.GetOpSymbol(Node.op) if isinstance(OldVal.type, ir.IntType) and isinstance(Value.type, ir.IntType): if OldVal.type.width < Value.type.width: OldVal = Gen.builder.zext(OldVal, Value.type, name="zext_left_sub") elif OldVal.type.width > Value.type.width: Value = Gen.builder.zext(Value, OldVal.type, name="zext_right_sub") elif isinstance(OldVal.type, ir.PointerType) and isinstance(Value.type, ir.IntType): if Op in ('+', '-'): if Value.type.width < 64: Value = Gen.builder.zext(Value, ir.IntType(64), name="zext_ptr_offset_sub") NewVal: ir.Value if Op == '-': NegValue: ir.Value = Gen.builder.neg(Value, name="neg_ptr_offset_sub") NewVal = Gen.builder.gep(OldVal, [NegValue], name="ptr_sub_sub") else: NewVal = Gen.builder.gep(OldVal, [Value], name="ptr_add_sub") Gen._store(NewVal, ElemPtr) return NewVal: ir.Value = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) if isinstance(NewVal.type, ir.IntType) and isinstance(ElemPtr.type.pointee, ir.IntType): if NewVal.type.width > ElemPtr.type.pointee.width: NewVal = Gen.builder.trunc(NewVal, ElemPtr.type.pointee, name="trunc_sub_result") Gen._store(NewVal, ElemPtr)