TermiNexus
ff7c0457ad
- 帧结构体增加参数槽与提升临时变量槽，create_<name> 带参数存入帧 - 嵌套 await（如 a = b + await f()）提升为顶层临时赋值 __await_N，帧预留槽 - 协程版 If/While：条件求值 + 分支内 await 就地挂起/恢复（动态状态号恢复块） - For 循环委托现有 handler（仅支持无 await/return，M3 扩展） - CoroTest 用例覆盖参数/嵌套 await/while+await/if+await（期望 coro result = 30）
764 lines
39 KiB
Python
764 lines
39 KiB
Python
from __future__ import annotations
|
||
from typing import TYPE_CHECKING, Any
|
||
if TYPE_CHECKING:
|
||
from lib.core.translator import Translator
|
||
from lib.core.LlvmCodeGenerator import LlvmCodeGenerator
|
||
import vast as ast
|
||
import llvmlite.ir as ir
|
||
from lib.core.Handles.HandlesBase import BaseHandle, CTypeInfo
|
||
from lib.includes import t
|
||
|
||
|
||
class AsyncFunctionHandle(BaseHandle):
|
||
"""语言内协程 M2：async def + await + coro_run 状态机翻译。
|
||
|
||
M2 在 M1 基础上扩展：
|
||
- async 函数参数：帧结构体加参数槽，create_<name> 带参数并存入帧。
|
||
- 嵌套 await 表达式：`a = b + await leaf()` 提升为顶层临时赋值
|
||
（`__await_N = await leaf(); a = b + __await_N`），帧预留临时变量槽。
|
||
- 控制流：If/While 语句（含内部 await 就地挂起/恢复）；For 仅支持无 await。
|
||
- resume 采用动态状态号：每个 await 点分配恢复块（state{N}），switch 分发。
|
||
|
||
帧布局：%Coro_<name> = type { i32 state, i32 done, i32 ret, <参数槽>, <局部槽>, <提升临时槽> }
|
||
状态机：create_<name>(args...) 分配帧 + 存参数；resume_<name>(frame) -> i32
|
||
（0=SUSPEND 1=DONE）按 state 分发到主入口 state0 与各恢复块。
|
||
"""
|
||
|
||
# ------------------------------------------------------------------
|
||
# 公共查询
|
||
# ------------------------------------------------------------------
|
||
|
||
def _IsAsyncFunc(self, name: str) -> bool:
|
||
Gen: LlvmCodeGenerator | None = getattr(self.Trans, 'LlvmGen', None)
|
||
if Gen is None:
|
||
return False
|
||
return name in getattr(Gen, '_async_funcs', {})
|
||
|
||
def _RegisterAsyncFunc(self, name: str, Node: ast.AsyncFunctionDef) -> None:
|
||
Gen: LlvmCodeGenerator | None = getattr(self.Trans, 'LlvmGen', None)
|
||
if Gen is None:
|
||
return
|
||
if not hasattr(Gen, '_async_funcs'):
|
||
Gen._async_funcs = {}
|
||
Gen._async_funcs[name] = {'node': Node}
|
||
|
||
def _GetAsyncCallName(self, call: ast.Call) -> str | None:
|
||
if isinstance(call.func, ast.Name):
|
||
name: str = call.func.id
|
||
if self._IsAsyncFunc(name):
|
||
return name
|
||
return None
|
||
|
||
def _IsCoroModule(self, name: str) -> bool:
|
||
if name == 'coro':
|
||
return True
|
||
aliases: Any = getattr(self.Trans.SymbolTable, 'import_aliases', {})
|
||
return aliases.get(name) == 'coro'
|
||
|
||
# ------------------------------------------------------------------
|
||
# 入口：async 函数定义
|
||
# ------------------------------------------------------------------
|
||
|
||
def _EmitAsyncFunctionLlvm(self, Node: ast.AsyncFunctionDef, Gen: LlvmCodeGenerator) -> None:
|
||
name: str = Node.name
|
||
# 防重复生成：模块可能被多次编译（泛型特化等触发），避免 create_/resume_ 重名
|
||
if f'create_{name}' in Gen.functions or name in getattr(Gen, '_coro_frame_types', {}):
|
||
return
|
||
self._RegisterAsyncFunc(name, Node)
|
||
# 预扫描嵌套 await，分配提升临时变量槽（帧布局需要）
|
||
hoist_count: int = self._PreScanHoistedVars(Node, Gen)
|
||
if not hasattr(Gen, '_coro_hoisted_count'):
|
||
Gen._coro_hoisted_count = {}
|
||
Gen._coro_hoisted_count[name] = hoist_count
|
||
# 保存外层 Gen 状态（create/resume 生成期间会切换 builder/func/variables）
|
||
saved_builder: ir.IRBuilder | None = Gen.builder
|
||
saved_func: Any = Gen.func
|
||
saved_variables: Any = Gen.variables
|
||
saved_reg_values: Any = Gen._reg_values
|
||
saved_direct_values: Any = Gen._direct_values
|
||
saved_var_signedness: Any = Gen.var_signedness
|
||
saved_var_struct_class: Any = Gen.var_struct_class
|
||
saved_var_type_info: Any = Gen.var_type_info
|
||
saved_var_type_assignments: Any = Gen.var_type_assignments
|
||
saved_var_scope_len: int = len(self.Trans.VarScopes)
|
||
saved_loop_len: int = len(Gen.loop_break_targets)
|
||
saved_cont_len: int = len(Gen.loop_continue_targets)
|
||
saved_coro_state = {
|
||
'state_ptr': getattr(Gen, '_coro_state_ptr', None),
|
||
'done_ptr': getattr(Gen, '_coro_done_ptr', None),
|
||
'ret_ptr': getattr(Gen, '_coro_ret_field_ptr', None),
|
||
'exit_bb': getattr(Gen, '_coro_exit_bb', None),
|
||
'next_state': getattr(Gen, '_coro_next_state', None),
|
||
'await_slot': getattr(Gen, '_coro_await_slot_ptr', None),
|
||
'frame_ptr': getattr(Gen, '_coro_frame_ptr', None),
|
||
'switch': getattr(Gen, '_coro_switch_instr', None),
|
||
'state_counter': getattr(Gen, '_coro_state_counter', None),
|
||
'hoist_counter': getattr(Gen, '_coro_hoist_counter', None),
|
||
}
|
||
try:
|
||
self.Trans.VarScopes.append({})
|
||
frame_ty: ir.IdentifiedStructType = self._GetOrCreateFrameType(name, Node, Gen, hoist_count)
|
||
Gen._coro_hoist_counter = 0
|
||
self._EmitCreateFunction(name, Node, frame_ty, Gen)
|
||
self._EmitResumeFunction(name, Node, frame_ty, Gen)
|
||
# 符号表注册（普通名称），供 _HandleCallLlvm 的类型解析
|
||
if not self.Trans.SymbolTable.has(name):
|
||
FuncInfo: CTypeInfo = CTypeInfo()
|
||
FuncInfo.Name = name
|
||
FuncInfo.IsFunction = True
|
||
FuncInfo.file = '<stdin>'
|
||
self.Trans.SymbolTable.insert(name, FuncInfo)
|
||
self.Trans.FunctionReturnTypes[name] = CTypeInfo.FromTypeName('int')
|
||
finally:
|
||
while len(self.Trans.VarScopes) > saved_var_scope_len:
|
||
self.Trans.VarScopes.pop()
|
||
del Gen.loop_break_targets[saved_loop_len:]
|
||
del Gen.loop_continue_targets[saved_cont_len:]
|
||
Gen.builder = saved_builder
|
||
Gen.func = saved_func
|
||
Gen.variables = saved_variables
|
||
Gen._reg_values = saved_reg_values
|
||
Gen._direct_values = saved_direct_values
|
||
Gen.var_signedness = saved_var_signedness
|
||
Gen.var_struct_class = saved_var_struct_class
|
||
Gen.var_type_info = saved_var_type_info
|
||
Gen.var_type_assignments = saved_var_type_assignments
|
||
Gen._coro_state_ptr = saved_coro_state['state_ptr']
|
||
Gen._coro_done_ptr = saved_coro_state['done_ptr']
|
||
Gen._coro_ret_field_ptr = saved_coro_state['ret_ptr']
|
||
Gen._coro_exit_bb = saved_coro_state['exit_bb']
|
||
Gen._coro_next_state = saved_coro_state['next_state']
|
||
Gen._coro_await_slot_ptr = saved_coro_state['await_slot']
|
||
Gen._coro_frame_ptr = saved_coro_state['frame_ptr']
|
||
Gen._coro_switch_instr = saved_coro_state['switch']
|
||
Gen._coro_state_counter = saved_coro_state['state_counter']
|
||
Gen._coro_hoist_counter = saved_coro_state['hoist_counter']
|
||
|
||
# ------------------------------------------------------------------
|
||
# 帧结构体与局部变量收集
|
||
# ------------------------------------------------------------------
|
||
|
||
def _ResolveSlotType(self, annotation: ast.expr | None, Gen: LlvmCodeGenerator) -> ir.Type:
|
||
"""从注解推断帧槽 LLVM 类型；失败/void/结构体值一律按安全标量处理。"""
|
||
if annotation is None:
|
||
return ir.IntType(32)
|
||
try:
|
||
TypeInfo: CTypeInfo | None = self.Trans.TypeMergeHandler.GetCTypeInfo(annotation)
|
||
except Exception:
|
||
TypeInfo = None
|
||
if TypeInfo is None:
|
||
return ir.IntType(32)
|
||
try:
|
||
llvm_ty: ir.Type | None = TypeInfo.ToLLVM(Gen)
|
||
except Exception:
|
||
return ir.IntType(32)
|
||
if llvm_ty is None or isinstance(llvm_ty, ir.VoidType):
|
||
return ir.IntType(32)
|
||
if isinstance(llvm_ty, ir.IdentifiedStructType):
|
||
llvm_ty = ir.PointerType(llvm_ty)
|
||
return llvm_ty
|
||
|
||
def _CollectFrameVars(self, Node: ast.AsyncFunctionDef, Gen: LlvmCodeGenerator) -> list[tuple[str, ir.Type]]:
|
||
"""递归收集协程体内所有 AnnAssign/Assign 的 Name 目标及其槽类型（含控制流内部）。"""
|
||
vars_: list[tuple[str, ir.Type]] = []
|
||
seen: set[str] = set()
|
||
|
||
def _walk(stmts: list[ast.stmt]) -> None:
|
||
for stmt in stmts:
|
||
if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
|
||
target: str = stmt.target.id
|
||
slot_ty: ir.Type = self._ResolveSlotType(stmt.annotation, Gen)
|
||
elif (isinstance(stmt, ast.Assign)
|
||
and len(stmt.targets) == 1
|
||
and isinstance(stmt.targets[0], ast.Name)):
|
||
target = stmt.targets[0].id
|
||
slot_ty = ir.IntType(32)
|
||
elif isinstance(stmt, (ast.If, ast.While, ast.For)):
|
||
_walk(stmt.body)
|
||
if stmt.orelse:
|
||
_walk(stmt.orelse)
|
||
continue
|
||
else:
|
||
continue
|
||
if target not in seen:
|
||
seen.add(target)
|
||
vars_.append((target, slot_ty))
|
||
|
||
_walk(Node.body)
|
||
return vars_
|
||
|
||
def _GetOrCreateFrameType(self, name: str, Node: ast.AsyncFunctionDef, Gen: LlvmCodeGenerator, hoist_count: int) -> ir.IdentifiedStructType:
|
||
if not hasattr(Gen, '_coro_frame_types'):
|
||
Gen._coro_frame_types = {}
|
||
if not hasattr(Gen, '_coro_slot_vars'):
|
||
Gen._coro_slot_vars = {}
|
||
if not hasattr(Gen, '_coro_param_vars'):
|
||
Gen._coro_param_vars = {}
|
||
if name in Gen._coro_frame_types:
|
||
return Gen._coro_frame_types[name]
|
||
# 参数：名字 + LLVM 类型
|
||
params: list[tuple[str, ir.Type]] = []
|
||
for a in Node.args.args:
|
||
pname: str = a.arg
|
||
pty: ir.Type = self._ResolveSlotType(a.annotation, Gen)
|
||
params.append((pname, pty))
|
||
locals_: list[tuple[str, ir.Type]] = self._CollectFrameVars(Node, Gen)
|
||
param_names: set[str] = {p for p, _ in params}
|
||
# 提升临时变量槽：__await_0 .. __await_{hoist_count-1}
|
||
fields: list[ir.Type] = [
|
||
ir.IntType(32), # 0: state
|
||
ir.IntType(32), # 1: done
|
||
ir.IntType(32), # 2: ret
|
||
]
|
||
slot_vars: list[str] = []
|
||
for pname, pty in params:
|
||
fields.append(pty)
|
||
slot_vars.append(pname)
|
||
for lname, lty in locals_:
|
||
if lname in param_names or lname in slot_vars:
|
||
continue
|
||
fields.append(lty)
|
||
slot_vars.append(lname)
|
||
for i in range(hoist_count):
|
||
fields.append(ir.IntType(32))
|
||
slot_vars.append(f'__await_{i}')
|
||
frame_ty: ir.IdentifiedStructType = ir.IdentifiedStructType(Gen.module, f'Coro_{name}')
|
||
frame_ty.set_body(*fields)
|
||
Gen._coro_frame_types[name] = frame_ty
|
||
Gen.structs[f'Coro_{name}'] = frame_ty
|
||
Gen._coro_slot_vars[name] = slot_vars
|
||
Gen._coro_param_vars[name] = params
|
||
return frame_ty
|
||
|
||
# ------------------------------------------------------------------
|
||
# 预扫描：嵌套 await 提升（分配 __await_N 临时槽）
|
||
# ------------------------------------------------------------------
|
||
|
||
def _PreScanHoistedVars(self, Node: ast.AsyncFunctionDef, Gen: LlvmCodeGenerator) -> int:
|
||
"""统计需提升的嵌套 await 数量（与翻译端提取顺序一致）。"""
|
||
counter: int = 0
|
||
for stmt in Node.body:
|
||
counter = self._CountNestedAwaitForStmt(stmt, counter, Gen)
|
||
return counter
|
||
|
||
def _CountNestedAwaitForStmt(self, stmt: ast.stmt, counter: int, Gen: LlvmCodeGenerator) -> int:
|
||
if isinstance(stmt, (ast.If, ast.While)):
|
||
for s in stmt.body:
|
||
counter = self._CountNestedAwaitForStmt(s, counter, Gen)
|
||
if stmt.orelse:
|
||
for s in stmt.orelse:
|
||
counter = self._CountNestedAwaitForStmt(s, counter, Gen)
|
||
return counter
|
||
if isinstance(stmt, ast.For):
|
||
if any(isinstance(n, ast.Await) for n in ast.walk(stmt)):
|
||
raise SyntaxError(
|
||
f"[coro-M2] For 循环暂不支持 await（M3 实现）(line {getattr(stmt, 'lineno', 0)})"
|
||
)
|
||
return counter
|
||
# 顶层 await 赋值（value 直接是 Await）不提升
|
||
if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.value, ast.Await):
|
||
return counter
|
||
if (isinstance(stmt, ast.Assign)
|
||
and len(stmt.targets) == 1
|
||
and isinstance(stmt.targets[0], ast.Name)
|
||
and isinstance(stmt.value, ast.Await)):
|
||
return counter
|
||
for n in ast.walk(stmt):
|
||
if isinstance(n, ast.Await):
|
||
counter += 1
|
||
return counter
|
||
|
||
# ------------------------------------------------------------------
|
||
# create_<name>(args...) -> %Coro_<name>*
|
||
# ------------------------------------------------------------------
|
||
|
||
def _GetMalloc(self, Gen: LlvmCodeGenerator) -> ir.Function:
|
||
for fn in Gen.module.global_values:
|
||
if getattr(fn, 'name', None) == 'malloc':
|
||
return fn
|
||
malloc_ty: ir.FunctionType = ir.FunctionType(ir.PointerType(ir.IntType(8)), [ir.IntType(64)])
|
||
fn: ir.Function = ir.Function(Gen.module, malloc_ty, name='malloc')
|
||
Gen.functions['malloc'] = fn
|
||
return fn
|
||
|
||
def _RegisterCoroFunction(self, Gen: LlvmCodeGenerator, raw_name: str, fn: ir.Function) -> None:
|
||
"""按 _EmitFunctionForwardDeclLlvm 惯例注册 create_/resume_ 函数。"""
|
||
mangled: str = Gen._mangle_func_name(raw_name)
|
||
Gen.functions[mangled] = fn
|
||
if mangled == raw_name or raw_name not in Gen._export_extern_funcs:
|
||
Gen.functions[raw_name] = fn
|
||
|
||
def _EmitCreateFunction(self, name: str, Node: ast.AsyncFunctionDef, frame_ty: ir.IdentifiedStructType, Gen: LlvmCodeGenerator) -> None:
|
||
malloc_fn: ir.Function = self._GetMalloc(Gen)
|
||
frame_size: int = Gen._get_struct_size(frame_ty)
|
||
params: list[tuple[str, ir.Type]] = getattr(Gen, '_coro_param_vars', {}).get(name, [])
|
||
param_tys: list[ir.Type] = [ty for _, ty in params]
|
||
create_ty: ir.FunctionType = ir.FunctionType(ir.PointerType(frame_ty), param_tys)
|
||
create_fn: ir.Function = ir.Function(Gen.module, create_ty, name=Gen._mangle_func_name(f'create_{name}'))
|
||
self._RegisterCoroFunction(Gen, f'create_{name}', create_fn)
|
||
create_fn.attributes.add('noredzone')
|
||
entry: ir.Block = create_fn.append_basic_block('entry')
|
||
Gen.builder = ir.IRBuilder(entry)
|
||
Gen.func = create_fn
|
||
raw: ir.Value = Gen.builder.call(malloc_fn, [ir.Constant(ir.IntType(64), frame_size)], name='frame_raw')
|
||
frame_ptr: ir.Value = Gen.builder.bitcast(raw, ir.PointerType(frame_ty), name='frame')
|
||
zero32: ir.Constant = ir.Constant(ir.IntType(32), 0)
|
||
for i in range(3):
|
||
slot_ptr: ir.Value = Gen.builder.gep(frame_ptr, [zero32, ir.Constant(ir.IntType(32), i)], name=f'init_slot{i}')
|
||
Gen._store(zero32, slot_ptr)
|
||
# 参数存入帧槽（3 + i）
|
||
for i in range(len(params)):
|
||
slot_ptr = Gen.builder.gep(frame_ptr, [zero32, ir.Constant(ir.IntType(32), 3 + i)], name=f'param{i}_slot')
|
||
arg_val: ir.Argument = create_fn.args[i]
|
||
arg_val.name = f'arg{i}'
|
||
Gen._store(arg_val, slot_ptr)
|
||
Gen.builder.ret(frame_ptr)
|
||
Gen.builder = None
|
||
Gen.func = None
|
||
|
||
# ------------------------------------------------------------------
|
||
# resume_<name>(%Coro_<name>*) -> i32 状态机（动态状态号）
|
||
# ------------------------------------------------------------------
|
||
|
||
def _EmitResumeFunction(self, name: str, Node: ast.AsyncFunctionDef, frame_ty: ir.IdentifiedStructType, Gen: LlvmCodeGenerator) -> None:
|
||
resume_ty: ir.FunctionType = ir.FunctionType(ir.IntType(32), [ir.PointerType(frame_ty)])
|
||
resume_fn: ir.Function = ir.Function(Gen.module, resume_ty, name=Gen._mangle_func_name(f'resume_{name}'))
|
||
self._RegisterCoroFunction(Gen, f'resume_{name}', resume_fn)
|
||
resume_fn.attributes.add('noredzone')
|
||
Gen.func = resume_fn
|
||
entry: ir.Block = resume_fn.append_basic_block('entry')
|
||
Gen.builder = ir.IRBuilder(entry)
|
||
frame_arg: ir.Argument = resume_fn.args[0]
|
||
frame_arg.name = 'frame'
|
||
zero32: ir.Constant = ir.Constant(ir.IntType(32), 0)
|
||
one32: ir.Constant = ir.Constant(ir.IntType(32), 1)
|
||
two32: ir.Constant = ir.Constant(ir.IntType(32), 2)
|
||
state_ptr: ir.Value = Gen.builder.gep(frame_arg, [zero32, zero32], name='state_ptr')
|
||
done_ptr: ir.Value = Gen.builder.gep(frame_arg, [zero32, one32], name='done_ptr')
|
||
ret_ptr: ir.Value = Gen.builder.gep(frame_arg, [zero32, two32], name='ret_ptr')
|
||
# 槽指针 → Gen.variables（读取/写入自动走帧字段 GEP）
|
||
local_slot_ptrs: dict[str, ir.Value] = {}
|
||
slot_vars: list[str] = getattr(Gen, '_coro_slot_vars', {}).get(name, [])
|
||
for idx, var in enumerate(slot_vars):
|
||
slot_ptr: ir.Value = Gen.builder.gep(
|
||
frame_arg,
|
||
[zero32, ir.Constant(ir.IntType(32), 3 + idx)],
|
||
name=f'{var}_slot',
|
||
)
|
||
local_slot_ptrs[var] = slot_ptr
|
||
state_val: ir.Value = Gen._load(state_ptr, name='state')
|
||
exit_bb: ir.Block = resume_fn.append_basic_block('coro_exit')
|
||
Gen._coro_state_ptr = state_ptr
|
||
Gen._coro_done_ptr = done_ptr
|
||
Gen._coro_ret_field_ptr = ret_ptr
|
||
Gen._coro_exit_bb = exit_bb
|
||
Gen._coro_frame_ptr = frame_arg
|
||
switch_instr: Any = Gen.builder.switch(state_val, exit_bb)
|
||
Gen._coro_switch_instr = switch_instr
|
||
Gen._coro_state_counter = 1
|
||
state0: ir.Block = resume_fn.append_basic_block('state0')
|
||
switch_instr.add_case(ir.Constant(ir.IntType(32), 0), state0)
|
||
# 主入口：翻译整个函数体（If/While 内 await 会动态分配恢复块）
|
||
Gen.builder.position_at_start(state0)
|
||
Gen.variables = dict(local_slot_ptrs)
|
||
self._HandleCoroSegment(Node.body, Gen)
|
||
if Gen.builder is not None and not Gen.builder.block.is_terminated:
|
||
self._EmitCoroImplicitReturn(Gen)
|
||
# 完成块：返回 done（0=SUSPEND 1=DONE）
|
||
Gen.builder.position_at_start(exit_bb)
|
||
done_val: ir.Value = Gen._load(done_ptr, name='done')
|
||
Gen.builder.ret(done_val)
|
||
Gen.builder = None
|
||
Gen.func = None
|
||
|
||
# ------------------------------------------------------------------
|
||
# 协程体段翻译
|
||
# ------------------------------------------------------------------
|
||
|
||
def _HandleCoroSegment(self, stmts: list[ast.stmt], Gen: LlvmCodeGenerator) -> None:
|
||
for stmt in stmts:
|
||
if Gen.builder is None or Gen.builder.block.is_terminated:
|
||
return
|
||
Gen._set_node_info(stmt)
|
||
if isinstance(stmt, ast.Return):
|
||
self._HandleCoroReturn(stmt, Gen)
|
||
continue
|
||
if (isinstance(stmt, ast.AnnAssign)
|
||
and isinstance(stmt.target, ast.Name)
|
||
and isinstance(stmt.value, ast.Await)):
|
||
self._HandleCoroAwaitAnnAssign(stmt, Gen)
|
||
continue
|
||
if (isinstance(stmt, ast.Assign)
|
||
and len(stmt.targets) == 1
|
||
and isinstance(stmt.targets[0], ast.Name)
|
||
and isinstance(stmt.value, ast.Await)):
|
||
self._HandleCoroAwaitAssign(stmt, Gen)
|
||
continue
|
||
if isinstance(stmt, ast.If):
|
||
self._HandleCoroIf(stmt, Gen)
|
||
continue
|
||
if isinstance(stmt, ast.While):
|
||
self._HandleCoroWhile(stmt, Gen)
|
||
continue
|
||
if isinstance(stmt, ast.For):
|
||
self._HandleCoroFor(stmt, Gen)
|
||
continue
|
||
if isinstance(stmt, ast.Break):
|
||
if Gen.loop_break_targets:
|
||
Gen.builder.branch(Gen.loop_break_targets[-1])
|
||
continue
|
||
if isinstance(stmt, ast.Continue):
|
||
if Gen.loop_continue_targets:
|
||
Gen.builder.branch(Gen.loop_continue_targets[-1])
|
||
continue
|
||
if isinstance(stmt, ast.Pass):
|
||
continue
|
||
# 语句内嵌套 await（非顶层）→ 提升
|
||
if any(isinstance(n, ast.Await) for n in ast.walk(stmt)):
|
||
self._HandleNestedAwaitStmt(stmt, Gen)
|
||
continue
|
||
self._TranslateOrdinaryStmt(stmt, Gen)
|
||
|
||
def _TranslateOrdinaryStmt(self, stmt: ast.stmt, Gen: LlvmCodeGenerator) -> None:
|
||
if isinstance(stmt, ast.AnnAssign):
|
||
self.Trans.AnnAssignHandler._HandleAnnAssignLlvm(stmt)
|
||
return
|
||
if isinstance(stmt, ast.Assign):
|
||
self.Trans.AssignHandler._HandleAssignLlvm(stmt)
|
||
return
|
||
if isinstance(stmt, ast.Expr):
|
||
self.Trans.ExprHandler.HandleExprLlvm(stmt.value)
|
||
return
|
||
raise SyntaxError(
|
||
f"[coro-M2] async 函数体暂不支持语句: {type(stmt).__name__} "
|
||
f"(line {getattr(stmt, 'lineno', 0)})"
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 嵌套 await 提升：a = b + await f() → __await_N = await f(); a = b + __await_N
|
||
# ------------------------------------------------------------------
|
||
|
||
def _FindFirstAwait(self, node: ast.AST) -> ast.Await | None:
|
||
for n in ast.walk(node):
|
||
if isinstance(n, ast.Await):
|
||
return n
|
||
return None
|
||
|
||
def _ReplaceAwaitWithName(self, node: ast.AST, temp_name: str) -> bool:
|
||
"""把语句树中第一个 Await 替换为 Name(temp_name, Load)（DFS 顺序）。"""
|
||
for field in node._fields:
|
||
val: Any = getattr(node, field, None)
|
||
if isinstance(val, ast.Await):
|
||
setattr(node, field, ast.Name(temp_name, ast.Load()))
|
||
return True
|
||
if isinstance(val, list):
|
||
for i, item in enumerate(val):
|
||
if isinstance(item, ast.Await):
|
||
val[i] = ast.Name(temp_name, ast.Load())
|
||
return True
|
||
if isinstance(item, ast.AST):
|
||
if self._ReplaceAwaitWithName(item, temp_name):
|
||
return True
|
||
elif isinstance(val, ast.AST):
|
||
if self._ReplaceAwaitWithName(val, temp_name):
|
||
return True
|
||
return False
|
||
|
||
def _HandleNestedAwaitStmt(self, stmt: ast.stmt, Gen: LlvmCodeGenerator) -> None:
|
||
"""把语句中所有嵌套 await 提升为顶层临时赋值（与预扫描顺序一致）。"""
|
||
counter: int = getattr(Gen, '_coro_hoist_counter', 0)
|
||
while True:
|
||
await_node: ast.Await | None = self._FindFirstAwait(stmt)
|
||
if await_node is None:
|
||
break
|
||
temp_name: str = f'__await_{counter}'
|
||
counter += 1
|
||
hoisted: ast.Assign = ast.Assign(
|
||
targets=[ast.Name(temp_name, ast.Store())],
|
||
value=await_node,
|
||
type_comment=None,
|
||
)
|
||
# 提升赋值挂起 + position 到恢复块
|
||
self._HandleCoroAwaitAssign(hoisted, Gen)
|
||
self._ReplaceAwaitWithName(stmt, temp_name)
|
||
Gen._coro_hoist_counter = counter
|
||
self._TranslateOrdinaryStmt(stmt, Gen)
|
||
|
||
# ------------------------------------------------------------------
|
||
# await 表达式翻译
|
||
# ------------------------------------------------------------------
|
||
|
||
def _HandleCoroAwaitAnnAssign(self, stmt: ast.AnnAssign, Gen: LlvmCodeGenerator) -> None:
|
||
target: str = stmt.target.id
|
||
self._HandleCoroAwaitCommon(target, stmt.value, Gen)
|
||
# 注册作用域信息（模拟 AnnAssignHandler 的局部变量路径）
|
||
TypeInfo: CTypeInfo | None = self.Trans.TypeMergeHandler.GetCTypeInfo(stmt.annotation)
|
||
if TypeInfo is None:
|
||
TypeInfo = CTypeInfo()
|
||
TypeInfo.BaseType = t.CInt()
|
||
if self.Trans.VarScopes:
|
||
self.Trans.VarScopes[-1][target] = TypeInfo
|
||
Gen._record_var_signedness(target, False)
|
||
|
||
def _HandleCoroAwaitAssign(self, stmt: ast.Assign, Gen: LlvmCodeGenerator) -> None:
|
||
target: str = stmt.targets[0].id
|
||
self._HandleCoroAwaitCommon(target, stmt.value, Gen)
|
||
Gen._record_var_signedness(target, False)
|
||
|
||
def _HandleCoroAwaitCommon(self, target: str, await_node: ast.Await, Gen: LlvmCodeGenerator) -> None:
|
||
"""await 顶层翻译：设置目标槽后调用 _HandleAwaitLlvm（生成子帧/resume/挂起）。"""
|
||
Gen._coro_await_slot_ptr = Gen.variables.get(target)
|
||
# 校验 await 目标必须是 async 函数调用，且子协程已生成
|
||
call: Any = await_node.value
|
||
async_name: str | None = self._GetAsyncCallName(call) if isinstance(call, ast.Call) else None
|
||
if not async_name:
|
||
raise SyntaxError(
|
||
f"[coro-M2] await 目标必须是 async 函数调用 (line {getattr(await_node, 'lineno', 0)})"
|
||
)
|
||
self.Trans.ExprHandler.HandleExprLlvm(await_node) # → _HandleAwaitLlvm
|
||
|
||
def _CoerceArg(self, v: ir.Value | None, target_ty: ir.Type, Gen: LlvmCodeGenerator, name: str = 'arg') -> ir.Value:
|
||
"""实参转换到 create 形参类型（int 宽转换 / 指针 bitcast）。"""
|
||
if v is None:
|
||
if isinstance(target_ty, ir.IntType):
|
||
return ir.Constant(target_ty, 0)
|
||
return ir.Constant(target_ty, None)
|
||
src_ty: ir.Type = v.type
|
||
if src_ty == target_ty:
|
||
return v
|
||
if isinstance(src_ty, ir.IntType) and isinstance(target_ty, ir.IntType):
|
||
if src_ty.width < target_ty.width:
|
||
return Gen.builder.zext(v, target_ty, name=f'zext_{name}')
|
||
if src_ty.width > target_ty.width:
|
||
return Gen.builder.trunc(v, target_ty, name=f'trunc_{name}')
|
||
return v
|
||
if isinstance(src_ty, ir.PointerType) and isinstance(target_ty, ir.PointerType):
|
||
return Gen.builder.bitcast(v, target_ty, name=f'bitcast_{name}')
|
||
return v
|
||
|
||
def _HandleAwaitLlvm(self, Node: ast.Await, Gen: LlvmCodeGenerator) -> ir.Value:
|
||
"""await 表达式翻译：
|
||
创建子帧（含参数）-> resume 到 DONE -> 取子帧 ret -> 存父帧槽
|
||
-> 父挂起（state=新恢复块编号、ret=0）-> position 到恢复块继续。
|
||
"""
|
||
zero32: ir.Constant = ir.Constant(ir.IntType(32), 0)
|
||
call: Any = Node.value
|
||
async_name: str | None = self._GetAsyncCallName(call) if isinstance(call, ast.Call) else None
|
||
if not async_name:
|
||
return zero32
|
||
frame_ty: ir.IdentifiedStructType | None = getattr(Gen, '_coro_frame_types', {}).get(async_name)
|
||
create_fn: ir.Function | None = Gen._get_function(f'create_{async_name}')
|
||
resume_fn: ir.Function | None = Gen._get_function(f'resume_{async_name}')
|
||
if not frame_ty or not create_fn or not resume_fn:
|
||
return zero32
|
||
# 子协程参数
|
||
call_args: list[ir.Value] = []
|
||
if call.args:
|
||
param_info: list[tuple[str, ir.Type]] = getattr(Gen, '_coro_param_vars', {}).get(async_name, [])
|
||
for i, arg_node in enumerate(call.args):
|
||
arg_val: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(arg_node)
|
||
if i < len(param_info):
|
||
arg_val = self._CoerceArg(arg_val, param_info[i][1], Gen, name=f'arg{i}')
|
||
call_args.append(arg_val)
|
||
child: ir.Value = Gen.builder.call(create_fn, call_args, name=f'create_{async_name}')
|
||
Gen.builder.call(resume_fn, [child], name=f'resume_{async_name}')
|
||
child_ret_ptr: ir.Value = Gen.builder.gep(
|
||
child,
|
||
[zero32, ir.Constant(ir.IntType(32), 2)],
|
||
name='child_ret_ptr',
|
||
)
|
||
child_ret: ir.Value = Gen._load(child_ret_ptr, name='child_ret')
|
||
# 存父帧局部变量槽
|
||
slot: ir.Value | None = getattr(Gen, '_coro_await_slot_ptr', None)
|
||
if slot is not None:
|
||
Gen._store(child_ret, slot)
|
||
# 父挂起：分配新恢复块（switch case），state=新块编号，ret=0(SUSPEND)
|
||
state_no: int = getattr(Gen, '_coro_state_counter', 1)
|
||
Gen._coro_state_counter = state_no + 1
|
||
resume_bb: ir.Block = Gen.func.append_basic_block(f'state{state_no}')
|
||
switch_instr: Any = getattr(Gen, '_coro_switch_instr', None)
|
||
if switch_instr is not None:
|
||
switch_instr.add_case(ir.Constant(ir.IntType(32), state_no), resume_bb)
|
||
state_ptr: ir.Value | None = getattr(Gen, '_coro_state_ptr', None)
|
||
if state_ptr is not None:
|
||
Gen._store(ir.Constant(ir.IntType(32), state_no), state_ptr)
|
||
ret_field_ptr: ir.Value | None = getattr(Gen, '_coro_ret_field_ptr', None)
|
||
if ret_field_ptr is not None:
|
||
Gen._store(zero32, ret_field_ptr)
|
||
exit_bb: ir.Block | None = getattr(Gen, '_coro_exit_bb', None)
|
||
if exit_bb is not None:
|
||
Gen.builder.branch(exit_bb)
|
||
Gen.builder.position_at_start(resume_bb)
|
||
return child_ret
|
||
|
||
def _HandleCoroReturn(self, stmt: ast.Return, Gen: LlvmCodeGenerator) -> None:
|
||
val: ir.Value | None = None
|
||
if stmt.value:
|
||
val = self.Trans.ExprHandler.HandleExprLlvm(stmt.value)
|
||
if val is None:
|
||
val = ir.Constant(ir.IntType(32), 0)
|
||
if isinstance(val.type, ir.IntType) and val.type.width != 32:
|
||
if val.type.width < 32:
|
||
val = Gen.builder.zext(val, ir.IntType(32), name='zext_coro_ret')
|
||
else:
|
||
val = Gen.builder.trunc(val, ir.IntType(32), name='trunc_coro_ret')
|
||
Gen._store(val, Gen._coro_ret_field_ptr)
|
||
Gen._store(ir.Constant(ir.IntType(32), 1), Gen._coro_done_ptr)
|
||
Gen.builder.branch(Gen._coro_exit_bb)
|
||
|
||
def _EmitCoroImplicitReturn(self, Gen: LlvmCodeGenerator) -> None:
|
||
"""函数体自然结束：同 return 路径（ret=0 + done=1）。"""
|
||
Gen._store(ir.Constant(ir.IntType(32), 0), Gen._coro_ret_field_ptr)
|
||
Gen._store(ir.Constant(ir.IntType(32), 1), Gen._coro_done_ptr)
|
||
Gen.builder.branch(Gen._coro_exit_bb)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 控制流（协程版：body 内 await 就地挂起/恢复）
|
||
# ------------------------------------------------------------------
|
||
|
||
def _HandleCoroIf(self, Node: ast.If, Gen: LlvmCodeGenerator) -> None:
|
||
if any(isinstance(n, ast.Await) for n in ast.walk(Node.test)):
|
||
raise SyntaxError(
|
||
f"[coro-M2] if 条件中不支持 await (line {getattr(Node, 'lineno', 0)})"
|
||
)
|
||
Cond: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(Node.test)
|
||
if not Cond:
|
||
return
|
||
if not isinstance(Cond.type, ir.IntType) or Cond.type.width != 1:
|
||
if isinstance(Cond.type, (ir.FloatType, ir.DoubleType)):
|
||
Cond = Gen.builder.fcmp_ordered('!=', Cond, ir.Constant(Cond.type, 0.0), name='coro_ifcond')
|
||
else:
|
||
Cond = Gen.builder.icmp_signed('!=', Cond, Gen._ZeroConst(Cond.type), name='coro_ifcond')
|
||
ThenBB: ir.Block = Gen.func.append_basic_block(name='coro_then')
|
||
ElseBB: ir.Block | None = Gen.func.append_basic_block(name='coro_else') if Node.orelse else None
|
||
MergeBB: ir.Block = Gen.func.append_basic_block(name='coro_endif')
|
||
if not Gen.builder.block.is_terminated:
|
||
Gen.builder.cbranch(Cond, ThenBB, ElseBB if ElseBB else MergeBB)
|
||
Gen.builder.position_at_start(ThenBB)
|
||
self._HandleCoroSegment(Node.body, Gen)
|
||
if Gen.builder is not None and not Gen.builder.block.is_terminated:
|
||
Gen.builder.branch(MergeBB)
|
||
if ElseBB:
|
||
Gen.builder.position_at_start(ElseBB)
|
||
self._HandleCoroSegment(Node.orelse, Gen)
|
||
if Gen.builder is not None and not Gen.builder.block.is_terminated:
|
||
Gen.builder.branch(MergeBB)
|
||
if not MergeBB.is_terminated:
|
||
Gen.builder.position_at_start(MergeBB)
|
||
else:
|
||
Gen.builder.position_at_end(MergeBB)
|
||
|
||
def _HandleCoroWhile(self, Node: ast.While, Gen: LlvmCodeGenerator) -> None:
|
||
if any(isinstance(n, ast.Await) for n in ast.walk(Node.test)):
|
||
raise SyntaxError(
|
||
f"[coro-M2] while 条件中不支持 await (line {getattr(Node, 'lineno', 0)})"
|
||
)
|
||
CondBB: ir.Block = Gen.func.append_basic_block(name='coro_wcond')
|
||
BodyBB: ir.Block = Gen.func.append_basic_block(name='coro_wbody')
|
||
ElseBB: ir.Block | None = Gen.func.append_basic_block(name='coro_welse') if Node.orelse else None
|
||
AfterBB: ir.Block = Gen.func.append_basic_block(name='coro_wend')
|
||
Gen.loop_break_targets.append(AfterBB)
|
||
Gen.loop_continue_targets.append(CondBB)
|
||
Gen.builder.branch(CondBB)
|
||
Gen.builder.position_at_start(CondBB)
|
||
Cond: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(Node.test)
|
||
if Cond:
|
||
if not isinstance(Cond.type, ir.IntType) or Cond.type.width != 1:
|
||
if isinstance(Cond.type, (ir.FloatType, ir.DoubleType)):
|
||
Cond = Gen.builder.fcmp_ordered('!=', Cond, ir.Constant(Cond.type, 0.0), name='coro_wcond')
|
||
else:
|
||
Cond = Gen.builder.icmp_signed('!=', Cond, Gen._ZeroConst(Cond.type), name='coro_wcond')
|
||
if not Gen.builder.block.is_terminated:
|
||
Gen.builder.cbranch(Cond, BodyBB, ElseBB if ElseBB else AfterBB)
|
||
else:
|
||
if not Gen.builder.block.is_terminated:
|
||
Gen.builder.branch(ElseBB if ElseBB else AfterBB)
|
||
Gen.builder.position_at_start(BodyBB)
|
||
self._HandleCoroSegment(Node.body, Gen)
|
||
if Gen.builder is not None and not Gen.builder.block.is_terminated:
|
||
Gen.builder.branch(CondBB)
|
||
Gen.loop_break_targets.pop()
|
||
Gen.loop_continue_targets.pop()
|
||
if ElseBB:
|
||
Gen.builder.position_at_start(ElseBB)
|
||
self._HandleCoroSegment(Node.orelse, Gen)
|
||
if Gen.builder is not None and not Gen.builder.block.is_terminated:
|
||
Gen.builder.branch(AfterBB)
|
||
if not AfterBB.is_terminated:
|
||
Gen.builder.position_at_start(AfterBB)
|
||
else:
|
||
Gen.builder.position_at_end(AfterBB)
|
||
|
||
def _HandleCoroFor(self, Node: ast.For, Gen: LlvmCodeGenerator) -> None:
|
||
if any(isinstance(n, ast.Await) for n in ast.walk(Node)):
|
||
raise SyntaxError(
|
||
f"[coro-M2] For 循环暂不支持 await（M3 实现）(line {getattr(Node, 'lineno', 0)})"
|
||
)
|
||
if any(isinstance(n, ast.Return) for n in ast.walk(Node.body)):
|
||
raise SyntaxError(
|
||
f"[coro-M2] For 循环体内暂不支持 return（M3 实现）(line {getattr(Node, 'lineno', 0)})"
|
||
)
|
||
# 无 await/return：委托现有 ForHandler（迭代索引等 alloca 不跨挂起，安全）
|
||
self.Trans.ForHandler._HandleForLlvm(Node)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 普通函数调用 async 函数：task(...) -> create_task(...)（返回帧指针）
|
||
# ------------------------------------------------------------------
|
||
|
||
def _HandleAsyncCreateCallLlvm(self, Node: ast.Call, FuncName: str) -> ir.Value:
|
||
Gen: LlvmCodeGenerator = self.Trans.LlvmGen
|
||
create_fn: ir.Function | None = Gen._get_function(f'create_{FuncName}')
|
||
if not create_fn:
|
||
return ir.Constant(ir.IntType(32), 0)
|
||
call_args: list[ir.Value] = []
|
||
if Node.args:
|
||
param_info: list[tuple[str, ir.Type]] = getattr(Gen, '_coro_param_vars', {}).get(FuncName, [])
|
||
for i, arg_node in enumerate(Node.args):
|
||
arg_val: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(arg_node)
|
||
if i < len(param_info):
|
||
arg_val = self._CoerceArg(arg_val, param_info[i][1], Gen, name=f'arg{i}')
|
||
call_args.append(arg_val)
|
||
result: ir.Value = Gen.builder.call(create_fn, call_args, name=f'create_{FuncName}')
|
||
return Gen.builder.bitcast(result, ir.IntType(8).as_pointer(), name='coro_frame_i8ptr')
|
||
|
||
# ------------------------------------------------------------------
|
||
# coro.coro_run(帧) 内联翻译
|
||
# ------------------------------------------------------------------
|
||
|
||
def _HandleCoroRunLlvm(self, Node: ast.Call, Gen: LlvmCodeGenerator) -> ir.Value:
|
||
if not Node.args:
|
||
return ir.Constant(ir.IntType(32), 0)
|
||
arg: ast.expr = Node.args[0]
|
||
async_name: str | None = None
|
||
if isinstance(arg, ast.Name):
|
||
async_name = getattr(Gen, '_coro_var_async', {}).get(arg.id)
|
||
if not async_name:
|
||
return ir.Constant(ir.IntType(32), 0)
|
||
frame_ty: ir.IdentifiedStructType | None = getattr(Gen, '_coro_frame_types', {}).get(async_name)
|
||
resume_fn: ir.Function | None = Gen._get_function(f'resume_{async_name}')
|
||
if not frame_ty or not resume_fn:
|
||
return ir.Constant(ir.IntType(32), 0)
|
||
frame_val: ir.Value | None = self.Trans.ExprHandler.HandleExprLlvm(arg)
|
||
if frame_val is None:
|
||
return ir.Constant(ir.IntType(32), 0)
|
||
if isinstance(frame_val.type, ir.PointerType):
|
||
frame_ptr: ir.Value = Gen.builder.bitcast(frame_val, ir.PointerType(frame_ty), name='coro_frame')
|
||
elif isinstance(frame_val.type, ir.IntType):
|
||
frame_ptr = Gen.builder.inttoptr(frame_val, ir.PointerType(frame_ty), name='coro_frame')
|
||
else:
|
||
return ir.Constant(ir.IntType(32), 0)
|
||
loop_bb: ir.Block = Gen.func.append_basic_block('coro_loop')
|
||
after_bb: ir.Block = Gen.func.append_basic_block('coro_after')
|
||
Gen.builder.branch(loop_bb)
|
||
Gen.builder.position_at_start(loop_bb)
|
||
status: ir.Value = Gen.builder.call(resume_fn, [frame_ptr], name=f'resume_{async_name}')
|
||
done_cond: ir.Value = Gen.builder.icmp_signed('!=', status, ir.Constant(ir.IntType(32), 0), name='coro_done_cond')
|
||
Gen.builder.cbranch(done_cond, after_bb, loop_bb)
|
||
Gen.builder.position_at_start(after_bb)
|
||
ret_ptr: ir.Value = Gen.builder.gep(
|
||
frame_ptr,
|
||
[ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), 2)],
|
||
name='coro_ret_ptr',
|
||
)
|
||
return Gen._load(ret_ptr, name='coro_result')
|