import t, c from stdint import * import ast import llvmlite import memhub import stdio import string import viperlib import lib.core.VLogger as VLogger import lib.core.Handles.HandlesBase as HandlesBase import lib.core.Handles.HandlesTranslator as HT import lib.core.Handles.HandlesVar as HandlesVar import lib.core.Handles.HandlesExpr as HandlesExpr import lib.core.Handles.HandlesNonlocal as HandlesNonlocal # ============================================================ # HandlesAugAssign - 增强赋值语句处理（Mixin 继承模式） # # 处理 += -= *= /= %= &= |= ^= <<= >>= # 流程: load target → apply binop → store result # # 支持 local/global/nonlocal 三种变量作用域 # ============================================================ @t.NoVTable class AugAssignHandle(HandlesBase.Mixin): """增强赋值处理器 (+=, -=, *=, etc.)：继承 Mixin 获得 Trans 回指针""" def __init__(self, trans: HT.Translator | t.CPtr): self.Trans = trans # ============================================================ # Handle - 处理 AugAssign 语句，返回新增变量数（始终为 0） # ============================================================ def Handle(self, node: ast.AST | t.CPtr) -> int: """翻译增强赋值语句 (x += 1, y -= 2, etc.)""" if node is None: return 0 aug: ast.AugAssign | t.CPtr = (ast.AugAssign | t.CPtr)(node) if aug is None: return 0 target: ast.AST | t.CPtr = aug.target if target is None: return 0 tk: int = target.kind() # Attribute 目标: self.field += 1 # 流程: get_attribute_ptr → load → binop → store if tk == ast.ASTKind.Attribute: pool: memhub.MemBuddy | t.CPtr = self.Trans.Pool builder: llvmlite.IRBuilder | t.CPtr = self.Trans._cur_builder mod: llvmlite.LLVMModule | t.CPtr = self.Trans.Module # 1. 获取字段指针 field_ptr: llvmlite.Value | t.CPtr = HandlesExpr.get_attribute_ptr( builder, pool, mod, target, self.Trans) if field_ptr is None: VLogger.error("attribute ptr is None", "AUGASGN") return 0 # 2. 确定字段类型并加载当前值 target_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool) if field_ptr.Ty is not None: target_ty = field_ptr.Ty.Pointee cur_val: llvmlite.Value | t.CPtr = llvmlite.build_load( builder, target_ty, field_ptr) if cur_val is None: VLogger.error("cannot load attribute", "AUGASGN") return 0 # 3. 翻译 RHS 值 rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value( builder, pool, mod, aug.value, None, 0, self.Trans) if rhs_val is None: VLogger.error("rhs is None", "AUGASGN") return 0 # 4. 应用二元运算 result: llvmlite.Value | t.CPtr = _apply_aug_op( pool, builder, aug.op, cur_val, rhs_val) if result is None: fb: t.CChar | t.CPtr = VLogger.fmt_buf() if fb is not None: viperlib.snprintf(fb, 1024, "binop failed for attr op=%d", aug.op) VLogger.error(fb, "AUGASGN") return 0 # 5. 类型对齐并存储 result = HandlesExpr.coerce_to_type(builder, result, target_ty) if result is None: return 0 llvmlite.build_store(builder, result, field_ptr) return 0 # Subscript 目标: self.state[i] += x / arr[i] += x # 流程: get_subscript_ptr → load → binop → store if tk == ast.ASTKind.Subscript: pool: memhub.MemBuddy | t.CPtr = self.Trans.Pool builder: llvmlite.IRBuilder | t.CPtr = self.Trans._cur_builder mod: llvmlite.LLVMModule | t.CPtr = self.Trans.Module # 1. 获取元素指针 elem_ptr: llvmlite.Value | t.CPtr = HandlesExpr.get_subscript_ptr( builder, pool, mod, target, self.Trans) if elem_ptr is None: VLogger.error("subscript ptr is None", "AUGASGN") return 0 # 2. 确定元素类型并加载当前值 target_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool) if elem_ptr.Ty is not None and elem_ptr.Ty.Pointee is not None: target_ty = elem_ptr.Ty.Pointee cur_val: llvmlite.Value | t.CPtr = llvmlite.build_load( builder, target_ty, elem_ptr) if cur_val is None: VLogger.error("cannot load subscript element", "AUGASGN") return 0 # 3. 翻译 RHS 值 rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value( builder, pool, mod, aug.value, None, 0, self.Trans) if rhs_val is None: VLogger.error("rhs is None", "AUGASGN") return 0 # 4. 应用二元运算 result: llvmlite.Value | t.CPtr = _apply_aug_op( pool, builder, aug.op, cur_val, rhs_val) if result is None: fb: t.CChar | t.CPtr = VLogger.fmt_buf() if fb is not None: viperlib.snprintf(fb, 1024, "binop failed for sub op=%d", aug.op) VLogger.error(fb, "AUGASGN") return 0 # 5. 类型对齐并存储 result = HandlesExpr.coerce_to_type(builder, result, target_ty) if result is None: return 0 llvmlite.build_store(builder, result, elem_ptr) return 0 if tk != ast.ASTKind.Name: VLogger.error("only Name/Attribute/Subscript target supported", "AUGASGN") return 0 nm: ast.Name | t.CPtr = (ast.Name | t.CPtr)(target) if nm is None or nm.id is None: return 0 pool: memhub.MemBuddy | t.CPtr = self.Trans.Pool builder: llvmlite.IRBuilder | t.CPtr = self.Trans._cur_builder mod: llvmlite.LLVMModule | t.CPtr = self.Trans.Module i32_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int32(pool) # 1. 确定变量作用域类型: 0=local, 1=global, 2=nonlocal scope_type: int = 0 if HT.is_global_name(self.Trans, nm.id) != 0: scope_type = 1 elif HT.is_nonlocal_name(self.Trans, nm.id) != 0: scope_type = 2 # 2. 加载当前值 cur_val: llvmlite.Value | t.CPtr = None target_alloca: llvmlite.Value | t.CPtr = None target_ty: llvmlite.LLVMType | t.CPtr = i32_ty if scope_type == 1: # global 变量 target_alloca = HandlesVar.lookup_module_var( self.Trans.SymTab, nm.id) if target_alloca is not None: if target_alloca.Ty is not None: target_ty = target_alloca.Ty.Pointee cur_val = llvmlite.build_load(builder, target_ty, target_alloca) elif scope_type == 2: # nonlocal 变量（通过闭包 env） cur_val = HandlesNonlocal.load_nonlocal_var(self.Trans, nm.id) if cur_val is not None: target_ty = cur_val.Ty else: # 普通局部变量 target_alloca = HandlesVar.lookup_var(self.Trans.SymTab, nm.id) if target_alloca is not None: if target_alloca.Ty is not None: target_ty = target_alloca.Ty.Pointee cur_val = llvmlite.build_load(builder, target_ty, target_alloca) if cur_val is None: fb: t.CChar | t.CPtr = VLogger.fmt_buf() if fb is not None: viperlib.snprintf(fb, 1024, "cannot load target %s", nm.id) VLogger.error(fb, "AUGASGN") return 0 # 3. 翻译 RHS 值 rhs_val: llvmlite.Value | t.CPtr = HandlesExpr.translate_value( builder, pool, mod, aug.value, None, 0, self.Trans) if rhs_val is None: VLogger.error("rhs is None", "AUGASGN") return 0 # 4. 类型对齐 + 应用二元运算 # 注意: AugAssign 不走运算符重载（语义上需要 __iadd__ 而非 __add__） result: llvmlite.Value | t.CPtr = _apply_aug_op( pool, builder, aug.op, cur_val, rhs_val) if result is None: fb: t.CChar | t.CPtr = VLogger.fmt_buf() if fb is not None: viperlib.snprintf(fb, 1024, "binop failed for op=%d", aug.op) VLogger.error(fb, "AUGASGN") return 0 # 5. 存储结果 result = HandlesExpr.coerce_to_type(builder, result, target_ty) if result is None: return 0 if scope_type == 1: # global 变量 if target_alloca is not None: llvmlite.build_store(builder, result, target_alloca) elif scope_type == 2: # nonlocal 变量 nl_ptr: llvmlite.Value | t.CPtr = HandlesNonlocal.get_nonlocal_var_ptr( self.Trans, nm.id) if nl_ptr is not None: llvmlite.build_store(builder, result, nl_ptr) else: # 普通局部变量 if target_alloca is not None: llvmlite.build_store(builder, result, target_alloca) return 0 # ============================================================ # _apply_aug_op - 应用增强赋值的二元运算 # # 支持指针算术: ptr += int / ptr -= int # 整数运算自动类型提升 # ============================================================ def _apply_aug_op(pool: memhub.MemBuddy | t.CPtr, builder: llvmlite.IRBuilder | t.CPtr, op: int, lhs: llvmlite.Value | t.CPtr, rhs: llvmlite.Value | t.CPtr) -> llvmlite.Value | t.CPtr: """应用增强赋值的二元运算（指针算术 + 整数运算）""" lhs_bits: int = HandlesExpr.get_llvm_type_bits(lhs.Ty) rhs_bits: int = HandlesExpr.get_llvm_type_bits(rhs.Ty) # 指针算术: ptr += int / ptr -= int if lhs_bits == 0 and rhs_bits != 0: if op == ast.OpKind.Add or op == ast.OpKind.Sub: i64_ty: llvmlite.LLVMType | t.CPtr = llvmlite.Int64(pool) ptr_as_int: llvmlite.Value | t.CPtr = llvmlite.build_ptrtoint(builder, lhs, i64_ty) int_val: llvmlite.Value | t.CPtr = HandlesExpr.coerce_to_type(builder, rhs, i64_ty) if ptr_as_int is None or int_val is None: return None if op == ast.OpKind.Add: result: llvmlite.Value | t.CPtr = llvmlite.build_add(builder, ptr_as_int, int_val) else: result = llvmlite.build_sub(builder, ptr_as_int, int_val) if result is None: return None return llvmlite.build_inttoptr(builder, result, lhs.Ty) return None # 整数运算：类型提升 if lhs_bits > rhs_bits and rhs_bits > 0: rhs = HandlesExpr.coerce_to_type(builder, rhs, lhs.Ty) elif rhs_bits > lhs_bits and lhs_bits > 0: lhs = HandlesExpr.coerce_to_type(builder, lhs, rhs.Ty) if op == ast.OpKind.Add: return llvmlite.build_add(builder, lhs, rhs) elif op == ast.OpKind.Sub: return llvmlite.build_sub(builder, lhs, rhs) elif op == ast.OpKind.Mult: return llvmlite.build_mul(builder, lhs, rhs) elif op == ast.OpKind.Div: return llvmlite.build_sdiv(builder, lhs, rhs) elif op == ast.OpKind.FloorDiv: return llvmlite.build_sdiv(builder, lhs, rhs) elif op == ast.OpKind.Mod: return llvmlite.build_srem(builder, lhs, rhs) elif op == ast.OpKind.BitAnd: return llvmlite.build_and(builder, lhs, rhs) elif op == ast.OpKind.BitOr: return llvmlite.build_or(builder, lhs, rhs) elif op == ast.OpKind.BitXor: return llvmlite.build_xor(builder, lhs, rhs) elif op == ast.OpKind.LShift: return llvmlite.build_shl(builder, lhs, rhs) elif op == ast.OpKind.RShift: return llvmlite.build_ashr(builder, lhs, rhs) return None # ============================================================ # NewAugAssignHandle - 工厂函数 # ============================================================ def NewAugAssignHandle(pool: memhub.MemBuddy | t.CPtr, trans: HT.Translator | t.CPtr) -> AugAssignHandle | t.CPtr: h: AugAssignHandle | t.CPtr = pool.alloc(AugAssignHandle.__sizeof__()) if h is None: return None string.memset(h, 0, AugAssignHandle.__sizeof__()) h.Trans = trans return h