clean_asm

#!/usr/bin/env python3
"""Strip ELF metadata from assembly, keeping only labels and instructions."""

import glob
import re
import sys
from pathlib import Path

def _detect_comment_prefix(lines):
    for line in lines:
        if line.lstrip().startswith("//"):
            return "//"
    return "#"


def _strip_comment(line, comment="//"):
    if comment in line:
        return line[: line.index(comment)]
    return line


def _is_label_line(line, comment="#"):
    text = _strip_comment(line, comment).strip()
    return text.endswith(":")


def _defined_label_name(line, comment="#"):
    text = _strip_comment(line, comment).strip()
    if not text.endswith(":"):
        return None
    return text[:-1].strip()


def _collect_referenced_labels(lines, comment="#"):
    refs = set()
    num_refs = set()
    for line in lines:
        text = _strip_comment(line, comment).strip()
        if not text or text.endswith(":") or text.startswith("."):
            continue
        parts = text.split(None, 1)
        if len(parts) < 2:
            continue
        operands = parts[1]
        # Numeric local label refs like 1f / 2b.
        for m in re.finditer(r"\b(\d+)[fb]\b", operands):
            num_refs.add(m.group(1))
        # Symbol refs in operands, e.g. .LBB0_1(%rip), foo, .Ltmp42
        for m in re.finditer(r"(?<![%$])([._A-Za-z][._A-Za-z0-9$]*)", operands):
            refs.add(m.group(1))
    return refs, num_refs


_DATA_DIRECTIVES = {
    ".byte", ".2byte", ".4byte", ".8byte",
    ".word", ".hword", ".short", ".long", ".quad", ".octa",
    ".ascii", ".asciz", ".string",
    ".float", ".double",
    ".zero", ".skip",
}

# Assembler mode directives that affect instruction encoding/parsing.
_MODE_DIRECTIVES = {
    ".set",
    ".option",
}


def _clean_lines(lines):
    comment = _detect_comment_prefix(lines)
    refs, num_refs = _collect_referenced_labels(lines, comment)
    out = []
    for line in lines:
        stripped = line.strip()
        # Skip empty lines between blocks (we'll add our own)
        if not stripped:
            if out and out[-1] != "\n":
                out.append("\n")
            continue
        # Skip directives (lines starting with .), but keep data directives.
        if stripped.startswith(".") and not _is_label_line(line, comment):
            directive = stripped.split()[0].rstrip(",")
            if directive not in _DATA_DIRECTIVES and directive not in _MODE_DIRECTIVES:
                continue
        if _is_label_line(line, comment):
            label = _defined_label_name(line, comment)
            if label is None:
                continue
            # Keep global/non-local labels.
            if not label.startswith("."):
                out.append(line)
                continue
            # Keep numeric local labels only if referenced (e.g. 1: with 1f/1b).
            if label.isdigit():
                if label in num_refs:
                    out.append(line)
                continue
            # Keep local assembler labels only when used by an instruction.
            if label in refs:
                out.append(line)
            continue
        # Skip inline comments that are the whole line
        if stripped.startswith(comment):
            continue
        # Strip trailing comments
        if comment in line:
            line = line[: line.index(comment)].rstrip() + "\n"
        out.append(line)

    # Remove trailing blank lines
    while out and out[-1] == "\n":
        out.pop()

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


def clean(path, dry_run=False):
    with open(path, encoding="utf-8") as f:
        original = f.read()
    cleaned = _clean_lines(original.splitlines(keepends=True))
    if dry_run:
        if cleaned != original:
            print(f"Would update {path}")
        else:
            print(f"No change {path}")
        return
    with open(path, "w", encoding="utf-8") as f:
        f.write(cleaned)


def clean_stream(dry_run=False):
    original = sys.stdin.read()
    cleaned = _clean_lines(original.splitlines(keepends=True))
    if dry_run:
        if cleaned != original:
            print("Would update stdin")
        else:
            print("No change stdin")
        return
    sys.stdout.write(cleaned)


def _iter_input_files(patterns):
    seen = set()
    for pattern in patterns:
        if pattern == "-":
            yield "-"
            continue
        matches = glob.glob(pattern)
        if not matches:
            matches = [pattern]
        for entry in matches:
            p = Path(entry)
            # Only clean files ending with .s; skip everything else.
            if p.suffix != ".s" or not p.is_file():
                print(f"Skipping {entry}: not a .s file", file=sys.stderr)
                continue
            resolved = str(p.resolve())
            if resolved in seen:
                continue
            seen.add(resolved)
            yield str(p)


if __name__ == "__main__":
    args = sys.argv[1:]
    dry_run = False
    if "--dry-run" in args:
        dry_run = True
        args = [a for a in args if a != "--dry-run"]
    if len(args) < 1:
        print(
            f"Usage: {sys.argv[0]} [--dry-run] <file.s|glob> [more files/globs...]",
            file=sys.stderr,
        )
        sys.exit(1)
    had_inputs = False
    for input_path in _iter_input_files(args):
        had_inputs = True
        if input_path == "-":
            clean_stream(dry_run=dry_run)
        else:
            clean(input_path, dry_run=dry_run)
    if not had_inputs:
        print(
            f"Usage: {sys.argv[0]} [--dry-run] <file.s|glob|- > [more files/globs/-...]",
            file=sys.stderr,
        )
        sys.exit(1)