or1-design / dfgraph / pipeline.py

"""Progressive pipeline runner for dfasm assembly.

Runs assembler passes individually, capturing the deepest successful IRGraph
even when later passes fail. This enables partial graph visualisation.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Optional

from lark import Lark
from lark.exceptions import UnexpectedInput
from pathlib import Path

from asm.ir import IRGraph
from asm.lower import lower
from asm.resolve import resolve
from asm.place import place
from asm.allocate import allocate
from asm.errors import AssemblyError


_GRAMMAR_PATH = Path(__file__).parent.parent / "dfasm.lark"
_parser: Optional[Lark] = None


def _get_parser() -> Lark:
    """Lazily initialize and cache the Lark parser."""
    global _parser
    if _parser is None:
        _parser = Lark(
            _GRAMMAR_PATH.read_text(),
            parser="earley",
            propagate_positions=True,
        )
    return _parser


class PipelineStage(Enum):
    """Enumeration of pipeline stages."""
    PARSE_ERROR = "parse_error"
    LOWER = "lower"
    RESOLVE = "resolve"
    PLACE = "place"
    ALLOCATE = "allocate"


@dataclass(frozen=True)
class PipelineResult:
    """Result of running the progressive pipeline.

    Attributes:
        graph: The IRGraph at the deepest successful stage, or None if parse failed
        stage: The stage where progress stopped
        errors: All AssemblyErrors accumulated from the pipeline
        parse_error: String representation of parse error, if stage is PARSE_ERROR
    """
    graph: Optional[IRGraph]
    stage: PipelineStage
    errors: list[AssemblyError]
    parse_error: Optional[str] = None


def run_progressive(source: str) -> PipelineResult:
    """Run the assembly pipeline progressively, capturing errors at each stage.

    Unlike asm._run_pipeline() which raises on first error, this function runs
    each pass independently and accumulates errors, allowing partial graphs to
    be captured for visualization.

    Pipeline stages:
    1. Parse: Convert source to Lark CST (may raise lark.exceptions.UnexpectedInput)
    2. Lower: CST to IRGraph (may accumulate errors but doesn't stop)
    3. Resolve: Validate edge endpoints and scope (may accumulate errors)
    4. Place: Validate/auto-place nodes on PEs (stops before this if resolve failed)
    5. Allocate: Assign IRAM offsets and context slots (stops before this if place failed)

    Args:
        source: dfasm source code as a string

    Returns:
        PipelineResult containing the deepest graph, stage reached, and all errors
    """
    # Stage 1: Parse
    try:
        tree = _get_parser().parse(source)
    except UnexpectedInput as exc:
        return PipelineResult(
            graph=None,
            stage=PipelineStage.PARSE_ERROR,
            errors=[],
            parse_error=str(exc),
        )

    # Stage 2: Lower
    graph = lower(tree)
    stage = PipelineStage.LOWER

    # Stage 3: Resolve
    # Note: resolve() always runs after lower (matches asm._run_pipeline behaviour)
    graph = resolve(graph)
    stage = PipelineStage.RESOLVE

    # Stage 4: Place (skip if resolve accumulated errors)
    if not graph.errors:
        graph = place(graph)
        stage = PipelineStage.PLACE

    # Stage 5: Allocate (skip if place accumulated errors)
    if not graph.errors:
        graph = allocate(graph)
        stage = PipelineStage.ALLOCATE

    return PipelineResult(
        graph=graph,
        stage=stage,
        errors=list(graph.errors),
    )