ritz/ritz/build.py

"""Build script for compiling Python components to WASM."""

import argparse
import subprocess
import sys
from pathlib import Path

from ritz.llvm_codegen_signals import generate_llvm_ir
from ritz.parser import parse_component


def build_component(component_file: Path, output_file: Path, debug: bool = False) -> None:
    """
    Build a Python component to WASM.

    Args:
        component_file: Path to Python component file
        output_file: Path for output WASM file
        debug: Whether to include debug symbols
    """
    print(f"[Ritz Build] Compiling {component_file} → {output_file}")

    # Step 1: Read component source
    print("  [1/4] Reading component...")
    source = component_file.read_text()

    # Step 2: Parse to IR
    print("  [2/4] Parsing Python to IR...")
    ir = parse_component(source)

    # Handle single component or list of components
    if isinstance(ir, list):
        for component in ir:
            print(f"    Component: {component['name']}")
    else:
        print(f"    Component: {ir['name']}")

    # Step 3: Generate LLVM IR
    print("  [3/4] Generating LLVM IR...")
    llvm_ir = generate_llvm_ir(ir)

    # Save LLVM IR for inspection (always Main.ll)
    ll_file = output_file.with_suffix(".ll")
    ll_file.write_text(llvm_ir)
    print(f"    Saved LLVM IR: {ll_file}")

    # Step 4: Compile to WASM with clang
    print("  [4/4] Compiling to WASM...")
    compile_to_wasm(ll_file, output_file, debug=debug)

    print(f"✅ Build complete: {output_file}")
    print(f"   Size: {output_file.stat().st_size} bytes")


def compile_to_wasm(ll_file: Path, output_file: Path, debug: bool = False) -> None:
    """
    Compile LLVM IR to WASM using clang.

    Args:
        ll_file: Path to LLVM IR source file
        output_file: Path for output WASM file
        debug: Whether to include debug symbols (-g)
    """
    # Check if clang is available
    try:
        subprocess.run(
            ["clang", "--version"],
            check=True,
            capture_output=True,
            text=True,
        )
    except (subprocess.CalledProcessError, FileNotFoundError):
        print("ERROR: clang not found. Install LLVM/clang")
        print("  Ubuntu/Debian: apt install clang lld")
        print("  macOS: brew install llvm")
        sys.exit(1)

    # Build clang command
    cmd = [
        "clang",
        "--target=wasm32",
        "-nostdlib",  # No standard library
        "-Wl,--no-entry",  # No _start entry point
        "-Wl,--export-all",  # Export all functions
        "-Wl,--allow-undefined",  # Allow undefined JS imports
        "-O2",  # Optimization
        "-o",
        str(output_file),
        str(ll_file),
    ]

    # Add debug symbols if requested
    if debug:
        cmd.extend(["-g", "-Wl,--export-dynamic"])
        print("    Debug symbols: enabled (DWARF in WASM)")

    # Run clang
    try:
        result = subprocess.run(
            cmd,
            check=True,
            capture_output=True,
            text=True,
        )
        if result.stderr:
            # Filter out common warnings
            stderr = result.stderr
            if "warning" in stderr.lower() and len(stderr) < 500:
                print(f"    Warnings: {stderr}")
    except subprocess.CalledProcessError as e:
        print("ERROR: Compilation failed")
        print(f"  Command: {' '.join(cmd)}")
        print(f"  stdout: {e.stdout}")
        print(f"  stderr: {e.stderr}")
        sys.exit(1)


def main() -> None:
    """CLI entry point."""
    parser = argparse.ArgumentParser(description="Build Ritz components to WASM")
    parser.add_argument("component", type=Path, help="Python component file")
    parser.add_argument(
        "-o", "--output", type=Path, help="Output WASM file (default: component.wasm)"
    )
    parser.add_argument("-g", "--debug", action="store_true", help="Include debug symbols")

    args = parser.parse_args()

    # Determine output file
    if args.output:
        output_file = args.output
    else:
        # Convention: always output Main.wasm in the same directory as the component
        output_file = args.component.parent / "Main.wasm"

    # Ensure output directory exists
    output_file.parent.mkdir(parents=True, exist_ok=True)

    # Build the component
    build_component(args.component, output_file, debug=args.debug)


if __name__ == "__main__":
    main()