diff --git a/Test/CoroTest/App/main.vp b/Test/CoroTest/App/main.vp index 6809a3f..5f77b7b 100644 --- a/Test/CoroTest/App/main.vp +++ b/Test/CoroTest/App/main.vp @@ -4,18 +4,25 @@ import stdio import coro -async def leaf() -> int: - return 7 +async def leaf(x: int) -> int: + return x + 1 -async def task() -> int: - a: int = await leaf() - b: int = await leaf() - return a + b +async def task(base: int) -> int: + a: int = await leaf(base) + b: int = a * 2 + await leaf(a) + s: int = 0 + i: int = 0 + while i < 3: + s = s + await leaf(i) + i = i + 1 + if s > 5: + s = s + await leaf(10) + return s + a + b def main() -> t.CInt: - tt: coro.Coroutine = task() + tt: coro.Coroutine = task(2) r: t.CInt = coro.coro_run(tt) stdio.printf("coro result = %d\n", r) - return 0 + return 0 \ No newline at end of file diff --git a/lib/core/Handles/HandlesAsyncFunctions.py b/lib/core/Handles/HandlesAsyncFunctions.py index b475404..216a734 100644 --- a/lib/core/Handles/HandlesAsyncFunctions.py +++ b/lib/core/Handles/HandlesAsyncFunctions.py @@ -10,20 +10,18 @@ from lib.includes import t class AsyncFunctionHandle(BaseHandle): - """语言内协程 M1：async def + await + coro_run 状态机翻译。 + """语言内协程 M2：async def + await + coro_run 状态机翻译。 - M1 语义约定： - - async def 无参数、返回 i32/int。 - - 每个 async 函数生成 3 部分： - 1. 帧结构体 %Coro_ = type { i32, i32, i32, <跨 await 局部变量> } - （前 3 槽：state, done, ret） - 2. create_() -> %Coro_*：stdlib.malloc 分配帧，state/done/ret 置 0 - 3. resume_(%Coro_*) -> i32：0=SUSPEND 1=DONE；switch(state) 切分 await 点 - - await ：创建子帧 -> call resume(child)（M1 假设一次完成） - -> 取子帧 ret -> 存父帧局部变量槽 -> 父挂起（state=下一切片、ret=0(SUSPEND)） - - return expr：写帧 ret + done=1 + br 完成块；函数体结尾隐式同 return 路径。 - - 普通函数调用 async 函数（如 main 里 task()）：生成 create 调用（返回帧指针）。 - - coro.coro_run(帧)：内联翻译为 resume 循环，返回帧 ret。 + 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 与各恢复块。 """ # ------------------------------------------------------------------ @@ -63,14 +61,15 @@ class AsyncFunctionHandle(BaseHandle): def _EmitAsyncFunctionLlvm(self, Node: ast.AsyncFunctionDef, Gen: LlvmCodeGenerator) -> None: name: str = Node.name - if Node.args.args: - raise SyntaxError( - f"[coro-M1] async 函数 '{name}' 暂不支持参数（M1 仅支持无参数 async def）" - ) # 防重复生成：模块可能被多次编译（泛型特化等触发），避免 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 @@ -82,6 +81,8 @@ class AsyncFunctionHandle(BaseHandle): 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), @@ -90,11 +91,15 @@ class AsyncFunctionHandle(BaseHandle): '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) - self._EmitCreateFunction(name, frame_ty, Gen) + 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): @@ -107,6 +112,8 @@ class AsyncFunctionHandle(BaseHandle): 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 @@ -123,52 +130,146 @@ class AsyncFunctionHandle(BaseHandle): 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 _GetOrCreateFrameType(self, name: str, Node: ast.AsyncFunctionDef, Gen: LlvmCodeGenerator) -> ir.IdentifiedStructType: + 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_local_vars'): - Gen._coro_local_vars = {} + 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] - local_names: list[str] = self._CollectLocalVars(Node) + # 参数：名字 + 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 ] - for _ in local_names: + 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_local_vars[name] = local_names + Gen._coro_slot_vars[name] = slot_vars + Gen._coro_param_vars[name] = params return frame_ty - def _CollectLocalVars(self, Node: ast.AsyncFunctionDef) -> list[str]: - """收集协程体内顶层 AnnAssign/Assign 的 Name 目标（M1 仅 i32 槽）。""" - names: list[str] = [] - seen: set[str] = set() + # ------------------------------------------------------------------ + # 预扫描：嵌套 await 提升（分配 __await_N 临时槽） + # ------------------------------------------------------------------ + + def _PreScanHoistedVars(self, Node: ast.AsyncFunctionDef, Gen: LlvmCodeGenerator) -> int: + """统计需提升的嵌套 await 数量（与翻译端提取顺序一致）。""" + counter: int = 0 for stmt in Node.body: - target: str | None = None - if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name): - target = stmt.target.id - elif (isinstance(stmt, ast.Assign) - and len(stmt.targets) == 1 - and isinstance(stmt.targets[0], ast.Name)): - target = stmt.targets[0].id - if target is not None and target not in seen: - seen.add(target) - names.append(target) - return names + 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_() -> %Coro_* + # create_(args...) -> %Coro_* # ------------------------------------------------------------------ def _GetMalloc(self, Gen: LlvmCodeGenerator) -> ir.Function: @@ -187,10 +288,12 @@ class AsyncFunctionHandle(BaseHandle): if mangled == raw_name or raw_name not in Gen._export_extern_funcs: Gen.functions[raw_name] = fn - def _EmitCreateFunction(self, name: str, frame_ty: ir.IdentifiedStructType, Gen: LlvmCodeGenerator) -> None: + 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) - create_ty: ir.FunctionType = ir.FunctionType(ir.PointerType(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') @@ -203,46 +306,26 @@ class AsyncFunctionHandle(BaseHandle): 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 状态机 + # resume_(%Coro_*) -> i32 状态机（动态状态号） # ------------------------------------------------------------------ - def _SplitBody(self, body: list[ast.stmt]) -> list[list[ast.stmt]]: - """按顶层 await 切分函数体。 - - await 仅识别 AnnAssign.value / Assign.value（Name 目标）顶层形式； - 含 await 的语句作为所在段的最后一条（翻译后 br exit 挂起）。 - """ - segments: list[list[ast.stmt]] = [] - current: list[ast.stmt] = [] - for stmt in body: - has_await: bool = False - if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.value, ast.Await): - has_await = True - elif (isinstance(stmt, ast.Assign) - and len(stmt.targets) == 1 - and isinstance(stmt.targets[0], ast.Name) - and isinstance(stmt.value, ast.Await)): - has_await = True - current.append(stmt) - if has_await: - segments.append(current) - current = [] - if current: - segments.append(current) - return segments - 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 - segments: list[list[ast.stmt]] = self._SplitBody(Node.body) entry: ir.Block = resume_fn.append_basic_block('entry') Gen.builder = ir.IRBuilder(entry) frame_arg: ir.Argument = resume_fn.args[0] @@ -253,9 +336,10 @@ class AsyncFunctionHandle(BaseHandle): 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） + # 槽指针 → Gen.variables（读取/写入自动走帧字段 GEP） local_slot_ptrs: dict[str, ir.Value] = {} - for idx, var in enumerate(Gen._coro_local_vars.get(name, [])): + 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)], @@ -270,19 +354,16 @@ class AsyncFunctionHandle(BaseHandle): Gen._coro_exit_bb = exit_bb Gen._coro_frame_ptr = frame_arg switch_instr: Any = Gen.builder.switch(state_val, exit_bb) - state_bbs: list[ir.Block] = [] - for i in range(len(segments)): - bb: ir.Block = resume_fn.append_basic_block(f'state{i}') - state_bbs.append(bb) - switch_instr.add_case(ir.Constant(ir.IntType(32), i), bb) - # 逐段生成 - for i, seg in enumerate(segments): - Gen.builder.position_at_start(state_bbs[i]) - Gen.variables = dict(local_slot_ptrs) - Gen._coro_next_state = i + 1 - self._HandleCoroSegment(seg, Gen) - if Gen.builder is not None and not Gen.builder.block.is_terminated: - self._EmitCoroImplicitReturn(Gen) + 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') @@ -296,37 +377,116 @@ class AsyncFunctionHandle(BaseHandle): def _HandleCoroSegment(self, stmts: list[ast.stmt], Gen: LlvmCodeGenerator) -> None: for stmt in stmts: - if isinstance(stmt, ast.AnnAssign) and isinstance(stmt.value, ast.Await): - self._HandleCoroAwaitAnnAssign(stmt, Gen) + 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) - return - if isinstance(stmt, ast.Return): - self._HandleCoroReturn(stmt, Gen) - return - # M1 防御：await 仅支持顶层（AnnAssign/Assign 的 value 直接是 Await） + 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)): - raise SyntaxError( - f"[coro-M1] await 仅支持顶层赋值形式 `var: int = await fn()` " - f"(line {getattr(stmt, 'lineno', 0)})" - ) - if isinstance(stmt, ast.AnnAssign): - self.Trans.AnnAssignHandler._HandleAnnAssignLlvm(stmt) + self._HandleNestedAwaitStmt(stmt, Gen) continue - if isinstance(stmt, ast.Assign): - self.Trans.AssignHandler._HandleAssignLlvm(stmt) - continue - if isinstance(stmt, ast.Expr): - self.Trans.ExprHandler.HandleExprLlvm(stmt.value) - continue - raise SyntaxError( - f"[coro-M1] async 函数体暂不支持语句: {type(stmt).__name__} " - f"(line {getattr(stmt, 'lineno', 0)})" + 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 @@ -353,13 +513,33 @@ class AsyncFunctionHandle(BaseHandle): async_name: str | None = self._GetAsyncCallName(call) if isinstance(call, ast.Call) else None if not async_name: raise SyntaxError( - f"[coro-M1] await 目标必须是 async 函数调用 (line {getattr(await_node, 'lineno', 0)})" + 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（M1 假设一次完成）-> 取子帧 ret -> 存父帧槽 -> 父挂起。 + 创建子帧（含参数）-> resume 到 DONE -> 取子帧 ret -> 存父帧槽 + -> 父挂起（state=新恢复块编号、ret=0）-> position 到恢复块继续。 """ zero32: ir.Constant = ir.Constant(ir.IntType(32), 0) call: Any = Node.value @@ -371,7 +551,16 @@ class AsyncFunctionHandle(BaseHandle): 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 - child: ir.Value = Gen.builder.call(create_fn, [], name=f'create_{async_name}') + # 子协程参数 + 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, @@ -383,16 +572,23 @@ class AsyncFunctionHandle(BaseHandle): slot: ir.Value | None = getattr(Gen, '_coro_await_slot_ptr', None) if slot is not None: Gen._store(child_ret, slot) - # 父挂起：state = 下一切片，ret = 0(SUSPEND) + # 父挂起：分配新恢复块（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), getattr(Gen, '_coro_next_state', 1)), state_ptr) + 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: @@ -417,7 +613,96 @@ class AsyncFunctionHandle(BaseHandle): Gen.builder.branch(Gen._coro_exit_bb) # ------------------------------------------------------------------ - # 普通函数调用 async 函数：task() -> create_task()（返回帧指针） + # 控制流（协程版：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: @@ -425,11 +710,19 @@ class AsyncFunctionHandle(BaseHandle): create_fn: ir.Function | None = Gen._get_function(f'create_{FuncName}') if not create_fn: return ir.Constant(ir.IntType(32), 0) - result: ir.Value = Gen.builder.call(create_fn, [], name=f'create_{FuncName}') + 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(帧) 内联翻译（M1） + # coro.coro_run(帧) 内联翻译 # ------------------------------------------------------------------ def _HandleCoroRunLlvm(self, Node: ast.Call, Gen: LlvmCodeGenerator) -> ir.Value: