464 lines
22 KiB
Python
464 lines
22 KiB
Python
from __future__ import annotations
|
||
import re
|
||
import ast
|
||
import llvmlite.ir as ir
|
||
from llvmlite.ir.values import _Undefined
|
||
|
||
|
||
class StructGenMixin:
|
||
|
||
def _get_align(self, llvm_type):
|
||
if isinstance(llvm_type, ir.IntType):
|
||
return max(1, llvm_type.width // 8)
|
||
if isinstance(llvm_type, ir.FloatType):
|
||
return 4
|
||
if isinstance(llvm_type, ir.DoubleType):
|
||
return 8
|
||
if isinstance(llvm_type, ir.PointerType):
|
||
return self.ptr_size
|
||
if isinstance(llvm_type, ir.ArrayType):
|
||
return self._get_align(llvm_type.element)
|
||
if isinstance(llvm_type, ir.LiteralStructType):
|
||
return max((self._get_align(f) for f in llvm_type.elements), default=1)
|
||
if isinstance(llvm_type, ir.IdentifiedStructType):
|
||
if llvm_type.elements is None or len(llvm_type.elements) == 0:
|
||
return self.ptr_size # opaque struct, assume pointer alignment
|
||
return max((self._get_align(f) for f in llvm_type.elements), default=1)
|
||
if isinstance(llvm_type, ir.VoidType):
|
||
return 0
|
||
return 4
|
||
|
||
def _get_struct_size(self, llvm_type):
|
||
if isinstance(llvm_type, ir.IntType):
|
||
return max(1, llvm_type.width // 8)
|
||
if isinstance(llvm_type, ir.FloatType):
|
||
return 4
|
||
if isinstance(llvm_type, ir.DoubleType):
|
||
return 8
|
||
if isinstance(llvm_type, ir.PointerType):
|
||
return self.ptr_size
|
||
if isinstance(llvm_type, ir.ArrayType):
|
||
return self._get_struct_size(llvm_type.element) * llvm_type.count
|
||
if isinstance(llvm_type, ir.LiteralStructType):
|
||
total = 0
|
||
max_align = 1
|
||
for f in llvm_type.fields:
|
||
fa = self._get_align(f)
|
||
fs = self._get_struct_size(f)
|
||
max_align = max(max_align, fa)
|
||
if fa > 1:
|
||
total = (total + fa - 1) // fa * fa
|
||
total += fs
|
||
if max_align > 1:
|
||
total = (total + max_align - 1) // max_align * max_align
|
||
return total
|
||
if isinstance(llvm_type, ir.IdentifiedStructType):
|
||
if llvm_type.elements is None or len(llvm_type.elements) == 0:
|
||
return self.ptr_size # opaque struct, assume pointer size
|
||
total = 0
|
||
max_align = 1
|
||
for f in llvm_type.elements:
|
||
fa = self._get_align(f)
|
||
fs = self._get_struct_size(f)
|
||
max_align = max(max_align, fa)
|
||
if fa > 1:
|
||
total = (total + fa - 1) // fa * fa
|
||
total += fs
|
||
if max_align > 1:
|
||
total = (total + max_align - 1) // max_align * max_align
|
||
return total
|
||
return 4
|
||
|
||
@staticmethod
|
||
def _strip_unnecessary_quotes(ir_str):
|
||
def _unquote(m):
|
||
prefix = m.group(1)
|
||
name = m.group(2)
|
||
if '.' in name:
|
||
return m.group(0)
|
||
if re.match(r'^[a-zA-Z_$][a-zA-Z0-9_$]*$', name):
|
||
return prefix + name
|
||
return m.group(0)
|
||
return re.sub(r'([@%])"([^"]+)"', _unquote, ir_str)
|
||
|
||
def finalize(self):
|
||
self._set_Vtable_initializers()
|
||
self._fix_global_var_init()
|
||
try:
|
||
ir_text = str(self.module)
|
||
except TypeError as e:
|
||
for func in self.module.functions:
|
||
for block in func.blocks:
|
||
for instr in block.instructions:
|
||
try:
|
||
str(instr)
|
||
except TypeError as te:
|
||
pass
|
||
raise
|
||
ir_text = self._strip_unnecessary_quotes(ir_text)
|
||
struct_defs = []
|
||
struct_names = set()
|
||
for ClassName, struct_type in self.structs.items():
|
||
if isinstance(struct_type, ir.IdentifiedStructType):
|
||
sname = struct_type.name
|
||
if sname[0].isdigit() or '.' in sname or not sname.replace('_', '').replace('.', '').isalnum():
|
||
sname_ir = f'%"{sname}"'
|
||
else:
|
||
sname_ir = f'%{sname}'
|
||
if struct_type.elements is None or len(struct_type.elements) == 0:
|
||
if ClassName not in self.typedef_int_widths:
|
||
struct_names.add(sname)
|
||
struct_defs.append(f'{sname_ir} = type opaque')
|
||
continue
|
||
elems = ', '.join(str(e) for e in struct_type.elements)
|
||
if struct_type.packed:
|
||
struct_defs.append(f'{sname_ir} = type <{{ {elems} }}>')
|
||
else:
|
||
struct_defs.append(f'{sname_ir} = type {{ {elems} }}')
|
||
struct_names.add(sname)
|
||
if struct_defs:
|
||
lines = ir_text.split('\n')
|
||
filtered = []
|
||
for line in lines:
|
||
stripped = line.strip()
|
||
skip = False
|
||
for name in struct_names:
|
||
if name[0].isdigit() or '.' in name or not name.replace('_', '').replace('.', '').isalnum():
|
||
patterns = [f'%"{name}" = type', f'%"{name}"=type', f'%{name} = type', f'%{name}=type']
|
||
else:
|
||
patterns = [f'%{name} = type', f'%{name}= type', f'%{name}=type']
|
||
for p in patterns:
|
||
if stripped.startswith(p):
|
||
skip = True
|
||
break
|
||
if skip:
|
||
break
|
||
if not skip:
|
||
filtered.append(line)
|
||
ir_text = '\n'.join(filtered)
|
||
header_end = ir_text.find('\ndefine')
|
||
if header_end == -1:
|
||
header_end = ir_text.find('\ndeclare')
|
||
if header_end == -1:
|
||
header_end = ir_text.find('\n@')
|
||
if header_end == -1:
|
||
header_end = len(ir_text)
|
||
insert_pos = ir_text.rfind('\n', 0, header_end)
|
||
if insert_pos == -1:
|
||
insert_pos = header_end
|
||
struct_block = '\n'.join(struct_defs) + '\n'
|
||
ir_text = ir_text[:insert_pos + 1] + struct_block + ir_text[insert_pos + 1:]
|
||
for name in struct_names:
|
||
if name[0].isdigit() or '.' in name:
|
||
ir_text = re.sub(
|
||
r'%' + re.escape(name) + r'(?=[^a-zA-Z0-9_$\.]|$)',
|
||
r'%"' + name + r'"',
|
||
ir_text
|
||
)
|
||
|
||
return ir_text
|
||
|
||
def _collect_classes_and_methods(self, ast_nodes, parent=None):
|
||
current_class = None
|
||
for node in ast_nodes:
|
||
if isinstance(node, str):
|
||
continue
|
||
node.parent = parent
|
||
node_type = type(node).__name__
|
||
if node_type == 'Struct':
|
||
if hasattr(node, 'name'):
|
||
current_class = node.name
|
||
self.class_methods[current_class] = []
|
||
self.class_members[current_class] = []
|
||
if hasattr(node, 'decls') and node.decls:
|
||
for decl in node.decls:
|
||
if type(decl).__name__ == 'Decl' and hasattr(decl, 'name'):
|
||
member_type = self._get_type(decl.type) if hasattr(decl, 'type') else ir.IntType(32)
|
||
self.class_members[current_class].append((decl.name, member_type))
|
||
elif node_type == 'FuncDef':
|
||
if current_class and current_class in self.class_methods:
|
||
MethodName = node.decl.name if hasattr(node.decl, 'name') else 'unknown_method'
|
||
if not MethodName.startswith(f"{current_class}."):
|
||
MethodName = f"{current_class}.{MethodName}"
|
||
self.class_methods[current_class].append(MethodName)
|
||
if hasattr(node, 'body') and isinstance(node.body, list):
|
||
self._collect_classes_and_methods(node.body, parent=node)
|
||
elif hasattr(node, 'BlockItems') and node.BlockItems:
|
||
self._collect_classes_and_methods(node.BlockItems, parent=node)
|
||
|
||
def _generate_structs(self):
|
||
all_classes = set(self.class_methods.keys()) | set(self.class_members.keys())
|
||
|
||
preserved_structs = {}
|
||
for name, st in self.structs.items():
|
||
if name not in all_classes:
|
||
preserved_structs[name] = st
|
||
elif isinstance(st, ir.IdentifiedStructType):
|
||
preserved_structs[name] = st
|
||
|
||
self.structs.clear()
|
||
|
||
for name, st in preserved_structs.items():
|
||
self.structs[name] = st
|
||
|
||
# Step 1: 先为所有类创建空的 struct，确保后面可以正确引用
|
||
all_classes = set(self.class_methods.keys()) | set(self.class_members.keys())
|
||
for ClassName in all_classes:
|
||
is_packed = ClassName in self.class_packed
|
||
if hasattr(self, 'SymbolTable') and self.SymbolTable and self.SymbolTable.has(ClassName):
|
||
Entry = self.SymbolTable[ClassName]
|
||
if Entry.IsExceptionClass:
|
||
continue
|
||
if Entry.IsEnum:
|
||
if not Entry.IsRenum:
|
||
continue
|
||
if Entry.IsTypedef:
|
||
if Entry.BaseType and not isinstance(Entry.BaseType, (type(None),)):
|
||
import lib.includes.t as _t
|
||
if not isinstance(Entry.BaseType, (_t._CTypedef,)):
|
||
continue
|
||
if ClassName in self.structs:
|
||
if is_packed:
|
||
self.structs[ClassName].packed = True
|
||
continue
|
||
mangled_name = self._mangle_name(ClassName)
|
||
struct_type = ir.IdentifiedStructType(self.module, mangled_name, packed=is_packed)
|
||
self.structs[ClassName] = struct_type
|
||
|
||
# Step 2: 解析跨模块的成员类型
|
||
for ClassName, annotations in self.class_member_annotations.items():
|
||
if ClassName not in self.class_members:
|
||
continue
|
||
for i, (member_name, member_type) in enumerate(self.class_members[ClassName]):
|
||
if member_name in annotations:
|
||
annotation = annotations[member_name]
|
||
ResolvedType = None
|
||
|
||
if isinstance(annotation, ast.Attribute):
|
||
ClassTypeName = annotation.attr
|
||
if ClassTypeName in self.structs:
|
||
ResolvedType = self.structs[ClassTypeName]
|
||
elif isinstance(annotation, ast.BinOp) and isinstance(annotation.op, ast.BitOr):
|
||
left = annotation.left
|
||
if isinstance(left, ast.Attribute):
|
||
ClassTypeName = left.attr
|
||
if ClassTypeName in self.structs:
|
||
ResolvedType = self.structs[ClassTypeName]
|
||
|
||
if ResolvedType:
|
||
self.class_members[ClassName][i] = (member_name, ResolvedType)
|
||
|
||
struct_name_map = {}
|
||
for struct_name, struct in self.structs.items():
|
||
try:
|
||
struct_name_map[struct.name] = struct
|
||
except AssertionError:
|
||
struct_name_map[struct_name] = struct
|
||
orig = getattr(struct, '_original_name', None)
|
||
if orig:
|
||
struct_name_map[orig] = struct
|
||
|
||
# Step 3: 更新成员类型，确保使用正确的 mangled struct 引用
|
||
for ClassName in all_classes:
|
||
if ClassName not in self.class_members:
|
||
continue
|
||
for i, (member_name, member_type) in enumerate(self.class_members[ClassName]):
|
||
if isinstance(member_type, ir.PointerType):
|
||
pointee = member_type.pointee
|
||
if isinstance(pointee, ir.IdentifiedStructType):
|
||
try:
|
||
pname = pointee.name
|
||
except AssertionError:
|
||
pname = None
|
||
struct = struct_name_map.get(pname) if pname else None
|
||
if struct is None:
|
||
orig = getattr(pointee, '_original_name', None)
|
||
struct = struct_name_map.get(orig) if orig else None
|
||
if struct is not None:
|
||
self.class_members[ClassName][i] = (member_name, ir.PointerType(struct))
|
||
else:
|
||
raw_name = pname or ''
|
||
if raw_name.startswith('struct.'):
|
||
raw_name = raw_name[7:]
|
||
new_struct = self._get_or_create_struct(raw_name)
|
||
self.class_members[ClassName][i] = (member_name, ir.PointerType(new_struct))
|
||
elif isinstance(member_type, ir.IdentifiedStructType):
|
||
try:
|
||
mname = member_type.name
|
||
except AssertionError:
|
||
mname = None
|
||
struct = struct_name_map.get(mname) if mname else None
|
||
if struct is None:
|
||
orig = getattr(member_type, '_original_name', None)
|
||
struct = struct_name_map.get(orig) if orig else None
|
||
if struct is not None:
|
||
self.class_members[ClassName][i] = (member_name, struct)
|
||
else:
|
||
raw_name = mname or ''
|
||
if raw_name.startswith('struct.'):
|
||
raw_name = raw_name[7:]
|
||
new_struct = self._get_or_create_struct(raw_name)
|
||
self.class_members[ClassName][i] = (member_name, new_struct)
|
||
|
||
# Step 4: 创建最终的 structs 并设置正确的成员类型
|
||
for ClassName in all_classes:
|
||
if hasattr(self, 'SymbolTable') and self.SymbolTable and self.SymbolTable.has(ClassName):
|
||
Entry = self.SymbolTable[ClassName]
|
||
if Entry.IsExceptionClass:
|
||
continue
|
||
if Entry.IsEnum:
|
||
if not Entry.IsRenum:
|
||
continue
|
||
existing_st = self.structs.get(ClassName)
|
||
if isinstance(existing_st, ir.IdentifiedStructType) and existing_st.elements is not None and len(existing_st.elements) > 0:
|
||
# 即使 struct 已有元素，也需要修正 packed 属性
|
||
if ClassName in self.class_packed and not existing_st.packed:
|
||
existing_st.packed = True
|
||
continue
|
||
has_vtable = ClassName in self.class_vtable or ClassName in self._cross_module_vtable_classes
|
||
if not has_vtable:
|
||
p = self.class_parent.get(ClassName)
|
||
while p:
|
||
if p in self.class_vtable or p in self._cross_module_vtable_classes:
|
||
has_vtable = True
|
||
self.class_vtable.add(ClassName)
|
||
break
|
||
p = self.class_parent.get(p)
|
||
has_bitfield = ClassName in self.class_member_bitfields and any(
|
||
self.class_member_bitfields[ClassName].get(name, 0) > 0
|
||
for name, _ in self.class_members.get(ClassName, [])
|
||
)
|
||
all_members = self.class_members.get(ClassName, [])
|
||
all_bitfield = all(
|
||
self.class_member_bitfields.get(ClassName, {}).get(name, 0) > 0
|
||
for name, _ in all_members
|
||
) if all_members else False
|
||
|
||
mangled_name = self._mangle_name(ClassName)
|
||
|
||
if has_bitfield and all_bitfield:
|
||
total_bits = sum(self.class_member_bitfields.get(ClassName, {}).get(name, 0)
|
||
for name, _ in all_members)
|
||
if total_bits <= 8:
|
||
storage_type = ir.IntType(8)
|
||
elif total_bits <= 16:
|
||
storage_type = ir.IntType(16)
|
||
elif total_bits <= 32:
|
||
storage_type = ir.IntType(32)
|
||
else:
|
||
storage_type = ir.IntType(64)
|
||
member_types = [storage_type]
|
||
bitfield_offsets = {}
|
||
current_bit_offset = 0
|
||
for name, _ in all_members:
|
||
bit_width = self.class_member_bitfields.get(ClassName, {}).get(name, 0)
|
||
bitfield_offsets[name] = current_bit_offset
|
||
current_bit_offset += bit_width
|
||
self.class_member_bitoffsets[ClassName] = bitfield_offsets
|
||
else:
|
||
member_types = []
|
||
if has_vtable:
|
||
member_types.append(ir.PointerType(ir.IntType(8)))
|
||
if ClassName in self.class_members:
|
||
for _, member_type in self.class_members[ClassName]:
|
||
if member_type is None:
|
||
member_type = ir.IntType(32)
|
||
member_types.append(member_type)
|
||
else:
|
||
member_types = [ir.IntType(8)]
|
||
|
||
# 获取之前创建的 struct 并设置成员
|
||
struct_type = self.structs[ClassName]
|
||
if ClassName in self.class_packed:
|
||
struct_type.packed = True
|
||
# 只在未设置 body 时设置，避免重复定义错误
|
||
if struct_type.elements is None:
|
||
struct_type.set_body(*member_types)
|
||
|
||
def _create_Vtable_globals(self):
|
||
for ClassName, methods in self.class_methods.items():
|
||
if ClassName in self.Vtables:
|
||
continue
|
||
if not methods:
|
||
continue
|
||
if ClassName not in self.class_vtable and ClassName not in self._cross_module_vtable_classes:
|
||
continue
|
||
VtableType = ir.ArrayType(ir.PointerType(ir.IntType(8)), len(methods))
|
||
Vtable = ir.GlobalVariable(self.module, VtableType, name=self._mangle_name(f"{ClassName}_Vtable"))
|
||
Vtable.initializer = ir.Constant(VtableType, [ir.Constant(ir.PointerType(ir.IntType(8)), None)] * len(methods))
|
||
Vtable.linkage = 'internal'
|
||
self.Vtables[ClassName] = Vtable
|
||
|
||
def _fix_global_var_init(self):
|
||
for gv_name, gv in list(self.module.globals.items()):
|
||
if not isinstance(gv, ir.GlobalVariable):
|
||
continue
|
||
if gv.linkage in ('external', 'extern_weak'):
|
||
continue
|
||
needs_fix = False
|
||
if gv.initializer is None:
|
||
needs_fix = True
|
||
elif hasattr(gv.initializer, 'constant') and isinstance(gv.initializer.constant, _Undefined):
|
||
needs_fix = True
|
||
elif isinstance(gv.initializer, ir.Constant):
|
||
if isinstance(gv.initializer.constant, list):
|
||
if any(hasattr(c, 'constant') and isinstance(c.constant, _Undefined) for c in gv.initializer.constant):
|
||
needs_fix = True
|
||
if not needs_fix:
|
||
continue
|
||
gv.initializer = ir.Constant(gv.value_type, None)
|
||
if gv_name.startswith('str_const_') or gv_name.endswith('_Vtable') or gv_name.endswith('_str'):
|
||
gv.linkage = 'internal'
|
||
|
||
def _zero_value_for_type(self, var_type):
|
||
if isinstance(var_type, ir.IntType):
|
||
return ir.Constant(var_type, 0)
|
||
elif isinstance(var_type, (ir.FloatType, ir.DoubleType)):
|
||
return ir.Constant(var_type, 0.0)
|
||
elif isinstance(var_type, ir.PointerType):
|
||
return ir.Constant(var_type, None)
|
||
elif isinstance(var_type, ir.ArrayType):
|
||
elem_zv = self._zero_value_for_type(var_type.element)
|
||
if elem_zv is None:
|
||
return None
|
||
return ir.Constant(var_type, [elem_zv] * var_type.count)
|
||
elif isinstance(var_type, ir.BaseStructType):
|
||
if var_type.elements is not None and len(var_type.elements) > 0:
|
||
elem_zvs = [self._zero_value_for_type(m) for m in var_type.elements]
|
||
if any(zv is None for zv in elem_zvs):
|
||
return None
|
||
return ir.Constant(var_type, elem_zvs)
|
||
return None
|
||
return None
|
||
|
||
def _set_Vtable_initializers(self):
|
||
for ClassName, methods in self.class_methods.items():
|
||
if ClassName not in self.Vtables:
|
||
continue
|
||
if not methods:
|
||
continue
|
||
struct_type = self.structs.get(ClassName)
|
||
if not struct_type:
|
||
continue
|
||
Vtable = self.Vtables[ClassName]
|
||
Vtable_entries = []
|
||
for MethodName in methods:
|
||
short_name = MethodName.split('.')[-1] if '.' in MethodName else MethodName
|
||
full_MethodName = f"{ClassName}.{short_name}"
|
||
func = None
|
||
if full_MethodName in self.functions:
|
||
func = self.functions[full_MethodName]
|
||
if func is None:
|
||
parent_cls = self.class_parent.get(ClassName)
|
||
while parent_cls:
|
||
parent_full = f"{parent_cls}.{short_name}"
|
||
if parent_full in self.functions:
|
||
func = self.functions[parent_full]
|
||
break
|
||
parent_cls = self.class_parent.get(parent_cls)
|
||
if func is None:
|
||
Vtable_entries.append(ir.Constant(ir.PointerType(ir.IntType(8)), None))
|
||
continue
|
||
Vtable_entries.append(func.bitcast(ir.PointerType(ir.IntType(8))))
|
||
Vtable.initializer = ir.Constant(Vtable.type.pointee, Vtable_entries)
|