resources.py

import io
import re
from functools import reduce
from typing import Dict, Iterable, List, Optional, Set, Tuple, TypeVar, Union

import matplotlib.pyplot as plt
import networkx as nx
from matplotlib.axes import Axes
from matplotlib.ticker import MaxNLocator

from b_asic._preferences import LATENCY_COLOR
from b_asic.process import MemoryVariable, PlainMemoryVariable, Process

# Default latency coloring RGB tuple
_LATENCY_COLOR = tuple(c / 255 for c in LATENCY_COLOR)

#
# Human-intuitive sorting:
# https://stackoverflow.com/questions/2669059/how-to-sort-alpha-numeric-set-in-python
#
# Typing '_T' to help Pyright propagate type-information
#
_T = TypeVar('_T')


def _sorted_nicely(to_be_sorted: Iterable[_T]) -> List[_T]:
    """Sort the given iterable in the way that humans expect."""

    def convert(text):
        return int(text) if text.isdigit() else text

    def alphanum_key(key):
        return [convert(c) for c in re.split('([0-9]+)', str(key))]

    return sorted(to_be_sorted, key=alphanum_key)


def _sanitize_port_option(
    read_ports: Optional[int] = None,
    write_ports: Optional[int] = None,
    total_ports: Optional[int] = None,
) -> Tuple[int, int, int]:
    """
    General port sanitization function used to test if a port specification makes sense.
    Raises ValueError if the port specification is in-proper.

    Parameters
    ----------
    read_ports : int, optional
        The number of read ports.
    write_ports : int, optional
        The number of write ports.
    total_ports : int, optional
        The total number of ports

    Returns
    -------
    Returns a triple int tuple (read_ports, write_ports, total_ports) equal to the input, or sanitized if one of the input equals None.
    If total_ports is set to None at the input, it is set to read_ports+write_ports at the output.
    If read_ports or write_ports is set to None at the input, it is set to total_ports at the output.

    """
    if total_ports is None:
        if read_ports is None or write_ports is None:
            raise ValueError(
                "If total_ports is unset, both read_ports and write_ports"
                " must be provided."
            )
        else:
            total_ports = read_ports + write_ports
    else:
        read_ports = total_ports if read_ports is None else read_ports
        write_ports = total_ports if write_ports is None else write_ports
    if total_ports < read_ports:
        raise ValueError(
            f'Total ports ({total_ports}) less then read ports ({read_ports})'
        )
    if total_ports < write_ports:
        raise ValueError(
            f'Total ports ({total_ports}) less then write ports ({write_ports})'
        )
    return (read_ports, write_ports, total_ports)


def draw_exclusion_graph_coloring(
    exclusion_graph: nx.Graph,
    color_dict: Dict[Process, int],
    ax: Optional[Axes] = None,
    color_list: Optional[Union[List[str], List[Tuple[float, float, float]]]] = None,
) -> None:
    """
    Draw a colored exclusion graph from the memory assignment.

    .. code-block:: python

        _, ax = plt.subplots(1, 1)
        collection = ProcessCollection(...)
        exclusion_graph = collection.create_exclusion_graph_from_overlap()
        color_dict = nx.greedy_color(exclusion_graph)
        draw_exclusion_graph_coloring(exclusion_graph, color_dict, ax=ax[0])
        plt.show()

    Parameters
    ----------
    exclusion_graph : nx.Graph
        A nx.Graph exclusion graph object that is to be drawn.
    color_dict : dict
        A dict where keys are :class:`~b_asic.process.Process` objects and values are
        integers representing colors. These dictionaries are automatically generated by
        :func:`networkx.algorithms.coloring.greedy_color`.
    ax : :class:`matplotlib.axes.Axes`, optional
        A Matplotlib :class:`~matplotlib.axes.Axes` object to draw the exclusion graph.
    color_list : iterable of color, optional
        A list of colors in Matplotlib format.

    Returns
    -------
    None
    """
    COLOR_LIST = [
        '#aa0000',
        '#00aa00',
        '#0000ff',
        '#ff00aa',
        '#ffaa00',
        '#ffffff',
        '#00ffaa',
        '#aaff00',
        '#aa00ff',
        '#00aaff',
        '#ff0000',
        '#00ff00',
        '#0000aa',
        '#aaaa00',
        '#aa00aa',
        '#00aaaa',
        '#666666',
    ]
    if color_list is None:
        node_color_dict = {k: COLOR_LIST[v] for k, v in color_dict.items()}
    else:
        node_color_dict = {k: color_list[v] for k, v in color_dict.items()}
    node_color_list = [node_color_dict[node] for node in exclusion_graph]
    nx.draw_networkx(
        exclusion_graph,
        node_color=node_color_list,
        ax=ax,
        pos=nx.spring_layout(exclusion_graph, seed=1),
    )


class _ForwardBackwardEntry:
    def __init__(
        self,
        inputs: Optional[List[Process]] = None,
        outputs: Optional[List[Process]] = None,
        regs: Optional[List[Optional[Process]]] = None,
        back_edge_to: Optional[Dict[int, int]] = None,
        back_edge_from: Optional[Dict[int, int]] = None,
        outputs_from: Optional[int] = None,
    ):
        """
        Single entry in a _ForwardBackwardTable. Aggregate type of input, output and list of registers.

        Parameters
        ----------
        inputs : List[Process], optional
            input
        outputs : List[Process], optional
            output
        regs : List[Optional[Process]], optional
            regs
        back_edge_to : dict, optional
            Dictionary containing back edges of this entry to registers in the next entry.
        back_edge_from : dict, optional
            Dictionary containing the back edge of the previous entry to registers in this entry.
        outputs_from : int, optional
        """
        self.inputs: List[Process] = [] if inputs is None else inputs
        self.outputs: List[Process] = [] if outputs is None else outputs
        self.regs: List[Optional[Process]] = [] if regs is None else regs
        self.back_edge_to: Dict[int, int] = {} if back_edge_to is None else back_edge_to
        self.back_edge_from: Dict[int, int] = (
            {} if back_edge_from is None else back_edge_from
        )
        self.outputs_from = outputs_from


class _ForwardBackwardTable:
    def __init__(self, collection: 'ProcessCollection'):
        """
        Forward-Backward allocation table for ProcessCollections. This structure implements the forward-backward
        register allocation algorithm, which is used to generate hardware from MemoryVariables in a ProcessCollection.

        Parameters
        ----------
        collection : ProcessCollection
            ProcessCollection to apply forward-backward allocation on
        """
        # Generate an alive variable list