241 lines
10 KiB
Python
241 lines
10 KiB
Python
"""共享工具函数：符号表处理中使用的 AST 分析工具"""
|
||
from __future__ import annotations
|
||
import vast as ast
|
||
import re
|
||
from lib.includes import t
|
||
|
||
_T_MODULE_NAMES: frozenset[str] = frozenset({'t', t.__name__})
|
||
|
||
# Viper 软关键字 fn/fun/func/function 是 def 的别名（行首 + 空白 + 标识符才算函数定义，
|
||
# 变量名/调用/注解中的 fn 不受影响）。展开为 def 后行号不变，供 ast.parse 使用。
|
||
_SOFT_KEYWORD_RE: re.Pattern[str] = re.compile(
|
||
r'^(\s*)(function|func|fun|fn)(\s+)(?=[A-Za-z_])', re.MULTILINE
|
||
)
|
||
|
||
|
||
def ExpandSoftKeywords(Code: str) -> str:
|
||
"""将 Viper 软关键字 fn/fun/func/function 展开为 def（函数定义行）。
|
||
|
||
仅在"行首缩进 + 关键字 + 空白 + 标识符"的函数定义形态匹配，
|
||
变量名（fn = 5）、调用（fn(x)）、注解（x: fn）与字符串内容不受影响。
|
||
展开为 def 与关键字等宽或更短，行号不变，AST 行号映射保持有效。
|
||
"""
|
||
if not Code:
|
||
return Code
|
||
return _SOFT_KEYWORD_RE.sub(lambda m: f'{m.group(1)}def{m.group(3)}', Code)
|
||
|
||
|
||
def IsTModule(module_id: str, symtable: object | None = None) -> bool:
|
||
"""检查模块标识符是否指向 t 类型模块（支持别名如 import t as tt）"""
|
||
if module_id in _T_MODULE_NAMES:
|
||
return True
|
||
if symtable is not None and hasattr(symtable, 'import_aliases'):
|
||
resolved: str = symtable.import_aliases.get(module_id, module_id)
|
||
if resolved in _T_MODULE_NAMES:
|
||
return True
|
||
return False
|
||
|
||
|
||
def IsTModulePath(module_path: str) -> bool:
|
||
"""检查模块路径是否指向 t 类型模块"""
|
||
if module_path in _T_MODULE_NAMES:
|
||
return True
|
||
if module_path:
|
||
t_name: str
|
||
for t_name in _T_MODULE_NAMES:
|
||
if module_path == t_name or module_path.startswith(t_name + '.'):
|
||
return True
|
||
return False
|
||
|
||
|
||
def AnnotationContainsTType(node: ast.AST, TypeName: str) -> bool:
|
||
"""检测 AST 节点中是否包含指定名称的 t 模块类型"""
|
||
if isinstance(node, ast.Attribute):
|
||
if (hasattr(node.value, 'id') and IsTModule(node.value.id) and
|
||
node.attr == TypeName):
|
||
return True
|
||
elif isinstance(node, ast.BinOp):
|
||
return AnnotationContainsTType(node.left, TypeName) or AnnotationContainsTType(node.right, TypeName)
|
||
elif isinstance(node, ast.Call):
|
||
if AnnotationContainsTType(node.func, TypeName):
|
||
return True
|
||
for arg in node.args:
|
||
if AnnotationContainsTType(arg, TypeName):
|
||
return True
|
||
elif isinstance(node, ast.Subscript):
|
||
if AnnotationContainsTType(node.value, TypeName):
|
||
return True
|
||
if AnnotationContainsTType(node.slice, TypeName):
|
||
return True
|
||
return False
|
||
|
||
|
||
def IsListAnnotation(annotation: ast.expr) -> bool:
|
||
"""检查注解是否为栈上固定数组类型（t.CArray[...] 或 list[type, count]）"""
|
||
if not isinstance(annotation, ast.Subscript):
|
||
return False
|
||
value = annotation.value
|
||
# t.CArray[...] / c.CArray[...]
|
||
if (isinstance(value, ast.Attribute)
|
||
and isinstance(value.value, ast.Name)
|
||
and value.value.id in ('t', 'c')
|
||
and value.attr == 'CArray'):
|
||
return True
|
||
# 裸名 CArray[...]（from t import *）/ 内建 rawarray[...]：与 t.CArray 等价
|
||
# rawarray 是 CArray 的内建复刻名（hashlib 等裸名化重构使用，免写 t. 前缀）
|
||
if (isinstance(value, ast.Name)
|
||
and value.id in ('CArray', 'rawarray')):
|
||
return True
|
||
# list[type, count] — 仅匹配双参数 Tuple 切片（固定数组），
|
||
# 单参数 list[type] 不匹配（走泛型类 list[T] 路径）
|
||
if (isinstance(value, ast.Name)
|
||
and value.id == 'list'
|
||
and isinstance(annotation.slice, ast.Tuple)
|
||
and len(annotation.slice.elts) == 2):
|
||
return True
|
||
return False
|
||
|
||
|
||
def ExtractListFromBinOp(annotation: ast.expr) -> ast.Subscript | None:
|
||
"""从 BinOp(BitOr) 注解中提取 list[...] 部分，如 list[i32] | t.CPtr"""
|
||
if isinstance(annotation, ast.BinOp) and isinstance(annotation.op, ast.BitOr):
|
||
for side in (annotation.left, annotation.right):
|
||
if IsListAnnotation(side):
|
||
return side
|
||
return None
|
||
|
||
|
||
class ListAnnotationParseResult:
|
||
"""list[...] 注解解析结果。
|
||
|
||
is_pointer=True 表示单参数 list[type] → 指针模式；
|
||
is_pointer=False 表示 Tuple(2) 模式 list[type, count]。
|
||
"""
|
||
__slots__ = ('elem_type_node', 'count_node', 'is_pointer')
|
||
|
||
def __init__(self, elem_type_node: ast.expr, count_node: ast.expr | None, is_pointer: bool) -> None:
|
||
self.elem_type_node: ast.expr = elem_type_node
|
||
self.count_node: ast.expr | None = count_node
|
||
self.is_pointer: bool = is_pointer
|
||
|
||
|
||
def ParseListAnnotation(annotation: ast.expr) -> ListAnnotationParseResult | None:
|
||
"""解析 list[...] 注解，统一处理 list[type, count] / list[type] / BinOp(BitOr) 形式。
|
||
|
||
自动处理 BinOp(BitOr) 包裹（如 list[i32, N] | t.CPtr）。
|
||
返回 None 表示不是 list 注解。
|
||
"""
|
||
list_node: ast.expr = annotation
|
||
if isinstance(annotation, ast.BinOp) and isinstance(annotation.op, ast.BitOr):
|
||
extracted: ast.Subscript | None = ExtractListFromBinOp(annotation)
|
||
if extracted is None:
|
||
return None
|
||
list_node = extracted
|
||
if not IsListAnnotation(list_node):
|
||
return None
|
||
slice_node: ast.AST = list_node.slice
|
||
if isinstance(slice_node, ast.Tuple) and len(slice_node.elts) == 2:
|
||
return ListAnnotationParseResult(slice_node.elts[0], slice_node.elts[1], is_pointer=False)
|
||
# 单参数 list[type] → 指针模式
|
||
return ListAnnotationParseResult(slice_node, None, is_pointer=True)
|
||
|
||
|
||
def CheckAnnotationHasCInline(annotation_node: ast.AST | None) -> bool:
|
||
"""检测注解节点中是否包含 CInline"""
|
||
if annotation_node is None:
|
||
return False
|
||
if isinstance(annotation_node, ast.Attribute):
|
||
if hasattr(annotation_node.value, 'id') and IsTModule(annotation_node.value.id) and annotation_node.attr == 'CInline':
|
||
return True
|
||
if isinstance(annotation_node, ast.Name):
|
||
if annotation_node.id == 'CInline':
|
||
return True
|
||
if isinstance(annotation_node, ast.BinOp) and isinstance(annotation_node.op, ast.BitOr):
|
||
return CheckAnnotationHasCInline(annotation_node.left) or CheckAnnotationHasCInline(annotation_node.right)
|
||
return False
|
||
|
||
|
||
def ExtractTypeNameFromBinOp(annotation: ast.expr) -> str | None:
|
||
"""从 BinOp(BitOr) 注解左侧提取类型名（如 TypeA | TypeB → TypeA）。
|
||
|
||
返回 Name.id 或 Attribute.attr，非 BinOp 返回 None。
|
||
"""
|
||
if isinstance(annotation, ast.BinOp) and isinstance(annotation.op, ast.BitOr):
|
||
left: ast.expr = annotation.left
|
||
if isinstance(left, ast.Name):
|
||
return left.id
|
||
if isinstance(left, ast.Attribute):
|
||
return left.attr
|
||
return None
|
||
|
||
|
||
def FindStructNameInAnnotation(node: ast.AST, struct_names: dict | set) -> str | None:
|
||
"""递归在注解中查找结构体名（支持 BinOp(BitOr) 递归）。
|
||
|
||
检查 Name.id、Constant.str、Attribute.attr 是否在 struct_names 中。
|
||
支持泛型特化类型（如 GSList[Param]），构造特化名后检查是否在 struct_names 中。
|
||
"""
|
||
if isinstance(node, ast.Name) and node.id in struct_names:
|
||
return node.id
|
||
if isinstance(node, ast.Constant) and isinstance(node.value, str) and node.value in struct_names:
|
||
return node.value
|
||
if isinstance(node, ast.Attribute) and node.attr in struct_names:
|
||
return node.attr
|
||
if isinstance(node, ast.Subscript):
|
||
base_name: str | None = None
|
||
if isinstance(node.value, ast.Name):
|
||
base_name = node.value.id
|
||
elif isinstance(node.value, ast.Attribute):
|
||
base_name = node.value.attr
|
||
if base_name:
|
||
slice_node: ast.AST = node.slice
|
||
if hasattr(ast, 'Index') and isinstance(slice_node, ast.Index):
|
||
slice_node = slice_node.value
|
||
type_args: list[str] = []
|
||
if isinstance(slice_node, ast.Name):
|
||
# 裸 CPtr（i8*）与 t.CPtr 同规范化为 'ptr'，与
|
||
# _resolve_type_arg_str/_mangle_generic_class_name 一致
|
||
type_args.append('ptr' if slice_node.id == 'CPtr' else slice_node.id)
|
||
elif isinstance(slice_node, ast.Attribute):
|
||
type_args.append('ptr' if slice_node.attr == 'CPtr' else slice_node.attr)
|
||
elif isinstance(slice_node, ast.Tuple):
|
||
for elt in slice_node.elts:
|
||
if isinstance(elt, ast.Name):
|
||
type_args.append('ptr' if elt.id == 'CPtr' else elt.id)
|
||
elif isinstance(elt, ast.Attribute):
|
||
type_args.append('ptr' if elt.attr == 'CPtr' else elt.attr)
|
||
if type_args:
|
||
spec_name: str = base_name + ''.join(f'[{arg}]' for arg in type_args)
|
||
if spec_name in struct_names:
|
||
return spec_name
|
||
if base_name in struct_names:
|
||
return base_name
|
||
if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):
|
||
result: str | None = FindStructNameInAnnotation(node.left, struct_names)
|
||
if result:
|
||
return result
|
||
return FindStructNameInAnnotation(node.right, struct_names)
|
||
return None
|
||
|
||
|
||
def AnnotationContainsName(annotation: ast.AST, name: str) -> bool:
|
||
"""检查注解中是否包含指定名称（支持 BinOp(BitOr) 递归）。
|
||
|
||
同时检查 Attribute.attr 和 Name.id，比 ast.dump 字符串匹配更精确。
|
||
"""
|
||
if isinstance(annotation, ast.Attribute):
|
||
return annotation.attr == name
|
||
if isinstance(annotation, ast.Name):
|
||
return annotation.id == name
|
||
if isinstance(annotation, ast.BinOp) and isinstance(annotation.op, ast.BitOr):
|
||
return AnnotationContainsName(annotation.left, name) or AnnotationContainsName(annotation.right, name)
|
||
return False
|
||
|
||
|
||
def IsCPtrNode(node: ast.AST) -> bool:
|
||
"""检查节点是否为 t.CPtr 属性引用。"""
|
||
return (isinstance(node, ast.Attribute)
|
||
and isinstance(node.value, ast.Name)
|
||
and node.value.id == 't'
|
||
and node.attr == 'CPtr')
|