vga_sim.py

#!/usr/bin/env python3
"""
VGA Simulation Script for Tiny Tapeout Projects

Simulates Verilog VGA projects using Verilator and captures frames as PNG images.
Reads project configuration from info.yaml and handles compilation, simulation,
and frame capture automatically.

Usage:
    python vga_sim.py [options] <project_dir>
    python vga_sim.py [options] -f <verilog_files...>

Examples:
    # Simulate project from info.yaml
    python vga_sim.py /path/to/tt-project

    # Capture 10 frames
    python vga_sim.py -n 10 /path/to/tt-project

    # Direct file mode (no info.yaml)
    python vga_sim.py -f project.v hvsync_generator.v -t tt_um_example
"""

import argparse
import os
import shutil
import struct
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Optional

# Check for yaml - provide helpful error if missing
try:
    import yaml
except ImportError:
    print("Error: PyYAML is required. Install with: pip install pyyaml", file=sys.stderr)
    sys.exit(1)


# VGA timing constants (640x480 @ 60Hz)
H_DISPLAY = 640
H_FRONT = 16
H_SYNC = 96
H_BACK = 48
H_TOTAL = H_DISPLAY + H_FRONT + H_SYNC + H_BACK  # 800

V_DISPLAY = 480
V_BOTTOM = 10
V_SYNC = 2
V_TOP = 33
V_TOTAL = V_DISPLAY + V_BOTTOM + V_SYNC + V_TOP  # 525


def generate_testbench(top_module: str) -> str:
    """Generate C++ testbench code for Verilator."""
    return f'''// Auto-generated Verilator VGA testbench
#include <verilated.h>
#include "V{top_module}.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cstdint>

#define H_DISPLAY   {H_DISPLAY}
#define H_FRONT     {H_FRONT}
#define H_SYNC      {H_SYNC}
#define H_BACK      {H_BACK}
#define H_TOTAL     {H_TOTAL}

#define V_DISPLAY   {V_DISPLAY}
#define V_BOTTOM    {V_BOTTOM}
#define V_SYNC      {V_SYNC}
#define V_TOP       {V_TOP}
#define V_TOTAL     {V_TOTAL}

static uint8_t framebuffer[H_DISPLAY * V_DISPLAY * 3];

inline uint8_t extend2to8(uint8_t val) {{ return val * 85; }}

void decode_vga(uint8_t uo_out, bool &hsync, bool &vsync, uint8_t &r, uint8_t &g, uint8_t &b) {{
    hsync = (uo_out >> 7) & 1;
    vsync = (uo_out >> 3) & 1;
    r = ((uo_out & 0x01) << 1) | ((uo_out >> 4) & 0x01);
    g = ((uo_out & 0x02) >> 0) | ((uo_out >> 5) & 0x01);
    b = ((uo_out & 0x04) >> 1) | ((uo_out >> 6) & 0x01);
}}

void save_ppm(const char *filename) {{
    FILE *f = fopen(filename, "wb");
    if (!f) {{ fprintf(stderr, "Error: Cannot open %s\\n", filename); return; }}
    fprintf(f, "P6\\n%d %d\\n255\\n", H_DISPLAY, V_DISPLAY);
    fwrite(framebuffer, 1, H_DISPLAY * V_DISPLAY * 3, f);
    fclose(f);
}}

void clock_tick(V{top_module} *top) {{
    top->clk = 0; top->eval();
    top->clk = 1; top->eval();
}}

int main(int argc, char **argv) {{
    int num_frames = 1;
    const char *output_prefix = "frame";
    if (argc > 1) num_frames = atoi(argv[1]);
    if (argc > 2) output_prefix = argv[2];

    Verilated::commandArgs(argc, argv);
    V{top_module} *top = new V{top_module};

    // Reset sequence
    top->clk = 0;
    top->rst_n = 0;
    top->ena = 1;
    top->ui_in = 0;
    top->uio_in = 0;

    for (int i = 0; i < 10; i++) clock_tick(top);
    top->rst_n = 1;

    bool hsync, vsync;
    uint8_t r, g, b;

    // Synchronize to frame start
    do {{ clock_tick(top); decode_vga(top->uo_out, hsync, vsync, r, g, b); }} while (!vsync);
    do {{ clock_tick(top); decode_vga(top->uo_out, hsync, vsync, r, g, b); }} while (vsync);

    // Capture frames
    for (int frame = 0; frame < num_frames; frame++) {{
        memset(framebuffer, 0, sizeof(framebuffer));

        // Skip V_TOP lines (vertical back porch)
        for (int line = 0; line < V_TOP; line++)
            for (int px = 0; px < H_TOTAL; px++) clock_tick(top);

        // Capture V_DISPLAY lines
        for (int y = 0; y < V_DISPLAY; y++) {{
            // Skip H_BACK pixels (horizontal back porch)
            for (int px = 0; px < H_BACK; px++) clock_tick(top);

            // Capture H_DISPLAY pixels
            for (int x = 0; x < H_DISPLAY; x++) {{
                clock_tick(top);
                decode_vga(top->uo_out, hsync, vsync, r, g, b);
                int idx = (y * H_DISPLAY + x) * 3;
                framebuffer[idx] = extend2to8(r);
                framebuffer[idx + 1] = extend2to8(g);
                framebuffer[idx + 2] = extend2to8(b);
            }}

            // Skip H_FRONT + H_SYNC pixels
            for (int px = 0; px < H_FRONT + H_SYNC; px++) clock_tick(top);
        }}

        // Skip V_BOTTOM + V_SYNC lines
        for (int line = 0; line < V_BOTTOM + V_SYNC; line++)
            for (int px = 0; px < H_TOTAL; px++) clock_tick(top);

        // Save frame
        char filename[256];
        snprintf(filename, sizeof(filename), "%s_%04d.ppm", output_prefix, frame);
        save_ppm(filename);
        fprintf(stderr, "Saved %s\\n", filename);
    }}

    top->final();
    delete top;
    return 0;
}}
'''


def ppm_to_png(ppm_path: Path, png_path: Path) -> bool:
    """Convert PPM to PNG using ImageMagick or Python."""
    # Try ImageMagick first (faster)
    if shutil.which("convert"):
        result = subprocess.run(
            ["convert", str(ppm_path), str(png_path)],
            capture_output=True
        )
        if result.returncode == 0:
            return True

    # Fallback to pure Python
    try:
        with open(ppm_path, "rb") as f:
            # Read PPM header
            magic = f.readline().strip()
            if magic != b"P6":
                return False

            # Skip comments
            line = f.readline()
            while line.startswith(b"#"):
                line = f.readline()

            width, height = map(int, line.split())
            maxval = int(f.readline().strip())

            # Read pixel data
            pixels = f.read()

        # Write PNG using pure Python (minimal PNG encoder)
        write_png(png_path, width, height, pixels)
        return True
    except Exception as e:
        print(f"Warning: PNG conversion failed: {e}", file=sys.stderr)
        return False


def write_png(path: Path, width: int, height: int, rgb_data: bytes):
    """Write a minimal PNG file from RGB data."""
    import zlib

    def png_chunk(chunk_type: bytes, data: bytes) -> bytes:
        chunk_len = struct.pack(">I", len(data))
        chunk_crc = struct.pack(">I", zlib.crc32(chunk_type + data) & 0xFFFFFFFF)
        return chunk_len + chunk_type + data + chunk_crc

    # PNG signature
    signature = b"\x89PNG\r\n\x1a\n"

    # IHDR chunk
    ihdr_data = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)
    ihdr = png_chunk(b"IHDR", ihdr_data)

    # IDAT chunk (image data)
    raw_data = b""
    for y in range(height):
        raw_data += b"\x00"  # Filter type: None
        raw_data += rgb_data[y * width * 3:(y + 1) * width * 3]

    compressed = zlib.compress(raw_data, 9)
    idat = png_chunk(b"IDAT", compressed)

    # IEND chunk
    iend = png_chunk(b"IEND", b"")

    with open(path, "wb") as f:
        f.write(signature + ihdr + idat + iend)


def parse_info_yaml(yaml_path: Path) -> tuple[str, list[str], Path]:
    """Parse info.yaml and return (top_module, source_files, src_dir)."""
    with open(yaml_path) as f:
        info = yaml.safe_load(f)

    project = info.get("project", {})
    top_module = project.get("top_module", "tt_um_example")
    source_files = project.get("source_files", [])

    # Source files are relative to src/ directory
    src_dir = yaml_path.parent / "src"

    return top_module, source_files, src_dir


def find_hvsync_generator(search_paths: list[Path]) -> Optional[Path]:
    """Find hvsync_generator.v in search paths."""
    for path in search_paths:
        if path.is_dir():
            hvsync = path / "hvsync_generator.v"
            if hvsync.exists():
                return hvsync
        elif path.name == "hvsync_generator.v" and path.exists():
            return path

    # Check in script directory
    script_dir = Path(__file__).parent
    hvsync = script_dir / "hvsync_generator.v"
    if hvsync.exists():
        return hvsync

    return None


def run_verilator(
    source_files: list[Path],
    top_module: str,
    build_dir: Path,
    num_frames: int,
    output_prefix: str,
    output_dir: Path,
    verbose: bool = False
) -> tuple[bool, list[Path]]:
    """
    Compile and run Verilator simulation.
    Returns (success, list_of_output_files).
    """
    build_dir.mkdir(parents=True, exist_ok=True)

    # Generate testbench
    tb_path = build_dir / "vga_tb.cpp"
    tb_path.write_text(generate_testbench(top_module))

    # Verilator command
    verilator_args = [
        "verilator",
        "--cc", "--exe", "--build",
        "-O3",
        "--x-assign", "fast",
        "-Wno-UNUSEDSIGNAL",
        "-Wno-DECLFILENAME",
        "-Wno-WIDTH",
        "-Wno-CASEINCOMPLETE",
        "-Wno-TIMESCALEMOD",
        "-Wno-PINMISSING",
        "-Wno-fatal",
        "--top-module", top_module,
        "-Mdir", str(build_dir),
        "-o", "vga_sim",
    ]

    # Add source files
    for src in source_files:
        verilator_args.append(str(src))

    # Add testbench
    verilator_args.append(str(tb_path))

    # Add include paths for all source directories
    include_dirs = set()
    for src in source_files:
        include_dirs.add(src.parent)
    for inc_dir in include_dirs:
        verilator_args.append(f"-I{inc_dir}")

    if verbose:
        print(f"Running: {' '.join(verilator_args)}", file=sys.stderr)

    # Run Verilator
    result = subprocess.run(
        verilator_args,
        capture_output=True,
        text=True
    )

    if verbose or result.returncode != 0:
        # Filter out perl locale warnings
        stderr_lines = [
            line for line in result.stderr.split("\n")
            if not line.startswith("perl: warning")
            and "LANGUAGE" not in line
            and "LC_ALL" not in line
            and "locale" not in line.lower()
            and line.strip()
        ]
        if stderr_lines:
            print("\n".join(stderr_lines), file=sys.stderr)

    sim_executable = build_dir / "vga_sim"
    if not sim_executable.exists():
        print("Error: Verilator compilation failed", file=sys.stderr)
        if result.stdout:
            print(result.stdout, file=sys.stderr)
        return False, []

    # Run simulation
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / output_prefix

    sim_args = [str(sim_executable), str(num_frames), str(output_path)]

    if verbose:
        print(f"Running: {' '.join(sim_args)}", file=sys.stderr)

    result = subprocess.run(
        sim_args,
        capture_output=True,
        text=True,
        cwd=output_dir
    )

    # Print simulation output (includes $display statements)
    if result.stdout:
        print(result.stdout)
    if result.stderr:
        # stderr has "Saved frame_XXXX.ppm" messages
        for line in result.stderr.split("\n"):
            if line.strip():
                print(line, file=sys.stderr)

    # Collect output files
    output_files = []
    for i in range(num_frames):
        ppm_file = output_dir / f"{output_prefix}_{i:04d}.ppm"
        if ppm_file.exists():
            output_files.append(ppm_file)

    return len(output_files) > 0, output_files


def main():
    parser = argparse.ArgumentParser(
        description="VGA Simulation for Tiny Tapeout Projects",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s /path/to/tt-project              # Simulate from info.yaml
  %(prog)s -n 10 /path/to/tt-project        # Capture 10 frames
  %(prog)s -f project.v hvsync_generator.v  # Direct file mode
  %(prog)s --keep-ppm /path/to/project      # Keep PPM files (don't convert to PNG)
"""
    )

    parser.add_argument(
        "project_dir",
        nargs="?",
        type=Path,
        help="Project directory containing info.yaml"
    )

    parser.add_argument(
        "-f", "--files",
        nargs="+",
        type=Path,
        metavar="FILE",
        help="Verilog source files (direct mode, no info.yaml)"
    )

    parser.add_argument(
        "-t", "--top-module",
        default="tt_um_vga_example",
        help="Top module name (default: tt_um_vga_example)"
    )

    parser.add_argument(
        "-n", "--num-frames",
        type=int,
        default=1,
        help="Number of frames to capture (default: 1)"
    )

    parser.add_argument(
        "-o", "--output",
        default="frame",
        help="Output file prefix (default: frame)"
    )

    parser.add_argument(
        "-d", "--output-dir",
        type=Path,
        default=Path("."),
        help="Output directory (default: current directory)"
    )

    parser.add_argument(
        "--keep-ppm",
        action="store_true",
        help="Keep PPM files instead of converting to PNG"
    )

    parser.add_argument(
        "--build-dir",
        type=Path,
        help="Build directory (default: temp directory)"
    )

    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Verbose output"
    )

    parser.add_argument(
        "-c", "--clean",
        action="store_true",
        help="Clean build directory before building"
    )

    args = parser.parse_args()

    # Determine source files and top module
    source_files: list[Path] = []
    top_module = args.top_module

    if args.files:
        # Direct file mode
        source_files = [f.resolve() for f in args.files]
        top_module = args.top_module
    elif args.project_dir:
        # Project directory mode
        project_dir = args.project_dir.resolve()
        info_yaml = project_dir / "info.yaml"

        if not info_yaml.exists():
            print(f"Error: info.yaml not found in {project_dir}", file=sys.stderr)
            sys.exit(1)

        top_module, src_names, src_dir = parse_info_yaml(info_yaml)

        if not src_dir.exists():
            print(f"Error: src directory not found: {src_dir}", file=sys.stderr)
            sys.exit(1)

        for name in src_names:
            src_path = src_dir / name
            if not src_path.exists():
                print(f"Error: Source file not found: {src_path}", file=sys.stderr)
                sys.exit(1)
            source_files.append(src_path)

        # Also look for hvsync_generator.v in src directory
        hvsync = src_dir / "hvsync_generator.v"
        if hvsync.exists() and hvsync not in source_files:
            source_files.append(hvsync)
    else:
        parser.print_help()
        sys.exit(1)

    # Check if hvsync_generator.v is included, add from script dir if not
    has_hvsync = any("hvsync_generator" in f.name for f in source_files)
    if not has_hvsync:
        hvsync = find_hvsync_generator([f.parent for f in source_files])
        if hvsync:
            source_files.append(hvsync)
            if args.verbose:
                print(f"Added hvsync_generator.v from {hvsync}", file=sys.stderr)
        else:
            print("Warning: hvsync_generator.v not found", file=sys.stderr)

    # Verify all source files exist
    for src in source_files:
        if not src.exists():
            print(f"Error: Source file not found: {src}", file=sys.stderr)
            sys.exit(1)

    if args.verbose:
        print(f"Top module: {top_module}", file=sys.stderr)
        print(f"Source files:", file=sys.stderr)
        for src in source_files:
            print(f"  {src}", file=sys.stderr)

    # Set up build directory
    if args.build_dir:
        build_dir = args.build_dir.resolve()
    else:
        # Use a persistent temp directory based on top module name
        build_dir = Path(tempfile.gettempdir()) / f"vga_sim_{top_module}"

    if args.clean and build_dir.exists():
        shutil.rmtree(build_dir)
        if args.verbose:
            print(f"Cleaned build directory: {build_dir}", file=sys.stderr)

    # Run simulation
    output_dir = args.output_dir.resolve()
    success, ppm_files = run_verilator(
        source_files=source_files,
        top_module=top_module,
        build_dir=build_dir,
        num_frames=args.num_frames,
        output_prefix=args.output,
        output_dir=output_dir,
        verbose=args.verbose
    )

    if not success:
        sys.exit(1)

    # Convert PPM to PNG
    if not args.keep_ppm:
        for ppm_file in ppm_files:
            png_file = ppm_file.with_suffix(".png")
            if ppm_to_png(ppm_file, png_file):
                ppm_file.unlink()  # Remove PPM file
                print(f"Created {png_file}")
            else:
                print(f"Warning: Could not convert {ppm_file} to PNG")
    else:
        for ppm_file in ppm_files:
            print(f"Created {ppm_file}")

    print(f"Done! Generated {len(ppm_files)} frame(s)")


if __name__ == "__main__":
    main()