ritz/tests/test_llvm_codegen.py

"""Tests for LLVM IR code generation."""

import subprocess
import tempfile
from pathlib import Path


def test_llvm_hello_world_structure():
    """Test that LLVM IR has correct structure for HelloWorld."""
    from ritz.llvm_codegen import generate_llvm_ir

    ir_dict = {
        "type": "component",
        "name": "HelloWorld",
        "args": [],
        "body": [
            {
                "type": "with",
                "element": {"tag": "P", "props": {}},
                "body": [{"type": "yield", "value": {"type": "constant", "value": "Hello World"}}],
            }
        ],
    }

    llvm_ir = generate_llvm_ir(ir_dict)

    # Check basic LLVM IR structure
    assert "define i32 @render_HelloWorld()" in llvm_ir
    assert "ret void" in llvm_ir

    # Should have string constants
    assert '"p\\00"' in llvm_ir or 'c"p\\00"' in llvm_ir
    assert "Hello World" in llvm_ir

    # Should call JS bridge functions
    assert "@js_create_element" in llvm_ir
    assert "@js_create_text" in llvm_ir
    assert "@js_append_child" in llvm_ir


def test_llvm_ir_has_function_declarations():
    """Test that LLVM IR declares imported JS functions."""
    from ritz.llvm_codegen import generate_llvm_ir

    ir_dict = {
        "type": "component",
        "name": "Simple",
        "args": [],
        "body": [
            {
                "type": "with",
                "element": {"tag": "P", "props": {}},
                "body": [{"type": "yield", "value": {"type": "constant", "value": "test"}}],
            }
        ],
    }

    llvm_ir = generate_llvm_ir(ir_dict)

    # Should declare imported functions
    assert "declare i32 @js_create_element" in llvm_ir
    assert "declare i32 @js_create_text" in llvm_ir
    assert "declare void @js_append_child" in llvm_ir
    assert "declare void @js_set_root" in llvm_ir


def test_llvm_ir_with_attributes():
    """Test LLVM IR generation with element attributes."""
    from ritz.llvm_codegen import generate_llvm_ir

    ir_dict = {
        "type": "component",
        "name": "Styled",
        "args": [],
        "body": [
            {
                "type": "with",
                "element": {
                    "tag": "Div",
                    "props": {
                        "class_": {"type": "constant", "value": "container"},
                        "id": {"type": "constant", "value": "main"},
                    },
                },
                "body": [],
            }
        ],
    }

    llvm_ir = generate_llvm_ir(ir_dict)

    # Should have js_set_attribute calls
    assert "@js_set_attribute" in llvm_ir
    assert "container" in llvm_ir
    assert "main" in llvm_ir


def test_llvm_ir_syntax_valid():
    """Test that generated LLVM IR has valid syntax (basic checks)."""
    from ritz.llvm_codegen import generate_llvm_ir

    ir_dict = {
        "type": "component",
        "name": "Test",
        "args": [],
        "body": [
            {
                "type": "with",
                "element": {"tag": "P", "props": {}},
                "body": [{"type": "yield", "value": {"type": "constant", "value": "x"}}],
            }
        ],
    }

    llvm_ir = generate_llvm_ir(ir_dict)

    # Basic syntax checks
    assert llvm_ir.count("define") == llvm_ir.count("ret")  # All functions have returns
    assert "%" in llvm_ir  # Has SSA variables
    assert "call" in llvm_ir  # Has function calls
    assert "@" in llvm_ir  # Has global symbols
    assert "define void @main()" in llvm_ir  # Has main entry point
    assert "define i32 @render_" in llvm_ir  # Component functions return i32


def test_llvm_ir_compiles_with_llc():
    """Test that generated LLVM IR actually compiles with llc."""
    from ritz.llvm_codegen import generate_llvm_ir

    ir_dict = {
        "type": "component",
        "name": "CompileTest",
        "args": [],
        "body": [
            {
                "type": "with",
                "element": {"tag": "P", "props": {}},
                "body": [{"type": "yield", "value": {"type": "constant", "value": "test"}}],
            }
        ],
    }

    llvm_ir = generate_llvm_ir(ir_dict)

    # Check if llc is available
    try:
        subprocess.run(["llc", "--version"], check=True, capture_output=True)
    except (subprocess.CalledProcessError, FileNotFoundError):
        # Skip if llc not installed
        import pytest

        pytest.skip("llc not available")
        return

    # Try to compile the IR
    with tempfile.NamedTemporaryFile(mode="w", suffix=".ll", delete=False) as f:
        f.write(llvm_ir)
        ll_file = Path(f.name)

    try:
        # Compile to object file (not WASM yet, just validate syntax)
        result = subprocess.run(
            ["llc", "-filetype=obj", str(ll_file)],
            capture_output=True,
            text=True,
        )

        # Should compile without errors
        assert result.returncode == 0, f"llc failed:\n{result.stderr}"

    finally:
        ll_file.unlink()


def test_llvm_ir_compiles_to_wasm():
    """Test that LLVM IR can be compiled to WASM."""
    from ritz.llvm_codegen import generate_llvm_ir

    ir_dict = {
        "type": "component",
        "name": "WasmTest",
        "args": [],
        "body": [
            {
                "type": "with",
                "element": {"tag": "P", "props": {}},
                "body": [{"type": "yield", "value": {"type": "constant", "value": "wasm"}}],
            }
        ],
    }

    llvm_ir = generate_llvm_ir(ir_dict)

    # Check if clang with wasm target is available
    try:
        result = subprocess.run(["clang", "--version"], check=True, capture_output=True, text=True)
        # Check if wasm target is supported
        if "wasm" not in result.stdout.lower():
            import pytest

            pytest.skip("clang wasm target not available")
            return
    except (subprocess.CalledProcessError, FileNotFoundError):
        import pytest

        pytest.skip("clang not available")
        return

    with tempfile.NamedTemporaryFile(mode="w", suffix=".ll", delete=False) as f:
        f.write(llvm_ir)
        ll_file = Path(f.name)

    wasm_file = ll_file.with_suffix(".wasm")

    try:
        # Compile to WASM
        result = subprocess.run(
            [
                "clang",
                "--target=wasm32",
                "-nostdlib",
                "-Wl,--no-entry",
                "-Wl,--export-all",
                "-o",
                str(wasm_file),
                str(ll_file),
            ],
            capture_output=True,
            text=True,
        )

        # Should compile successfully
        assert result.returncode == 0, f"clang failed:\n{result.stderr}"
        assert wasm_file.exists(), "WASM file not created"
        assert wasm_file.stat().st_size > 0, "WASM file is empty"

    finally:
        ll_file.unlink()
        if wasm_file.exists():
            wasm_file.unlink()