Move coreblocks utils to transactron (kuznia-rdzeni/coreblocks#503)

Co-authored-by: Lekcyjna <[email protected]>
Co-authored-by: Marek Materzok <[email protected]>

transactron/_utils.py

@@ -4,8 +4,8 @@
 from typing import Any, Concatenate, Optional, TypeAlias, TypeGuard, TypeVar
 from collections.abc import Callable, Iterable, Mapping
 from amaranth import *
-from coreblocks.utils._typing import LayoutLike, ShapeLike
-from coreblocks.utils import OneHotSwitchDynamic
+from transactron.utils._typing import LayoutLike, ShapeLike
+from transactron.utils import OneHotSwitchDynamic
 
 __all__ = [
     "Scheduler",

transactron/core.py

@@ -21,11 +21,10 @@
 from itertools import count, chain, filterfalse, product
 from amaranth.hdl.dsl import FSM, _ModuleBuilderDomain
 
-from coreblocks.utils import AssignType, assign, ModuleConnector
-from coreblocks.utils.utils import OneHotSwitchDynamic
+from transactron.utils import AssignType, assign, ModuleConnector, silence_mustuse
+from transactron.utils.utils import OneHotSwitchDynamic
 from ._utils import *
-from coreblocks.utils import silence_mustuse
-from coreblocks.utils._typing import ValueLike, SignalBundle, HasElaborate, SwitchKey, ModuleLike
+from transactron.utils._typing import ValueLike, SignalBundle, HasElaborate, SwitchKey, ModuleLike
 from .graph import Owned, OwnershipGraph, Direction
 
 __all__ = [

transactron/lib/reqres.py

@@ -1,7 +1,7 @@
 from amaranth import *
 from ..core import *
 from .connectors import Forwarder, FIFO
-from coreblocks.utils.fifo import BasicFifo
+from transactron.utils.fifo import BasicFifo
 from amaranth.utils import *
 
 __all__ = [

transactron/lib/simultaneous.py

@@ -3,7 +3,7 @@
 from ..core import TransactionBase
 from contextlib import contextmanager
 from typing import Optional
-from coreblocks.utils import ValueLike
+from transactron.utils import ValueLike
 
 __all__ = [
     "condition",

transactron/lib/storage.py

@@ -2,7 +2,7 @@
 from amaranth.utils import *
 from ..core import *
 from typing import Optional
-from coreblocks.utils import assign, AssignType
+from transactron.utils import assign, AssignType
 from .reqres import ArgumentsToResultsZipper
 
 __all__ = ["MemoryBank"]

transactron/lib/transformers.py

@@ -3,7 +3,7 @@
 from ..core import RecordDict
 from typing import Optional
 from collections.abc import Callable
-from coreblocks.utils import ValueLike, assign, AssignType
+from transactron.utils import ValueLike, assign, AssignType
 from .connectors import Forwarder, ManyToOneConnectTrans, ConnectTrans
 
 __all__ = [

transactron/tracing.py

@@ -7,7 +7,7 @@
 from amaranth.hdl.ir import Elaboratable, Fragment, Instance
 from amaranth.hdl.xfrm import FragmentTransformer
 from amaranth.hdl import dsl, ir, mem, xfrm
-from coreblocks.utils import HasElaborate
+from transactron.utils import HasElaborate
 from . import core
 
 

transactron/utils/__init__.py

@@ -0,0 +1,3 @@
+from .utils import *  # noqa: F401
+from ._typing import *  # noqa: F401
+from .debug_signals import *  # noqa: F401

transactron/utils/_typing.py

@@ -0,0 +1,105 @@
+from typing import (
+    Generic,
+    NoReturn,
+    Optional,
+    Protocol,
+    TypeAlias,
+    TypeVar,
+    runtime_checkable,
+    Union,
+    Any,
+)
+from collections.abc import Iterable, Mapping, Sequence
+from contextlib import AbstractContextManager
+from enum import Enum
+from amaranth import *
+from amaranth.lib.data import View
+from amaranth.hdl.ast import ShapeCastable, Statement, ValueCastable
+from amaranth.hdl.dsl import _ModuleBuilderSubmodules, _ModuleBuilderDomainSet, _ModuleBuilderDomain, FSM
+from amaranth.hdl.rec import Direction, Layout
+
+# Types representing Amaranth concepts
+FragmentLike: TypeAlias = Fragment | Elaboratable
+ValueLike: TypeAlias = Value | int | Enum | ValueCastable
+ShapeLike: TypeAlias = Shape | ShapeCastable | int | range | type[Enum]
+StatementLike: TypeAlias = Statement | Iterable["StatementLike"]
+LayoutLike: TypeAlias = (
+    Layout | Sequence[tuple[str, "ShapeLike | LayoutLike"] | tuple[str, "ShapeLike | LayoutLike", Direction]]
+)
+SwitchKey: TypeAlias = str | int | Enum
+
+# Internal Coreblocks types
+SignalBundle: TypeAlias = Signal | Record | View | Iterable["SignalBundle"] | Mapping[str, "SignalBundle"]
+LayoutList: TypeAlias = list[tuple[str, "ShapeLike | LayoutList"]]
+
+RecordIntDict: TypeAlias = Mapping[str, Union[int, "RecordIntDict"]]
+RecordIntDictRet: TypeAlias = Mapping[str, Any]  # full typing hard to work with
+RecordValueDict: TypeAlias = Mapping[str, Union[ValueLike, "RecordValueDict"]]
+
+
+class _ModuleBuilderDomainsLike(Protocol):
+    def __getattr__(self, name: str) -> _ModuleBuilderDomain:
+        ...
+
+    def __getitem__(self, name: str) -> _ModuleBuilderDomain:
+        ...
+
+    def __setattr__(self, name: str, value: _ModuleBuilderDomain) -> None:
+        ...
+
+    def __setitem__(self, name: str, value: _ModuleBuilderDomain) -> None:
+        ...
+
+
+_T_ModuleBuilderDomains = TypeVar("_T_ModuleBuilderDomains", bound=_ModuleBuilderDomainsLike)
+
+
+class ModuleLike(Protocol, Generic[_T_ModuleBuilderDomains]):
+    submodules: _ModuleBuilderSubmodules
+    domains: _ModuleBuilderDomainSet
+    d: _T_ModuleBuilderDomains
+
+    def If(self, cond: ValueLike) -> AbstractContextManager[None]:  # noqa: N802
+        ...
+
+    def Elif(self, cond: ValueLike) -> AbstractContextManager[None]:  # noqa: N802
+        ...
+
+    def Else(self) -> AbstractContextManager[None]:  # noqa: N802
+        ...
+
+    def Switch(self, test: ValueLike) -> AbstractContextManager[None]:  # noqa: N802
+        ...
+
+    def Case(self, *patterns: SwitchKey) -> AbstractContextManager[None]:  # noqa: N802
+        ...
+
+    def Default(self) -> AbstractContextManager[None]:  # noqa: N802
+        ...
+
+    def FSM(  # noqa: N802
+        self, reset: Optional[str] = ..., domain: str = ..., name: str = ...
+    ) -> AbstractContextManager[FSM]:
+        ...
+
+    def State(self, name: str) -> AbstractContextManager[None]:  # noqa: N802
+        ...
+
+    @property
+    def next(self) -> NoReturn:
+        ...
+
+    @next.setter
+    def next(self, name: str) -> None:
+        ...
+
+
+class HasElaborate(Protocol):
+    def elaborate(self, platform) -> "HasElaborate":
+        ...
+
+
+@runtime_checkable
+class HasDebugSignals(Protocol):
+    def debug_signals(self) -> SignalBundle:
+        ...

transactron/utils/debug_signals.py

@@ -0,0 +1,77 @@
+from typing import Optional
+from amaranth import *
+from ._typing import SignalBundle, HasDebugSignals
+from collections.abc import Collection, Mapping
+
+
+def auto_debug_signals(thing) -> SignalBundle:
+    """Automatic debug signal generation.
+
+    Exposes class attributes with debug signals (Amaranth `Signal`\\s,
+    `Record`\\s, `Array`\\s and `Elaboratable`\\s, `Method`\\s, classes
+    which define `debug_signals`). Used for generating ``gtkw`` files in
+    tests, for use in ``gtkwave``.
+    """
+
+    def auto_debug_signals_internal(thing, *, _visited: set) -> Optional[SignalBundle]:
+        # Please note, that the set `_visited` is used to memorise visited elements
+        # to break reference cycles. There is only one instance of this set, for whole
+        # `auto_debug_signals` recursion stack. It is being mutated by adding to it more
+        # elements id, so that caller know what was visited by callee.
+        smap: dict[str, SignalBundle] = {}
+
+        # Check for reference cycles e.g. Amaranth's MustUse
+        if id(thing) in _visited:
+            return None
+        _visited.add(id(thing))
+
+        match thing:
+            case HasDebugSignals():
+                return thing.debug_signals()
+            # avoid infinite recursion (strings are `Collection`s of strings)
+            case str():
+                return None
+            case Collection() | Mapping():
+                match thing:
+                    case Collection():
+                        f_iter = enumerate(thing)
+                    case Mapping():
+                        f_iter = thing.items()
+                for i, e in f_iter:
+                    sublist = auto_debug_signals_internal(e, _visited=_visited)
+                    if sublist is not None:
+                        smap[f"[{i}]"] = sublist
+                if smap:
+                    return smap
+                return None
+            case Array():
+                for i, e in enumerate(thing):
+                    if isinstance(e, Record):
+                        e.name = f"[{i}]"
+                return thing
+            case Signal() | Record():
+                return thing
+            case _:
+                try:
+                    vs = vars(thing)
+                except (KeyError, AttributeError, TypeError):
+                    return None
+
+                for v in vs:
+                    a = getattr(thing, v)
+
+                    # ignore private fields (mostly to ignore _MustUse_context to get pretty print)
+                    if v[0] == "_":
+                        continue
+
+                    dsignals = auto_debug_signals_internal(a, _visited=_visited)
+                    if dsignals is not None:
+                        smap[v] = dsignals
+                if smap:
+                    return smap
+                return None
+
+    ret = auto_debug_signals_internal(thing, _visited=set())
+    if ret is None:
+        return []
+    return ret