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_ 带参数并存入帧。 - 嵌套 await 表达式：`a = b + await leaf()` 提升为顶层临时赋值 （`__await_N = await leaf(); a = b + __await_N`），帧预留临时变量槽。 - 控制流：If/While 语句（含内部 await 就地挂起/恢复）；For 仅支持无 await。 - resume 采用动态状态号：每个 await 点分配恢复块（state{N}），switch 分发。 帧布局：%Coro_ = type { i32 state, i32 done, i32 ret, <参数槽>, <局部槽>, <提升临时槽> } 状态机：create_(args...) 分配帧 + 存参数；resume_(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 = '' 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_(args...) -> %Coro_* # ------------------------------------------------------------------ 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_(%Coro_*) -> 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')