Files
TransPyC/lib/core/Handles/HandlesAugAssign.py
2026-07-18 19:25:40 +08:00

224 lines
14 KiB
Python

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)