259 lines
9.6 KiB
Python
259 lines
9.6 KiB
Python
"""合并 IR 的元数据规整（sanitize）。
|
||
|
||
场景：多个模块的 .ll 文本被直接拼接成一个文件（每个模块都自带一小段
|
||
DWARF 元数据：!llvm.dbg.cu / !llvm.module.flags / !0 = !DIFile / ...
|
||
以及内联在 `define ... !dbg !N`、指令行上的 !N 引用）时，各模块的编号
|
||
元数据都从 !0 重新开始 —— 合并文件里同一个 !N 被定义多次、命名元数据
|
||
也被定义多次，LLVM 解析即报：
|
||
|
||
error: Metadata id is already used
|
||
error: ... !llvm.dbg.cu ... defined more than once
|
||
|
||
本模块提供 SanitizeIRMetadata(text) -> text：把全文的 `!N = ` 定义按
|
||
「元数据块」为单位重新编号成全局唯一 id，并同步改写全文引用；命名元数据
|
||
按名字聚合/去重成一份。无元数据的文本逐字节返回原样（严格 no-op），对
|
||
已规整的文本再跑一次结果一致（幂等）。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
|
||
# `!N = ` 形式的编号元数据定义行
|
||
_Definition = re.compile(r'^!(\d+) = ')
|
||
# 编号元数据引用
|
||
_Reference = re.compile(r'!(\d+)')
|
||
# `!name = ...` 形式的命名元数据（首字符为字母，故不会命中 `!N = `）
|
||
_NamedDefinition = re.compile(r'^!([A-Za-z][\w.]*)\s*=\s*(.*)$')
|
||
# `!{ ... }` 列表形式
|
||
_ListForm = re.compile(r'^!\{(.*)\}$', re.DOTALL)
|
||
# 纯 id 引用 `!N`
|
||
_PureId = re.compile(r'^!(\d+)$')
|
||
|
||
|
||
def _SplitTopLevel(Content: str) -> list[str]:
|
||
"""按顶层逗号切分 !{ ... } 的条目（跳过字符串字面量与嵌套括号内的逗号）。"""
|
||
Items: list[str] = []
|
||
Depth: int = 0
|
||
Buf: list[str] = []
|
||
i: int = 0
|
||
n: int = len(Content)
|
||
while i < n:
|
||
Ch: str = Content[i]
|
||
if Ch == '"':
|
||
j: int = i + 1
|
||
while j < n:
|
||
if Content[j] == '\\':
|
||
j += 2
|
||
continue
|
||
if Content[j] == '"':
|
||
j += 1
|
||
break
|
||
j += 1
|
||
Buf.append(Content[i:j])
|
||
i = j
|
||
continue
|
||
if Ch in '{[(':
|
||
Depth += 1
|
||
elif Ch in '}])':
|
||
Depth -= 1
|
||
elif Ch == ',' and Depth == 0:
|
||
Items.append(''.join(Buf).strip())
|
||
Buf = []
|
||
i += 1
|
||
continue
|
||
Buf.append(Ch)
|
||
i += 1
|
||
Tail: str = ''.join(Buf).strip()
|
||
if Tail:
|
||
Items.append(Tail)
|
||
return Items
|
||
|
||
|
||
def _ListItems(Rhs: str) -> list[str] | None:
|
||
"""`!{ a, b }` -> ['a', 'b']；非列表形式返回 None。"""
|
||
m = _ListForm.match(Rhs.strip())
|
||
if not m:
|
||
return None
|
||
return _SplitTopLevel(m.group(1))
|
||
|
||
|
||
def _RewriteReferences(Line: str, Mapping: dict[int, int]) -> str:
|
||
"""改写一行里对编号元数据的引用，跳过字符串字面量与 `;` 注释。"""
|
||
if '!' not in Line:
|
||
return Line
|
||
if '"' not in Line and ';' not in Line:
|
||
# 元数据行绝大多数不含字符串/注释：整体正则替换，避免逐字符扫描
|
||
return _Reference.sub(
|
||
lambda m: '!' + str(Mapping.get(int(m.group(1)), int(m.group(1)))), Line)
|
||
Out: list[str] = []
|
||
i: int = 0
|
||
n: int = len(Line)
|
||
while i < n:
|
||
Ch: str = Line[i]
|
||
if Ch == '"':
|
||
j: int = i + 1
|
||
while j < n:
|
||
if Line[j] == '\\':
|
||
j += 2
|
||
continue
|
||
if Line[j] == '"':
|
||
j += 1
|
||
break
|
||
j += 1
|
||
Out.append(Line[i:j])
|
||
i = j
|
||
continue
|
||
if Ch == ';':
|
||
Out.append(Line[i:])
|
||
break
|
||
if Ch == '!':
|
||
m = _Reference.match(Line, i)
|
||
if m:
|
||
New: int | None = Mapping.get(int(m.group(1)))
|
||
Out.append('!' + str(New) if New is not None else m.group(0))
|
||
i = m.end()
|
||
continue
|
||
Out.append(Ch)
|
||
i += 1
|
||
continue
|
||
Out.append(Ch)
|
||
i += 1
|
||
return ''.join(Out)
|
||
|
||
|
||
def SanitizeIRMetadata(text: str) -> str:
|
||
"""让拼接后的 IR 元数据合法：编号全局唯一 + 命名元数据只留一份。
|
||
|
||
无元数据（没有任何以 `!` 开头的行）时原样返回 —— 严格 no-op。
|
||
"""
|
||
if '!' not in text:
|
||
return text
|
||
Segments: list[str] = text.split('\n')
|
||
LineCount: int = len(Segments)
|
||
|
||
# 1) 元数据块 = 连续的以 '!' 开头的行。
|
||
# 每个模块的元数据（命名行 + 编号定义）是连续的一段；IR 行把它断开，
|
||
# 所以块边界即模块边界，块内共享同一套局部编号空间。
|
||
Blocks: list[tuple[int, int]] = [] # [Start, End)
|
||
i: int = 0
|
||
while i < LineCount:
|
||
if Segments[i].startswith('!'):
|
||
j: int = i
|
||
while j < LineCount and Segments[j].startswith('!'):
|
||
j += 1
|
||
Blocks.append((i, j))
|
||
i = j
|
||
else:
|
||
i += 1
|
||
if not Blocks:
|
||
return text
|
||
|
||
# 2) 每行归属的块：块内行归本块；块之间的行（IR、命名行）归「其后第一个
|
||
# 块」—— 因为每个模块的 IR 之后紧跟它自己的元数据块（含被引用的 !N）。
|
||
Owner: list[int] = [-1] * LineCount
|
||
PrevEnd: int = 0
|
||
for Index, (Start, End) in enumerate(Blocks):
|
||
for k in range(PrevEnd, End):
|
||
Owner[k] = Index
|
||
PrevEnd = End
|
||
|
||
# 3) 块内按定义出现顺序重新编号（局部 id -> 全局 id）
|
||
Lines: list[str] = list(Segments)
|
||
BlockMaps: list[dict[int, int]] = []
|
||
BlockIsIdentity: list[bool] = []
|
||
Counter: int = 0
|
||
for Start, End in Blocks:
|
||
Mapping: dict[int, int] = {}
|
||
Identity: bool = True
|
||
for k in range(Start, End):
|
||
m = _Definition.match(Lines[k])
|
||
if not m:
|
||
continue
|
||
Local: int = int(m.group(1))
|
||
Mapping[Local] = Counter
|
||
if Local != Counter:
|
||
Identity = False
|
||
Lines[k] = '!' + str(Counter) + ' = ' + Lines[k][m.end():]
|
||
Counter += 1
|
||
BlockMaps.append(Mapping)
|
||
BlockIsIdentity.append(Identity)
|
||
|
||
# 4) 改写全文引用（定义体右值、define/指令行的 !dbg、命名元数据列表）
|
||
# 定义行左侧的 id 是上一步刚分配的新编号，绝不能再当引用改写：本块的
|
||
# 新编号区间可能与它自己的旧局部编号区间重叠（例如新编号恰好也是本块
|
||
# 旧 id 843 → Mapping[843]=1686），一旦改写左侧就会出现两个定义共用同一
|
||
# id，opt/llc 随即报 "Metadata id is already used"（本 bug 实测残留 276 个
|
||
# 重复 id）。故定义行只改写 '=' 之后的右值。
|
||
for k in range(LineCount):
|
||
Index: int = Owner[k]
|
||
if Index < 0 or BlockIsIdentity[Index] or not BlockMaps[Index]:
|
||
continue
|
||
Line: str = Lines[k]
|
||
Head: re.Match | None = _Definition.match(Line)
|
||
if Head:
|
||
Fixed: str = Line[:Head.end()] + _RewriteReferences(Line[Head.end():], BlockMaps[Index])
|
||
else:
|
||
Fixed = _RewriteReferences(Line, BlockMaps[Index])
|
||
if Fixed != Line:
|
||
Lines[k] = Fixed
|
||
|
||
# 5) 定义体索引（新 id -> 右值文本），供命名元数据判重
|
||
BodyOf: dict[int, str] = {}
|
||
for k in range(LineCount):
|
||
m = _Definition.match(Lines[k])
|
||
if m:
|
||
BodyOf[int(m.group(1))] = Lines[k][m.end():].strip()
|
||
|
||
# 6) 命名元数据：同名多份合并成一份。
|
||
# - !llvm.dbg.cu：全部条目聚合（每个 DICompileUnit 都必须被列出，不能去重）
|
||
# - 其它（!llvm.module.flags、!llvm.ident ...）：条目内容（id 引用解析成
|
||
# 节点正文）完全相同的多份只留一份，不同则聚合。
|
||
NamedLines: dict[str, list[int]] = {}
|
||
for k in range(LineCount):
|
||
m = _NamedDefinition.match(Lines[k])
|
||
if m:
|
||
NamedLines.setdefault(m.group(1), []).append(k)
|
||
|
||
Dropped: set[int] = set()
|
||
for Name, Indices in NamedLines.items():
|
||
if len(Indices) == 1 and Name != 'llvm.dbg.cu':
|
||
continue
|
||
First: int = Indices[0]
|
||
Cr: str = Segments[First][len(Segments[First].rstrip('\r')):]
|
||
Items: list[str] = []
|
||
SeenItems: set[str] = set()
|
||
SeenBodies: set[tuple[str, ...]] = set()
|
||
Single: bool = True
|
||
for k in Indices:
|
||
Rhs: str = _NamedDefinition.match(Lines[k]).group(2).strip()
|
||
ListItems: list[str] | None = _ListItems(Rhs)
|
||
if ListItems is None:
|
||
Single = False
|
||
break
|
||
Bodies: tuple[str, ...] = tuple(
|
||
BodyOf.get(int(_PureId.match(Item).group(1)), Item) if _PureId.match(Item) else Item
|
||
for Item in ListItems
|
||
)
|
||
if Name != 'llvm.dbg.cu':
|
||
if Bodies in SeenBodies:
|
||
continue
|
||
SeenBodies.add(Bodies)
|
||
for Item in ListItems:
|
||
if Item not in SeenItems:
|
||
SeenItems.add(Item)
|
||
Items.append(Item)
|
||
if not Single:
|
||
# 非 `!{ ... }` 列表形式（实际 IR 里不出现）：保留首份，其余丢弃，
|
||
# 避免出现同名命名元数据被定义多次。
|
||
Dropped.update(Indices[1:])
|
||
continue
|
||
Rebuilt: str = '!%s = !{ %s }%s' % (Name, ', '.join(Items), Cr)
|
||
if Rebuilt != Segments[First]:
|
||
Lines[First] = Rebuilt
|
||
Dropped.update(Indices[1:])
|
||
|
||
if not Dropped and Lines == Segments:
|
||
return text
|
||
return '\n'.join(Line for k, Line in enumerate(Lines) if k not in Dropped)
|