signal_flow_graph.py

"""
B-ASIC Signal Flow Graph Module.

Contains the signal flow graph operation.
"""

import itertools as it
import re
from collections import defaultdict, deque
from io import StringIO
from numbers import Number
from queue import PriorityQueue
from typing import (
    DefaultDict,
    Deque,
    Dict,
    Iterable,
    List,
    MutableSet,
    Optional,
    Sequence,
    Set,
    Tuple,
    Union,
    cast,
)

from graphviz import Digraph

from b_asic.graph_component import (
    GraphComponent,
    GraphID,
    GraphIDNumber,
    Name,
    TypeName,
)
from b_asic.operation import (
    AbstractOperation,
    MutableDelayMap,
    MutableResultMap,
    Operation,
    ResultKey,
)
from b_asic.port import InputPort, OutputPort, SignalSourceProvider
from b_asic.signal import Signal
from b_asic.special_operations import Delay, Input, Output

DelayQueue = List[Tuple[str, ResultKey, OutputPort]]


class GraphIDGenerator:
    """Generates Graph IDs for objects."""

    _next_id_number: DefaultDict[TypeName, GraphIDNumber]

    def __init__(self, id_number_offset: GraphIDNumber = GraphIDNumber(0)):
        """Construct a GraphIDGenerator."""
        self._next_id_number = defaultdict(lambda: id_number_offset)

    def next_id(
        self, type_name: TypeName, used_ids: MutableSet = set()
    ) -> GraphID:
        """Get the next graph id for a certain graph id type."""
        self._next_id_number[type_name] += 1
        new_id = type_name + str(self._next_id_number[type_name])
        while new_id in used_ids:
            self._next_id_number[type_name] += 1
            new_id = type_name + str(self._next_id_number[type_name])
        return GraphID(new_id)

    @property
    def id_number_offset(self) -> GraphIDNumber:
        """Get the graph id number offset of this generator."""
        return GraphIDNumber(
            self._next_id_number.default_factory()
        )  # pylint: disable=not-callable


class SFG(AbstractOperation):
    """
    Signal flow graph.

    Contains a set of connected operations, forming a new operation.
    Used as a base for simulation, scheduling, etc.
    """

    _components_by_id: Dict[GraphID, GraphComponent]
    _components_by_name: DefaultDict[Name, List[GraphComponent]]
    _components_dfs_order: List[GraphComponent]
    _operations_dfs_order: List[Operation]
    _operations_topological_order: List[Operation]
    _graph_id_generator: GraphIDGenerator
    _input_operations: List[Input]
    _output_operations: List[Output]
    _original_components_to_new: Dict[GraphComponent, GraphComponent]
    _original_input_signals_to_indices: Dict[Signal, int]
    _original_output_signals_to_indices: Dict[Signal, int]
    _precedence_list: Optional[List[List[OutputPort]]]
    _used_ids: Set[GraphID] = set()

    def __init__(
        self,
        inputs: Optional[Sequence[Input]] = None,
        outputs: Optional[Sequence[Output]] = None,
        input_signals: Optional[Sequence[Signal]] = None,
        output_signals: Optional[Sequence[Signal]] = None,
        id_number_offset: GraphIDNumber = GraphIDNumber(0),
        name: Name = Name(""),
        input_sources: Optional[
            Sequence[Optional[SignalSourceProvider]]
        ] = None,
    ):
        """
        Construct an SFG given its inputs and outputs.

        Inputs/outputs may be specified using either Input/Output operations
        directly with the inputs/outputs parameters, or using signals with the
        input_signals/output_signals parameters. If signals are used, the
        corresponding Input/Output operations will be created automatically.

        The id_number_offset parameter specifies what number graph IDs will be
        offset by for each new graph component type. IDs start at 1 by default,
        so the default offset of 0 will result in IDs like "c1", "c2", etc.
        while an offset of 3 will result in "c4", "c5", etc.
        """

        input_signal_count = 0 if input_signals is None else len(input_signals)
        input_operation_count = 0 if inputs is None else len(inputs)
        output_signal_count = (
            0 if output_signals is None else len(output_signals)
        )
        output_operation_count = 0 if outputs is None else len(outputs)
        super().__init__(
            input_count=input_signal_count + input_operation_count,
            output_count=output_signal_count + output_operation_count,
            name=name,
            input_sources=input_sources,
        )

        self._components_by_id = {}
        self._used_ids = set()
        self._components_by_name = defaultdict(list)
        self._components_dfs_order = []
        self._operations_dfs_order = []
        self._operations_topological_order = []
        self._graph_id_generator = GraphIDGenerator(
            GraphIDNumber(id_number_offset)
        )
        self._input_operations = []
        self._output_operations = []
        self._original_components_to_new = {}
        self._original_input_signals_to_indices = {}
        self._original_output_signals_to_indices = {}
        self._precedence_list = None

        # Setup input signals.
        if input_signals is not None:
            for input_index, signal in enumerate(input_signals):
                if signal in self._original_components_to_new:
                    raise ValueError(
                        f"Duplicate input signal {signal!r} in SFG"
                    )
                new_input_op = cast(
                    Input, self._add_component_unconnected_copy(Input())
                )
                new_signal = cast(
                    Signal, self._add_component_unconnected_copy(signal)
                )
                new_signal.set_source(new_input_op.output(0))
                self._input_operations.append(new_input_op)
                self._original_input_signals_to_indices[signal] = input_index

        # Setup input operations, starting from indices after input signals.
        if inputs is not None:
            for input_index, input_op in enumerate(inputs, input_signal_count):
                if input_op in self._original_components_to_new:
                    raise ValueError(
                        f"Duplicate input operation {input_op!r} in SFG"
                    )
                new_input_op = self._add_component_unconnected_copy(input_op)
                for signal in input_op.output(0).signals:
                    assert signal not in self._original_components_to_new, (
                        "Duplicate input signals connected to input ports"
                        " supplied to SFG constructor."
                    )
                    new_signal = self._add_component_unconnected_copy(signal)
                    new_signal.set_source(new_input_op.output(0))
                    self._original_input_signals_to_indices[
                        signal
                    ] = input_index

                self._input_operations.append(new_input_op)

        # Setup output signals.
        if output_signals is not None:
            for output_index, signal in enumerate(output_signals):
                new_output_op = cast(
                    Output, self._add_component_unconnected_copy(Output())
                )
                if signal in self._original_components_to_new: