or1-design / asm / serialize.py

"""Serialize IRGraph back to dfasm source text.

This module provides the serialize() function that converts an IRGraph (at any
pipeline stage) back to valid dfasm source, enabling inspection of IR after
lowering, resolution, placement, or allocation.

The serializer emits:
- inst_def lines for nodes: {qualified_ref}|pe{N} <| {mnemonic}[, {const}]
- plain_edge lines for edges: {source} |> {dest}:{port}
- data_def lines: {name}|sm{id}:{cell} = {value}
- FUNCTION regions: $name |> { ...body... }
- LOCATION regions: bare directive tag followed by body

Names inside FUNCTION regions are unqualified (prefix stripped).
Anonymous nodes (__anon_*) are always emitted as inst_def + plain_edge,
never as inline syntax.
"""

from asm.ir import (
    IRGraph, IRNode, IREdge, IRRegion, RegionKind, IRDataDef
)
from asm.opcodes import OP_TO_MNEMONIC
from cm_inst import Port


def serialize(graph: IRGraph) -> str:
    """Serialize an IRGraph to dfasm source text.

    Converts an IRGraph back to valid dfasm source at any pipeline stage.
    Useful for inspecting IR after lowering, resolution, placement, or allocation.

    Args:
        graph: The IRGraph to serialize

    Returns:
        A string containing valid dfasm source text
    """
    if not graph.nodes and not graph.regions and not graph.data_defs and not graph.edges:
        return ""

    lines = []

    # Collect all nodes that are inside regions for later exclusion from top-level output
    nodes_in_regions: set[str] = set()
    edges_in_regions: set[tuple[str, str, Port]] = set()  # Track edges by (source, dest, port) tuple
    data_defs_in_regions: set[str] = set()

    # First pass: collect what's inside regions
    for region in graph.regions:
        for node_name in region.body.nodes.keys():
            nodes_in_regions.add(node_name)
        for edge in region.body.edges:
            edges_in_regions.add((edge.source, edge.dest, edge.port))
        for data_def in region.body.data_defs:
            data_defs_in_regions.add(data_def.name)

    # Emit regions in order
    for region in graph.regions:
        lines.append(_serialize_region(region, graph))

    # Emit top-level nodes (not inside any region)
    for name, node in graph.nodes.items():
        if name not in nodes_in_regions:
            lines.append(_serialize_node(name, node, func_scope=None))

    # Emit top-level edges (not inside any region)
    for edge in graph.edges:
        edge_key = (edge.source, edge.dest, edge.port)
        if edge_key not in edges_in_regions:
            lines.append(_serialize_edge(edge))

    # Emit top-level data_defs (not inside any region)
    for data_def in graph.data_defs:
        if data_def.name not in data_defs_in_regions:
            lines.append(_serialize_data_def(data_def))

    # Filter out empty lines and join
    output = '\n'.join(line for line in lines if line.strip())
    return output + '\n' if output else ""


def _serialize_region(region: IRRegion, parent_graph: IRGraph) -> str:
    """Serialize a single region (FUNCTION or LOCATION).

    Args:
        region: The IRRegion to serialize
        parent_graph: Parent IRGraph (for context)

    Returns:
        String containing the serialized region
    """
    lines = []

    if region.kind == RegionKind.FUNCTION:
        # FUNCTION regions: $name |> { ...body... }
        lines.append(f"{region.tag} |> {{")

        # Serialize body with function scope for name unqualification
        func_scope = region.tag
        for name, node in region.body.nodes.items():
            lines.append(_serialize_node(name, node, func_scope=func_scope))

        # Edges inside function
        for edge in region.body.edges:
            lines.append(_serialize_edge(edge, func_scope=func_scope))

        # Data defs inside function
        for data_def in region.body.data_defs:
            lines.append(_serialize_data_def(data_def))

        lines.append("}")

    elif region.kind == RegionKind.LOCATION:
        # LOCATION regions: bare directive tag with trailing colon, then body
        lines.append(f"{region.tag}:")

        # Serialize body (no function scope for locations)
        for name, node in region.body.nodes.items():
            lines.append(_serialize_node(name, node, func_scope=None))

        for edge in region.body.edges:
            lines.append(_serialize_edge(edge))

        for data_def in region.body.data_defs:
            lines.append(_serialize_data_def(data_def))

    return '\n'.join(lines)


def _serialize_node(name: str, node: IRNode, func_scope: str | None) -> str:
    """Serialize a single IR node as an inst_def line.

    Args:
        name: The node name
        node: The IRNode
        func_scope: If provided, unqualify the name by stripping this prefix

    Returns:
        String containing the inst_def line
    """
    # Unqualify the name if inside a function
    if func_scope and name.startswith(f"{func_scope}."):
        display_name = name[len(func_scope) + 1:]
    else:
        display_name = name

    # Get mnemonic
    try:
        mnemonic = OP_TO_MNEMONIC[node.opcode]
    except (KeyError, TypeError):
        # Fallback if mnemonic not found
        mnemonic = str(node.opcode).lower()

    # Build inst_def line: {ref}|pe{N} <| {mnemonic}[, {const}]
    pe_part = f"|pe{node.pe}" if node.pe is not None else ""
    line = f"{display_name}{pe_part} <| {mnemonic}"

    # Add const if present
    if node.const is not None:
        line += f", {node.const}"

    return line


def _serialize_edge(edge: IREdge, func_scope: str | None = None) -> str:
    """Serialize a single edge as a plain_edge line.

    Args:
        edge: The IREdge
        func_scope: If provided, unqualify names by stripping this prefix

    Returns:
        String containing the plain_edge line
    """
    # Unqualify names if inside a function scope
    source = edge.source
    dest = edge.dest

    if func_scope:
        if source.startswith(f"{func_scope}."):
            source = source[len(func_scope) + 1:]
        if dest.startswith(f"{func_scope}."):
            dest = dest[len(func_scope) + 1:]

    # Format port as :L or :R
    port_str = ":L" if edge.port == Port.L else ":R"

    return f"{source} |> {dest}{port_str}"


def _serialize_data_def(data_def: IRDataDef) -> str:
    """Serialize a single data definition.

    Args:
        data_def: The IRDataDef

    Returns:
        String containing the data_def line
    """
    sm_part = f"|sm{data_def.sm_id}" if data_def.sm_id is not None else ""
    cell_part = f":{data_def.cell_addr}" if data_def.cell_addr is not None else ""

    # Format value as hex if larger than a byte, decimal otherwise
    value_str = f"0x{data_def.value:x}" if data_def.value > 255 else str(data_def.value)

    return f"{data_def.name}{sm_part}{cell_part} = {value_str}"