""" DecoratorPass - Viper 语言自定义装饰器的 IR 层包装转换 (v4) v4 增强版特性： 1. 栈帧局部上下文 (ctx)：替代全局变量，解决并发/递归状态错乱 2. 流程劫持：handler 返回 i32 控制原函数调用次数 - 返回 0：跳过原函数调用（@cache, @mock） - 返回 1：正常调用一次（默认） - 返回 N > 1：循环调用 N 次（@repeat） 3. alwaysinline 属性：消除 wrapper 调用开销 4. 参数打包与修改：pre-phase 后从 args_struct 读回参数，装饰器可修改入参 5. 返回值修改：ret_ptr 可写，装饰器可在后置阶段修改返回值 6. 递归装饰保护：最外层 wrapper 检查全局标志位，递归调用时跳过所有装饰逻辑 装饰器调用约定 (v4)： i32 decor_name(i8* ctx, i8* func_name, i32 phase, i8* args_ptr, i8* ret_ptr, i8* decor_args_ptr) 参数说明： ctx : i8* - 栈帧局部上下文缓冲区（32字节），每次调用独立分配 前置阶段可写入状态，后置阶段可读取，生命周期绑定本次调用 多线程/递归完全隔离，替代全局变量 func_name : i8* - 原始函数名字符串 phase : i32 - 执行阶段（0=前置, 1=后置） args_ptr : i8* - 指向函数参数打包结构体（可读写，无参数时为 null） 前置阶段可写入修改参数值，wrapper 会从结构体读回修改后的参数 ret_ptr : i8* - 指向返回值（前置阶段为 null，void 函数始终为 null） 后置阶段可读写，装饰器可修改返回值 decor_args_ptr : i8* - 指向装饰器参数打包结构体（无参数装饰器为 null） 返回值（仅前置阶段有效，后置阶段忽略）： 0 : 跳过原函数调用（@cache, @mock, @skip） 1 : 调用原函数一次（默认行为，@log, @timing, @trace） N>1: 循环调用原函数 N 次（@repeat, @benchmark） 编译器生成的 wrapper（以 add(i32, i32) -> i32, @log 为例）： define i32 @__decor_wrap_add(i32 %a, i32 %b) alwaysinline { entry: ; 递归保护：检查标志位 %rec_val = Load i8, i8* @__decor_rec___decor_wrap_add %is_rec = icmp ne i8 %rec_val, 0 br i1 %is_rec, label %recursive_call, label %decorated_entry recursive_call: ; 递归调用：跳过所有装饰逻辑，直接调用原函数 %rec_result = call i32 @add(i32 %a, i32 %b) ret i32 %rec_result decorated_entry: ; 设置递归保护标志 store i8 1, i8* @__decor_rec___decor_wrap_add ; 分配栈帧局部上下文 %ctx = alloca [32 x i8] %ctx_ptr = bitcast [32 x i8]* %ctx to i8* ; 打包函数参数到结构体（装饰器可通过 args_ptr 修改） %args = alloca {i32, i32} store i32 %a, i32* getelementptr({i32, i32}* %args, i32 0, i32 0) store i32 %b, i32* getelementptr({i32, i32}* %args, i32 0, i32 1) %args_ptr = bitcast {i32, i32}* %args to i8* ; 分配返回值存储（零初始化） %ret = alloca i32 store i32 0, i32* %ret ; 前置阶段：handler 返回调用次数 %n = call i32 @log(i8* %ctx_ptr, i8* "add", i32 0, i8* %args_ptr, i8* null, i8* null) br label %loop.header loop.header: %i = phi i32 [0, %decorated_entry], [%i.next, %loop.body] %cond = icmp slt i32 %i, %n br i1 %cond, label %loop.body, label %post loop.body: ; 从结构体读回参数（装饰器可能已修改） %a.Loaded = Load i32, i32* getelementptr({i32, i32}* %args, i32 0, i32 0) %b.Loaded = Load i32, i32* getelementptr({i32, i32}* %args, i32 0, i32 1) %result = call i32 @add(i32 %a.Loaded, i32 %b.Loaded) store i32 %result, i32* %ret %i.next = add i32 %i, 1 br label %loop.header post: %ret_ptr = bitcast i32* %ret to i8* call i32 @log(i8* %ctx_ptr, i8* "add", i32 1, i8* %args_ptr, i8* %ret_ptr, i8* null) ; 清除递归保护标志 store i8 0, i8* @__decor_rec___decor_wrap_add %final = Load i32, i32* %ret ret i32 %final } 链式装饰器（从下到上嵌套）： @log @timing def f(x) -> int: 等价于 log(timing(f)) 执行顺序：log_pre → timing_pre → f → timing_post → log_post 生成：__decor_wrap_f (log) → __decor_wrap_f_timing (timing) → f 递归保护仅在最外层 wrapper 生效，递归调用直接跳到原始 f """ from __future__ import annotations import llvmlite.ir as ir from typing import TYPE_CHECKING if TYPE_CHECKING: from lib.core.LlvmCodeGenerator import LlvmCodeGenerator # 栈帧局部上下文缓冲区大小（字节） _CTX_SIZE: int = 32 # 装饰器处理函数的 LLVM 签名：i32 (i8*, i8*, i32, i8*, i8*, i8*) _DECOR_HANDLER_PARAM_TYPES: ir.FunctionType | None = None # 延迟初始化 def _get_decor_handler_func_type() -> ir.FunctionType: """获取装饰器处理函数的 LLVM FunctionType: i32 (i8*, i8*, i32, i8*, i8*, i8*)""" global _DECOR_HANDLER_PARAM_TYPES if _DECOR_HANDLER_PARAM_TYPES is None: i8_ptr: ir.PointerType = ir.IntType(8).as_pointer() _DECOR_HANDLER_PARAM_TYPES = ir.FunctionType( ir.IntType(32), [i8_ptr, i8_ptr, ir.IntType(32), i8_ptr, i8_ptr, i8_ptr] ) return _DECOR_HANDLER_PARAM_TYPES def _find_or_declare_handler(module: ir.Module, handler_name: str) -> ir.Function: """在 module 中查找或声明装饰器处理函数""" f: ir.Function for f in module.functions: if f.name == handler_name: return f func_type: ir.FunctionType = _get_decor_handler_func_type() handler: ir.Function = ir.Function(module, func_type, name=handler_name) return handler def _make_string_constant(module: ir.Module, text: str, name_prefix: str = "decor.str") -> ir.Constant: """在 module 中创建一个全局字符串常量，返回 i8* 指针""" counter: int = getattr(_make_string_constant, '_counter', 0) + 1 _make_string_constant._counter = counter const_name: str = f"{name_prefix}.{counter}" encoded: bytes = text.encode('utf-8') + b'\x00' const_type: ir.ArrayType = ir.ArrayType(ir.IntType(8), len(encoded)) const_val: ir.Constant = ir.Constant(const_type, bytearray(encoded)) global_var: ir.GlobalVariable = ir.GlobalVariable(module, const_type, name=const_name) global_var.global_constant = True global_var.linkage = 'private' global_var.initializer = const_val return global_var.bitcast(ir.IntType(8).as_pointer()) def _make_decor_args_constant(module: ir.Module, decorator_info: dict) -> ir.Constant: """ 为带参数的装饰器生成全局常量结构体，返回 i8* 指针。 无参数装饰器返回 null。 支持的参数类型： - int → i32 / i64 - float → f64 - bool → i1 - str → i8*（全局字符串常量） """ deco_args: list = decorator_info.get('args', []) deco_kwargs: dict = decorator_info.get('kwargs', {}) if not deco_args and not deco_kwargs: return ir.Constant(ir.IntType(8).as_pointer(), None) field_types: list = [] field_values: list = [] arg: int | str | float | bool for arg in deco_args: llvm_type: ir.Type | None llvm_val: ir.Constant | None llvm_type, llvm_val = _python_value_to_llvm(module, arg) if llvm_type is None: return ir.Constant(ir.IntType(8).as_pointer(), None) field_types.append(llvm_type) field_values.append(llvm_val) kw: str for kw in sorted(deco_kwargs.keys()): llvm_type: ir.Type | None llvm_val: ir.Constant | None llvm_type, llvm_val = _python_value_to_llvm(module, deco_kwargs[kw]) if llvm_type is None: return ir.Constant(ir.IntType(8).as_pointer(), None) field_types.append(llvm_type) field_values.append(llvm_val) struct_type: ir.LiteralStructType = ir.LiteralStructType(field_types) struct_const: ir.Constant = ir.Constant(struct_type, field_values) counter: int = getattr(_make_decor_args_constant, '_counter', 0) + 1 _make_decor_args_constant._counter = counter global_name: str = f"__decor_args.{counter}" global_var: ir.GlobalVariable = ir.GlobalVariable(module, struct_type, name=global_name) global_var.global_constant = True global_var.linkage = 'private' global_var.initializer = struct_const return global_var.bitcast(ir.IntType(8).as_pointer()) def _python_value_to_llvm(module: ir.Module, value: int | str | float | bool) -> tuple[ir.Type | None, ir.Constant | None]: """将 Python 值转换为 (LLVM类型, LLVM常量) 元组。""" if isinstance(value, bool): return ir.IntType(1), ir.Constant(ir.IntType(1), int(value)) elif isinstance(value, int): if -2**31 <= value < 2**31: return ir.IntType(32), ir.Constant(ir.IntType(32), value) else: return ir.IntType(64), ir.Constant(ir.IntType(64), value) elif isinstance(value, float): return ir.DoubleType(), ir.Constant(ir.DoubleType(), value) elif isinstance(value, str): str_const = _make_string_constant(module, value, name_prefix="decor.arg") return ir.IntType(8).as_pointer(), str_const else: return None, None def _ZeroConstant(llvm_type: ir.Type) -> ir.Constant | None: """为给定 LLVM 类型创建零值常量""" if isinstance(llvm_type, ir.IntType): return ir.Constant(llvm_type, 0) elif isinstance(llvm_type, (ir.FloatType, ir.DoubleType)): return ir.Constant(llvm_type, 0.0) elif isinstance(llvm_type, ir.PointerType): return ir.Constant(llvm_type, None) elif isinstance(llvm_type, ir.ArrayType): elem_zero: ir.Constant | None = _ZeroConstant(llvm_type.element) if elem_zero is None: return None return ir.Constant(llvm_type, [elem_zero] * llvm_type.count) else: return None def _generate_single_wrapper(module: ir.Module, original_func: ir.Function, wrapper_name: str, func_name_str: str, decorator_info: dict, is_export: bool = False, is_outermost: bool = False, true_original_func: ir.Function | None = None) -> ir.Function: """ 为单个装饰器生成一层 wrapper 函数。 v4 Wrapper 结构（最外层含递归保护）： entry → 递归保护检查（仅最外层） recursive_call → 递归时直接调用原函数，跳过所有装饰（仅最外层） decorated_entry → 设置递归标志，分配 ctx/args/ret，调用 handler 前置 loop.header → phi 计数器，比较 i < n_calls loop.body → 从 args_struct 读回参数（支持修改），调用原函数 post → 调用 handler 后置，清除递归标志，返回结果 """ decor_name: str = decorator_info['name'] handler: ir.Function = _find_or_declare_handler(module, decor_name) func_name_const: ir.Constant = _make_string_constant(module, func_name_str, name_prefix="decor.fn") func_type: ir.FunctionType = original_func.ftype return_type: ir.Type = func_type.return_type param_types: list[ir.Type] = [p.type for p in original_func.args] is_void: bool = isinstance(return_type, ir.VoidType) wrapper_type: ir.FunctionType = ir.FunctionType(return_type, param_types) wrapper: ir.Function = ir.Function(module, wrapper_type, name=wrapper_name) # 添加 alwaysinline 属性，优化器完全内联消除调用开销 wrapper.attributes.add('alwaysinline') i32: ir.IntType = ir.IntType(32) i8: ir.IntType = ir.IntType(8) i8_ptr: ir.PointerType = ir.IntType(8).as_pointer() phase_pre: ir.Constant = ir.Constant(i32, 0) phase_post: ir.Constant = ir.Constant(i32, 1) null_ptr: ir.Constant = ir.Constant(i8_ptr, None) zero_i32: ir.Constant = ir.Constant(i32, 0) one_i32: ir.Constant = ir.Constant(i32, 1) one_i8: ir.Constant = ir.Constant(i8, 1) zero_i8: ir.Constant = ir.Constant(i8, 0) # 递归保护：仅最外层 wrapper 生成 rec_flag: ir.GlobalVariable | None = None if is_outermost and true_original_func is not None: rec_flag_name: str = f"__decor_rec_{wrapper_name}" # 检查是否已存在 rec_flag = None g: ir.GlobalValue for g in module.global_values: if g.name == rec_flag_name: rec_flag = g break if rec_flag is None: rec_flag = ir.GlobalVariable(module, i8, name=rec_flag_name) rec_flag.global_constant = False rec_flag.linkage = 'internal' rec_flag.initializer = zero_i8 # 创建基本块 entry_block: ir.Block = wrapper.append_basic_block("entry") recursive_call_block: ir.Block | None decorated_entry_block: ir.Block | None if rec_flag is not None: recursive_call_block = wrapper.append_basic_block("recursive_call") decorated_entry_block = wrapper.append_basic_block("decorated_entry") else: recursive_call_block = None decorated_entry_block = None loop_header_block: ir.Block = wrapper.append_basic_block("loop.header") loop_body_block: ir.Block = wrapper.append_basic_block("loop.body") post_block: ir.Block = wrapper.append_basic_block("post") # ==== Entry block ==== builder: ir.IRBuilder = ir.IRBuilder(entry_block) loop_predecessor: ir.Block if rec_flag is not None: # 递归保护：检查标志位 rec_val: ir.LoadInstr = builder.load(rec_flag, name="rec_val") is_recursive: ir.ICMPInstr = builder.icmp_signed('!=', rec_val, zero_i8) builder.cbranch(is_recursive, recursive_call_block, decorated_entry_block) # ==== Recursive call block ==== builder = ir.IRBuilder(recursive_call_block) # 递归调用：直接调用真正的原函数，跳过所有装饰逻辑 rec_call_args: list[ir.Argument] = [arg for arg in wrapper.args] rec_ret_val: ir.CallInstr = builder.call(true_original_func, rec_call_args) if is_void: builder.ret_void() else: builder.ret(rec_ret_val) # ==== Decorated entry block ==== builder = ir.IRBuilder(decorated_entry_block) # 设置递归保护标志 builder.store(one_i8, rec_flag) # 记录 loop.header 的前置块为 decorated_entry_block loop_predecessor = decorated_entry_block else: loop_predecessor = entry_block # 1. 分配栈帧局部上下文（每次调用独立，线程安全，递归安全） ctx_type: ir.ArrayType = ir.ArrayType(ir.IntType(8), _CTX_SIZE) ctx_alloca: ir.AllocaInstr = builder.alloca(ctx_type, name="ctx") ctx_ptr: ir.BitCastInstr = builder.bitcast(ctx_alloca, i8_ptr) # 2. 打包函数参数到结构体（装饰器可通过 args_ptr 读写修改） args_alloca: ir.AllocaInstr | None = None args_struct_type: ir.LiteralStructType | None = None if param_types: args_struct_type = ir.LiteralStructType(param_types) args_alloca = builder.alloca(args_struct_type, name="args") zero: ir.Constant = ir.Constant(i32, 0) i: int arg: ir.Argument for i, arg in enumerate(wrapper.args): field_idx: ir.Constant = ir.Constant(i32, i) gep: ir.GEPInstr = builder.gep(args_alloca, [zero, field_idx], inbounds=True) builder.store(arg, gep) args_ptr: ir.BitCastInstr | ir.Constant = builder.bitcast(args_alloca, i8_ptr) else: args_ptr = null_ptr # 3. 分配返回值存储（零初始化，确保跳过原函数时返回安全默认值） ret_alloca: ir.AllocaInstr | None if not is_void: ret_alloca = builder.alloca(return_type, name="ret") zero_val: ir.Constant | None = _ZeroConstant(return_type) if zero_val is not None: builder.store(zero_val, ret_alloca) else: ret_alloca = None # 4. 获取装饰器参数指针 decor_args_ptr: ir.Constant = _make_decor_args_constant(module, decorator_info) # 5. 前置阶段：handler 返回调用次数 n_calls: ir.CallInstr = builder.call(handler, [ctx_ptr, func_name_const, phase_pre, args_ptr, null_ptr, decor_args_ptr]) # 6. 跳转到循环头 builder.branch(loop_header_block) # ==== Loop header block ==== builder = ir.IRBuilder(loop_header_block) i_phi: ir.PhiInstr = builder.phi(i32, "i") i_phi.add_incoming(zero_i32, loop_predecessor) cond: ir.ICMPInstr = builder.icmp_signed('<', i_phi, n_calls) builder.cbranch(cond, loop_body_block, post_block) # ==== Loop body block ==== builder = ir.IRBuilder(loop_body_block) # v4: 从 args_struct 读回参数（装饰器可能在 pre-phase 中修改了参数） call_args: list[ir.LoadInstr | ir.Argument] = [] if args_alloca is not None and args_struct_type is not None: zero = ir.Constant(i32, 0) for i in range(len(param_types)): field_idx: ir.Constant = ir.Constant(i32, i) gep: ir.GEPInstr = builder.gep(args_alloca, [zero, field_idx], inbounds=True) Loaded_arg: ir.LoadInstr = builder.load(gep, name=f"arg.{i}") call_args.append(Loaded_arg) ret_val: ir.CallInstr = builder.call(original_func, call_args) if ret_alloca is not None: builder.store(ret_val, ret_alloca) i_next: ir.AddInstr = builder.add(i_phi, one_i32) i_phi.add_incoming(i_next, loop_body_block) builder.branch(loop_header_block) # ==== Post block ==== builder = ir.IRBuilder(post_block) # 后置阶段 ret_ptr: ir.BitCastInstr | ir.Constant if ret_alloca is not None: ret_ptr = builder.bitcast(ret_alloca, i8_ptr) else: ret_ptr = null_ptr builder.call(handler, [ctx_ptr, func_name_const, phase_post, args_ptr, ret_ptr, decor_args_ptr]) # 递归保护：清除标志（必须在 post-phase 之后、return 之前） if rec_flag is not None: builder.store(zero_i8, rec_flag) # 返回 if is_void: builder.ret_void() else: final_ret: ir.LoadInstr = builder.load(ret_alloca, name="final_ret") builder.ret(final_ret) return wrapper def _redirect_calls(module: ir.Module, old_func: ir.Function, new_func: ir.Function, skip_func_names: set[str] | None = None) -> None: """ 替换 module 中所有对 old_func 的调用指令为对 new_func 的调用。 跳过 skip_func_names 中的函数（wrapper 函数内部不应被重定向）。 """ if skip_func_names is None: skip_func_names = set() func: ir.Function for func in module.functions: if func.name in skip_func_names: continue if func.name == new_func.name: continue block: ir.Block for block in func.blocks: instr: ir.Instruction for instr in list(block.instructions): if not isinstance(instr, ir.CallInstr): continue callee: ir.Function = instr.callee if isinstance(callee, ir.Function) and callee is old_func: instr.callee = new_func def run(Gen: LlvmCodeGenerator) -> None: """ 执行 DecoratorPass。 处理流程: 1. 遍历 _decorated_funcs 中所有带装饰标记的函数 2. 对每个函数，按装饰器列表从底到顶生成嵌套 wrapper 3. 将原函数改为 internal 链接 4. 替换所有调用点 """ if not Gen._decorated_funcs: return module: ir.Module = Gen.module decorated: dict = Gen._decorated_funcs.copy() mangled_name: str info: dict for mangled_name, info in decorated.items(): decorators: list = info['decorators'] func_name: str = info['func_name'] is_export: bool = info['is_export'] # 查找原始函数 original_func: ir.Function | None = None f: ir.Function for f in module.functions: if f.name == mangled_name: original_func = f break if original_func is None: continue # 原函数改为 internal 链接（仅 wrapper 调用） original_func.linkage = 'internal' # 按装饰器列表从底到顶生成嵌套 wrapper # Python 装饰器顺序：@a @b def f() => a(b(f)) # 执行顺序：先应用最靠近函数的 @b，再应用 @a current_func: ir.Function = original_func all_wrapper_names: set[str] = set() i: int for i in range(len(decorators) - 1, -1, -1): deco: dict = decorators[i] is_outermost: bool = (i == 0) deco_name: str = deco['name'] wrapper_name: str if is_outermost: wrapper_name = f"__decor_wrap_{mangled_name}" else: wrapper_name = f"__decor_wrap_{mangled_name}_{deco_name}" all_wrapper_names.add(wrapper_name) wrapper: ir.Function = _generate_single_wrapper( module, current_func, wrapper_name, func_name, deco, is_export=is_export, is_outermost=is_outermost, true_original_func=original_func if is_outermost else None ) current_func = wrapper # 更新 functions 映射，使原函数名指向最外层 wrapper Gen.functions[mangled_name] = current_func if func_name in Gen.functions: Gen.functions[func_name] = current_func # 替换所有非 wrapper 函数中对原函数的调用 _redirect_calls(module, original_func, current_func, all_wrapper_names)