"""合并 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+) = ')
# `!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+)$')
# Decimal digit set for the no-regex fast rewrite path.
_DIGITS: frozenset[str] = frozenset('0123456789')


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 _RewritePlain(Line: str, StrMap: dict[str, str]) -> str:
    """Rewrite `!N` references in a line that has no `"` and no `;`.

    Equivalent to a regex substitution of `!` + digits but without the regex
    engine callback: split on '!' (C-level), map each digit run by full token,
    keep unknown ids verbatim.
    """
    Parts: list[str] = Line.split('!')
    if len(Parts) == 1:
        return Line
    Out: list[str] = [Parts[0]]
    Ap = Out.append
    Get = StrMap.get
    for Part in Parts[1:]:
        d: int = 0
        n: int = len(Part)
        while d < n and Part[d] in _DIGITS:
            d += 1
        if not d:
            Ap('!')
            Ap(Part)
            continue
        Digits: str = Part[:d]
        # Leading zeros normalise the way int()/str() did in the old path.
        Tok: str = '!' + Digits if (d == 1 or Digits[0] != '0') else '!' + str(int(Digits))
        New: str | None = Get(Tok)
        Ap(Tok if New is None else New)
        Ap(Part[d:])
    return ''.join(Out)


def _RewriteQuoted(Line: str, StrMap: dict[str, str]) -> str:
    """Rewrite `!N` refs in a line that contains `"` and/or `;`.

    String literals and the trailing `;` comment are copied verbatim; the
    remaining (quote/comment free) chunks go through `_RewritePlain`.  Uses
    C-level str.find instead of the old per-character Python scan.
    """
    Out: list[str] = []; Ap = Out.append
    i: int = 0
    n: int = len(Line)
    while i < n:
        q: int = Line.find('"', i)
        s: int = Line.find(';', i, q if q >= 0 else n)
        if s >= 0:
            Ap(_RewritePlain(Line[i:s], StrMap))
            Ap(Line[s:])          # ';' comment: copy verbatim to end of line
            break
        if q < 0:
            Ap(_RewritePlain(Line[i:], StrMap))
            break
        Ap(_RewritePlain(Line[i:q], StrMap))
        j: int = q + 1
        while j < n:
            if Line[j] == '\\':
                j += 2
                continue
            if Line[j] == '"':
                j += 1
                break
            j += 1
        Ap(Line[q:j])             # string literal: copy verbatim
        i = j
    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 + 1
            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]] = []
    BlockStrMaps: list[dict[str, str]] = []
    BlockIsIdentity: list[bool] = []
    # Definition line index -> (new global id, length of its rewritten "!N = " prefix).
    # Lets step 4/5 skip a per-line regex to re-detect definition lines.
    DefInfo: dict[int, tuple[int, int]] = {}
    # Named metadata lines are collected here (only block lines can start with '!').
    NamedLines: dict[str, list[int]] = {}
    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:
                nm = _NamedDefinition.match(Lines[k])
                if nm is not None:
                    NamedLines.setdefault(nm.group(1), []).append(k)
                continue
            Local: int = int(m.group(1))
            Mapping[Local] = Counter
            if Local != Counter:
                Identity = False
            Prefix: str = '!' + str(Counter) + ' = '
            DefInfo[k] = (Counter, len(Prefix))
            Lines[k] = Prefix + Lines[k][m.end():]
            Counter += 1
        BlockMaps.append(Mapping)
        BlockStrMaps.append({'!' + str(LocalId): '!' + str(GlobId)
                             for LocalId, GlobId in Mapping.items()})
        BlockIsIdentity.append(Identity)

    # 4) 改写全文引用（定义体右值、define/指令行的 !dbg、命名元数据列表）
    #    定义行左侧的 id 是上一步刚分配的新编号，绝不能再当引用改写：本块的
    #    新编号区间可能与它自己的旧局部编号区间重叠（例如新编号恰好也是本块
    #    旧 id 843 → Mapping[843]=1686），一旦改写左侧就会出现两个定义共用同一
    #    id，opt/llc 随即报 "Metadata id is already used"（本 bug 实测残留 276 个
    #    重复 id）。故定义行只改写 '=' 之后的右值。
    GetDef = DefInfo.get
    for k in range(LineCount):
        Index: int = Owner[k]
        if Index < 0 or BlockIsIdentity[Index] or not BlockMaps[Index]:
            continue
        Line: str = Lines[k]
        if '!' not in Line:
            continue
        StrMap: dict[str, str] = BlockStrMaps[Index]
        Def: tuple[int, int] | None = GetDef(k)
        if Def is not None:
            Split: int = Def[1]
            Tail: str = Line[Split:]
            Fixed: str = Line[:Split] + (_RewritePlain(Tail, StrMap)
                                         if ('"' not in Tail and ';' not in Tail)
                                         else _RewriteQuoted(Tail, StrMap))
        else:
            Fixed = (_RewritePlain(Line, StrMap)
                     if ('"' not in Line and ';' not in Line)
                     else _RewriteQuoted(Line, StrMap))
        if Fixed != Line:
            Lines[k] = Fixed

    # 5) 定义体索引（新 id -> 右值文本），供命名元数据判重。
    #    只遍历定义行（步骤 3 已登记），不再逐行做正则匹配。
    BodyOf: dict[int, str] = {}
    for k, (NewId, Split) in DefInfo.items():
        BodyOf[NewId] = Lines[k][Split:].strip()

    # 6) 命名元数据：同名多份合并成一份。
    #    - !llvm.dbg.cu：全部条目聚合（每个 DICompileUnit 都必须被列出，不能去重）
    #    - 其它（!llvm.module.flags、!llvm.ident ...）：条目内容（id 引用解析成
    #      节点正文）完全相同的多份只留一份，不同则聚合。
    #    NamedLines 已在步骤 3 随块扫描一并收集（命名行的名字不受引用改写影响）。
    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
        IsDbgCu: bool = Name == 'llvm.dbg.cu'
        for k in Indices:
            NamedMatch = _NamedDefinition.match(Lines[k])
            Rhs: str = (NamedMatch.group(2) if NamedMatch is not None else '').strip()
            ListItems: list[str] | None = _ListItems(Rhs)
            if ListItems is None:
                Single = False
                break
            if not IsDbgCu:
                # !llvm.dbg.cu never de-dupes by body, so skip resolving bodies for it.
                Bodies: list[str] = []
                for Item in ListItems:
                    PureMatch = _PureId.match(Item)
                    Bodies.append(BodyOf.get(int(PureMatch.group(1)), Item)
                                  if PureMatch else Item)
                Key: tuple[str, ...] = tuple(Bodies)
                if Key in SeenBodies:
                    continue
                SeenBodies.add(Key)
            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)