snapshot before regression test
This commit is contained in:
t
547
lib/core/DecoratorPass.py
Normal file
547
lib/core/DecoratorPass.py
Normal file
@@ -0,0 +1,547 @@
|
||||
"""
|
||||
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)
|
||||
Reference in New Issue
Block a user