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 AugAssignHandle(BaseHandle): _AUGOP_MAP = { ast.Add: '__iadd__', ast.Sub: '__isub__', ast.Mult: '__imul__', ast.Div: '__itruediv__', ast.FloorDiv: '__ifloordiv__', ast.Mod: '__imod__', ast.Pow: '__ipow__', } _AUGOP_FALLBACK = { 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, OldVal, Value): """判断增量赋值操作是否应使用无符号指令""" Gen = 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): Gen = self.Trans.LlvmGen if isinstance(Node.target, ast.Name): VarName = Node.target.id if VarName in Gen.var_const_flags: src_info = Gen._get_node_info() print(f"[错误] 不能对 const 变量 '{VarName}' 增量赋值{src_info}") import sys sys.exit(1) ClassName = Gen.var_struct_class.get(VarName) iop_name = self._AUGOP_MAP.get(type(Node.op)) if ClassName and iop_name: iFuncName = f'{ClassName}.{iop_name}' FuncName = f'{ClassName}.{self._AUGOP_FALLBACK.get(type(Node.op), "")}' if iFuncName in Gen.functions or FuncName in Gen.functions: OldVal = self.Trans.ExprHandler.HandleExprLlvm(ast.Name(id=VarName, ctx=ast.Load())) Value = self.HandleExprLlvm(Node.value) if OldVal and Value: 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 = 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 = Gen.builder.bitcast(result, VarPtr.type.pointee, name=f"cast_{VarName}") Gen._store(CastedVal, VarPtr) else: NewVar = 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 = self.HandleExprLlvm(Node.value) if not Value: return is_unsigned = self._is_aug_unsigned(Node, None, Value) if isinstance(Node.target, ast.Name): VarName = Node.target.id Op = self.Trans.ExprHandler.GetOpSymbol(Node.op) if VarName in Gen.global_vars and VarName in Gen.module.globals: GVar = Gen.module.globals[VarName] Gen.variables[VarName] = GVar OldVal = Gen._load(GVar, name=VarName) NewVal = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) Gen._store(NewVal, GVar) return if VarName in Gen._reg_values: OldVal = Gen._reg_values[VarName] saved_block = Gen.builder.block entry_block = Gen.func.entry_basic_block Gen.builder.position_at_start(entry_block) var = 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 = Gen._direct_values.pop(VarName) saved_block = Gen.builder.block entry_block = Gen.func.entry_basic_block Gen.builder.position_at_start(entry_block) var = 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 = Gen._load(Gen.variables[VarName], name=VarName) store_ptr = 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") if Op == '-' or Op == '+': if Op == '-': NegValue = 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 = OldVal OldVal = Gen._load(OldVal, name="deref_ptr_for_op") orig_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") if Op == '-': NegValue = 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 = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) if isinstance(store_ptr.type, ir.PointerType): TargetType = store_ptr.type.pointee if NewVal.type != TargetType: Coerced = Gen._coerce_value(NewVal, TargetType) if Coerced is not None: NewVal = Coerced Gen._store(NewVal, store_ptr) elif VarName == 'pos': var = Gen._allocaEntry(Value.type, name=VarName) Gen.variables[VarName] = var OldVal = Gen._load(var, name=VarName) NewVal = Gen.emit_binary_op(Op, OldVal, Value, is_unsigned=is_unsigned) Gen._store(NewVal, var) elif isinstance(Node.target, ast.Attribute): AttrName = Node.target.attr OldVal = self.Trans.ExprHandler.HandleExprLlvm(Node.target) if OldVal: def _deref_if_ptr(val): if isinstance(val.type, ir.PointerType): pointee = 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 = 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 = 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 = Node.target ElemPtr = self.Trans.ExprAttrHandle.HandleSubscriptPtrLlvm(SubTarget) if ElemPtr: if isinstance(ElemPtr.type, ir.PointerType): OldVal = Gen._load(ElemPtr, name="old_subscript_val") Op = 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") if Op == '-': NegValue = 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 = 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)