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): """语言内协程 M1：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。 """ # ------------------------------------------------------------------ # 公共查询 # ------------------------------------------------------------------ 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 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) # 保存外层 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_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), } try: self.Trans.VarScopes.append({}) frame_ty: ir.IdentifiedStructType = self._GetOrCreateFrameType(name, Node, Gen) self._EmitCreateFunction(name, 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() 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'] # ------------------------------------------------------------------ # 帧结构体与局部变量收集 # ------------------------------------------------------------------ def _GetOrCreateFrameType(self, name: str, Node: ast.AsyncFunctionDef, Gen: LlvmCodeGenerator) -> 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 name in Gen._coro_frame_types: return Gen._coro_frame_types[name] local_names: list[str] = self._CollectLocalVars(Node) fields: list[ir.Type] = [ ir.IntType(32), # 0: state ir.IntType(32), # 1: done ir.IntType(32), # 2: ret ] for _ in local_names: fields.append(ir.IntType(32)) frame_ty: ir.IdentifiedStructType = ir.IdentifiedStructType(Gen.module, f'Coro_{name}') frame_ty.set_body(*fields) Gen._coro_frame_types[name] = frame_ty Gen._coro_local_vars[name] = local_names return frame_ty def _CollectLocalVars(self, Node: ast.AsyncFunctionDef) -> list[str]: """收集协程体内顶层 AnnAssign/Assign 的 Name 目标（M1 仅 i32 槽）。""" names: list[str] = [] seen: set[str] = set() 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 # ------------------------------------------------------------------ # create_() -> %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, 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), []) 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) Gen.builder.ret(frame_ptr) Gen.builder = None Gen.func = None # ------------------------------------------------------------------ # 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] 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] = {} for idx, var in enumerate(Gen._coro_local_vars.get(name, [])): 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) 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) # 完成块：返回 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 isinstance(stmt, ast.AnnAssign) and isinstance(stmt.value, ast.Await): self._HandleCoroAwaitAnnAssign(stmt, Gen) return 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） 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) 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)})" ) 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-M1] await 目标必须是 async 函数调用 (line {getattr(await_node, 'lineno', 0)})" ) self.Trans.ExprHandler.HandleExprLlvm(await_node) # → _HandleAwaitLlvm def _HandleAwaitLlvm(self, Node: ast.Await, Gen: LlvmCodeGenerator) -> ir.Value: """await 表达式翻译： 创建子帧 -> resume 到 DONE（M1 假设一次完成）-> 取子帧 ret -> 存父帧槽 -> 父挂起。 """ 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 child: ir.Value = Gen.builder.call(create_fn, [], 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) # 父挂起：state = 下一切片，ret = 0(SUSPEND) 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) 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) 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) # ------------------------------------------------------------------ # 普通函数调用 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) result: ir.Value = Gen.builder.call(create_fn, [], name=f'create_{FuncName}') return Gen.builder.bitcast(result, ir.IntType(8).as_pointer(), name='coro_frame_i8ptr') # ------------------------------------------------------------------ # coro.coro_run(帧) 内联翻译（M1） # ------------------------------------------------------------------ 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')