clean_cir

Cleans out debug information from clang IR output.

#!/usr/bin/env python3
"""Reduce CIR/ClangIR to globals, function declarations, and function bodies.

The cleaner retains:

  * Every cir.global definition.
  * Every cir.func declaration.
  * Every cir.func definition and its body.

It removes:

  * The enclosing module.
  * Type aliases such as !s32i = !cir.int<s, 32>.
  * Location definitions such as #loc1 = loc(...).
  * Function modifiers such as private, no_inline, and dso_local.
  * Function-level attributes, including target CPU and target features.
  * Trailing loc(...) debug-location annotations.

Function declarations and definitions are reduced to their symbol and
signature. For example:

    cir.func private @puts(!cir.ptr<!s8i>) -> !s32i loc(#loc1)

becomes:

    @puts(!cir.ptr<!s8i>) -> !s32i

Input files are rewritten in place. Use "-" to read from stdin and write the
cleaned CIR to stdout.

Examples:

    clean_cir input.cir
    clean_cir --dry-run input.cir
    clean_cir "*.cir"
    clang ... -emit-cir | clean_cir -
"""

import argparse
import glob
import re
import sys
from pathlib import Path

_FUNCTION_SYMBOL_RE = re.compile(r'@(?:"(?:\\.|[^"\\])*"|[-A-Za-z$._0-9]+)')
_FUNCTION_ATTRIBUTES_RE = re.compile(r"\battributes\s*\{")


def _brace_delta(text):
    """Count braces while ignoring strings and // comments."""
    depth = 0
    in_string = False
    escaped = False
    i = 0

    while i < len(text):
        ch = text[i]

        if in_string:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_string = False

            i += 1
            continue

        if ch == '"':
            in_string = True
        elif ch == "/" and i + 1 < len(text) and text[i + 1] == "/":
            break
        elif ch == "{":
            depth += 1
        elif ch == "}":
            depth -= 1

        i += 1

    return depth


def _matching_delimiter(text, opening, left, right):
    """Return the index of the delimiter matching text[opening]."""
    depth = 0
    in_string = False
    escaped = False
    i = opening

    while i < len(text):
        ch = text[i]

        if in_string:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_string = False

            i += 1
            continue

        if ch == '"':
            in_string = True
        elif ch == "/" and i + 1 < len(text) and text[i + 1] == "/":
            break
        elif ch == left:
            depth += 1
        elif ch == right:
            depth -= 1

            if depth == 0:
                return i

        i += 1

    return None


def _body_open_index(line):
    """Find the unmatched opening brace of a function body.

    Braces belonging to dictionaries such as:

        attributes {...}
        {llvm.noundef}

    are balanced on the same line. The function body's opening brace is the
    one that remains unmatched.
    """
    stack = []
    in_string = False
    escaped = False
    i = 0

    while i < len(line):
        ch = line[i]

        if in_string:
            if escaped:
                escaped = False
            elif ch == "\\":
                escaped = True
            elif ch == '"':
                in_string = False

            i += 1
            continue

        if ch == '"':
            in_string = True
        elif ch == "/" and i + 1 < len(line) and line[i + 1] == "/":
            break
        elif ch == "{":
            stack.append(i)
        elif ch == "}" and stack:
            stack.pop()

        i += 1

    return stack[0] if stack else None


def _strip_trailing_location(line):
    """Remove a trailing loc(...), including nested location expressions."""
    text = line.rstrip()

    location_starts = [match.start() for match in re.finditer(r"\bloc\s*\(", text)]

    for start in reversed(location_starts):
        opening = text.find("(", start)

        if opening == -1:
            continue

        closing = _matching_delimiter(
            text,
            opening,
            "(",
            ")",
        )

        if closing is not None and not text[closing + 1 :].strip():
            return text[:start].rstrip()

    return text


def _remove_function_attributes(header):
    """Remove a function-level `attributes {...}` dictionary."""
    match = _FUNCTION_ATTRIBUTES_RE.search(header)

    if match is None:
        return header

    opening = header.find("{", match.start())

    if opening == -1:
        return header

    closing = _matching_delimiter(
        header,
        opening,
        "{",
        "}",
    )

    if closing is None:
        return header

    before = header[: match.start()].rstrip()
    after = header[closing + 1 :].strip()

    if before and after:
        return before + " " + after

    return before or after


def _function_signature(line, body_open=None):
    """Extract the function symbol, parameters, and return type."""
    if body_open is None:
        header = line.rstrip()
    else:
        header = line[:body_open].rstrip()

    header = _strip_trailing_location(header)
    header = _remove_function_attributes(header)

    symbol_match = _FUNCTION_SYMBOL_RE.search(header)

    if symbol_match is None:
        raise ValueError(f"could not identify function symbol in: {line.strip()}")

    return header[symbol_match.start() :].strip()


def _remove_module_indent(line, module_indent):
    """Remove indentation contributed by the enclosing module."""
    if module_indent and line.startswith(module_indent):
        return line[len(module_indent) :]

    return line


def _clean_global(line):
    """Clean one cir.global definition."""
    return _strip_trailing_location(line).lstrip()


def _clean_declaration(line):
    """Clean one cir.func declaration."""
    return _function_signature(line)


def _clean_function(block):
    """Clean one complete cir.func definition."""
    first_line = block[0]
    body_open = _body_open_index(first_line)

    if body_open is None:
        raise ValueError("function definition has no body")

    signature = _function_signature(
        first_line,
        body_open=body_open,
    )

    module_indent = first_line[: len(first_line) - len(first_line.lstrip())]

    cleaned = [signature + " {"]

    for line in block[1:]:
        line = _strip_trailing_location(line)
        line = _remove_module_indent(line, module_indent)
        cleaned.append(line.rstrip())

    return "\n".join(cleaned).rstrip()


def _clean_text(original):
    """Extract and clean globals, declarations, and definitions."""
    lines = original.splitlines()
    output = []
    i = 0

    while i < len(lines):
        line = lines[i]
        stripped = line.lstrip()

        if stripped.startswith("cir.global "):
            output.append(_clean_global(line))
            i += 1
            continue

        if not stripped.startswith("cir.func "):
            i += 1
            continue

        body_open = _body_open_index(line)

        if body_open is None:
            output.append(_clean_declaration(line))
            i += 1
            continue

        depth = _brace_delta(line)
        block = [line]
        i += 1

        while i < len(lines) and depth > 0:
            block.append(lines[i])
            depth += _brace_delta(lines[i])
            i += 1

        if depth != 0:
            signature = _function_signature(
                block[0],
                body_open=body_open,
            )
            raise ValueError(f"unterminated function body for {signature}")

        output.append(_clean_function(block))

    if not output:
        raise ValueError("no cir.global or cir.func entries were found")

    return "\n\n".join(output) + "\n"


def clean(path, dry_run=False):
    """Clean a CIR file in place."""
    with open(path, encoding="utf-8") as file:
        original = file.read()

    try:
        cleaned = _clean_text(original)
    except ValueError as exc:
        print(
            f"Skipping {path}: {exc}",
            file=sys.stderr,
        )
        return False

    if dry_run:
        if cleaned != original:
            print(f"Would update {path}")
        else:
            print(f"No change {path}")

        return True

    with open(path, "w", encoding="utf-8") as file:
        file.write(cleaned)

    return True


def clean_stream(dry_run=False):
    """Clean CIR read from standard input."""
    original = sys.stdin.read()

    try:
        cleaned = _clean_text(original)
    except ValueError as exc:
        print(
            f"Could not clean stdin: {exc}",
            file=sys.stderr,
        )
        return False

    if dry_run:
        if cleaned != original:
            print("Would update stdin")
        else:
            print("No change stdin")

        return True

    sys.stdout.write(cleaned)
    return True


def _iter_input_files(patterns):
    """Expand input paths and glob patterns."""
    seen = set()

    for pattern in patterns:
        if pattern == "-":
            if "-" not in seen:
                seen.add("-")
                yield "-"

            continue

        matches = glob.glob(pattern)

        if not matches:
            matches = [pattern]

        for entry in matches:
            path = Path(entry)

            if not path.is_file():
                print(
                    f"Skipping {entry}: not a file",
                    file=sys.stderr,
                )
                continue

            resolved = str(path.resolve())

            if resolved in seen:
                continue

            seen.add(resolved)
            yield str(path)


def _parse_args():
    parser = argparse.ArgumentParser(
        description=(
            "retain CIR globals and functions while removing "
            "debug locations and function attributes"
        )
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="report files that would change without rewriting them",
    )
    parser.add_argument(
        "inputs",
        nargs="+",
        metavar="FILE|GLOB|-",
        help="CIR file, glob pattern, or - for standard input",
    )

    return parser.parse_args()


def main():
    args = _parse_args()
    had_inputs = False
    succeeded = True

    for input_path in _iter_input_files(args.inputs):
        had_inputs = True

        if input_path == "-":
            ok = clean_stream(
                dry_run=args.dry_run,
            )
        else:
            ok = clean(
                input_path,
                dry_run=args.dry_run,
            )

        succeeded = succeeded and ok

    if not had_inputs:
        print(
            "No input files matched.",
            file=sys.stderr,
        )
        return 1

    return 0 if succeeded else 1


if __name__ == "__main__":
    sys.exit(main())
download clean_cir