Swift .tbd symbol demangler

demangle_tbd.py

import yaml
import subprocess
import sys
import re
import os
import concurrent.futures
from concurrent.futures import ProcessPoolExecutor
from tqdm import tqdm

def demangle_symbol(symbol):
    # Handle special prefixes like '$ld$previous$...'
    prefix_match = re.match(r'^(\$ld\$previous\$[^$]+\$\$[^$]+\$[^$]+\$_)(.+)$', symbol)
    if prefix_match:
        prefix, mangled = prefix_match.groups()
    else:
        prefix = ''
        mangled = symbol

    mangled = mangled.lstrip('_')

    # Use `swift demangle` to demangle the symbol
    try:
        result = subprocess.run(['swift', 'demangle', '--compact', mangled], capture_output=True, text=True)
        demangled = result.stdout.strip()
        if result.returncode != 0 or not demangled:
            # If demangling fails, return the original symbol
            return symbol
        demangled = demangled.replace(' (extension in SwiftUI):', '')
        demangled = demangled.replace(' in SwiftUI', '')
        # Break long lines for better readability
        demangled = re.sub(r' -> ', r' ->\n    ', demangled)
        demangled = re.sub(r', ', r',\n    ', demangled)
        return prefix + demangled
    except Exception as e:
        return symbol

def process_tbd_file(input_file, output_file):
    with open(input_file, 'r') as f:
        content = f.read()

    # Add a constructor for the '!tapi-tbd' tag lest we fail early
    def tbd_constructor(loader, node):
        return loader.construct_mapping(node)

    yaml.add_constructor('!tapi-tbd', tbd_constructor, Loader=yaml.FullLoader)

    data = yaml.load(content, Loader=yaml.FullLoader)

    symbols_to_demangle = []
    export_sections = []

    if 'exports' in data:
        for export in data['exports']:
            if 'symbols' in export:
                symbols_to_demangle.extend(export['symbols'])
                export_sections.append(export)

    if not symbols_to_demangle:
        print("No symbols to demangle found in the .tbd file.")
        return

    max_workers = os.cpu_count() or 1
    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        # Show a progress bar
        with tqdm(total=len(symbols_to_demangle), desc="Demangling symbols") as pbar:
            # Map symbols to demangled symbols
            future_to_symbol = {executor.submit(demangle_symbol, symbol): symbol for symbol in symbols_to_demangle}
            demangled_symbols = []
            for future in concurrent.futures.as_completed(future_to_symbol):
                demangled_symbol = future.result()
                demangled_symbols.append(demangled_symbol)
                pbar.update(1)

    # Replace the original symbols with demangled symbols
    symbol_iter = iter(demangled_symbols)
    for export in export_sections:
        if 'symbols' in export:
            num_symbols = len(export['symbols'])
            export['symbols'] = [next(symbol_iter) for _ in range(num_symbols)]

    # Write the modified data back to a file
    # Use custom Dumper to control formatting
    class CustomDumper(yaml.SafeDumper):
        pass

    def str_representer(dumper, data):
        if '\n' in data:
            return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
        return dumper.represent_scalar('tag:yaml.org,2002:str', data)

    CustomDumper.add_representer(str, str_representer)

    with open(output_file, 'w') as f:
        yaml.dump(data, f, Dumper=CustomDumper, default_flow_style=False, sort_keys=False)

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python demangle_tbd.py input.tbd output.tbd")
        sys.exit(1)
    input_file = sys.argv[1]
    output_file = sys.argv[2]
    if not os.path.isfile(input_file):
        print(f"Input file {input_file} does not exist.")
        sys.exit(1)
    process_tbd_file(input_file, output_file)
    print(f"Demangled TBD file written to {output_file}")

Just for fun, I compared this script on my 2021 10-core M1 Pro vs a base model 2024 Mac mini
M1 Pro: 47205/47205 [03:28<00:00, 226.10it/s]
M4: 47205/47205 [02:55<00:00, 268.79it/s]

Anyway, sample input and output here: https://gist.github.com/samhenrigold/460ba99cd0ae409cb9591b572711a57b