From f242864779e4f874eff33cc3493b13a8ed8a6d28 Mon Sep 17 00:00:00 2001
From: Marek Materzok <tilk@tilk.eu>
Date: Sat, 6 Jan 2024 11:52:19 +0100
Subject: [PATCH] Move dependency related classes to Transactron (#544)

---
 coreblocks/core.py                            |  2 +-
 coreblocks/fu/exception.py                    |  2 +-
 coreblocks/fu/jumpbranch.py                   |  1 +
 coreblocks/fu/priv.py                         |  1 +
 coreblocks/lsu/dummyLsu.py                    |  2 +-
 coreblocks/params/__init__.py                 |  1 -
 coreblocks/params/genparams.py                | 42 +-----------------
 coreblocks/params/keys.py                     |  2 +-
 coreblocks/stages/func_blocks_unifier.py      |  7 ++-
 coreblocks/stages/retirement.py               |  2 +-
 coreblocks/structs_common/csr.py              |  2 +-
 coreblocks/structs_common/exception.py        |  2 +-
 .../structs_common/interrupt_controller.py    |  2 +-
 coreblocks/utils/protocols.py                 |  9 +---
 test/fu/functional_common.py                  |  2 +-
 test/lsu/test_dummylsu.py                     |  2 +-
 test/lsu/test_pma.py                          |  2 +-
 test/structs_common/test_csr.py               |  2 +-
 transactron/lib/dependencies.py               | 40 +++++++++++++++++
 transactron/lib/transformers.py               | 25 +++++++----
 transactron/utils/__init__.py                 |  2 +
 transactron/utils/depcache.py                 | 44 +++++++++++++++++++
 .../utils}/dependencies.py                    | 32 +-------------
 23 files changed, 123 insertions(+), 105 deletions(-)
 create mode 100644 transactron/lib/dependencies.py
 create mode 100644 transactron/utils/depcache.py
 rename {coreblocks/params => transactron/utils}/dependencies.py (75%)

diff --git a/coreblocks/core.py b/coreblocks/core.py
index 7db5d21b1..509cc7339 100644
--- a/coreblocks/core.py
+++ b/coreblocks/core.py
@@ -1,6 +1,6 @@
 from amaranth import *
 
-from coreblocks.params.dependencies import DependencyManager
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.stages.func_blocks_unifier import FuncBlocksUnifier
 from coreblocks.structs_common.instr_counter import CoreInstructionCounter
 from coreblocks.structs_common.interrupt_controller import InterruptController
diff --git a/coreblocks/fu/exception.py b/coreblocks/fu/exception.py
index 811423b0b..2927c0399 100644
--- a/coreblocks/fu/exception.py
+++ b/coreblocks/fu/exception.py
@@ -1,6 +1,6 @@
 from typing import Sequence
 from amaranth import *
-from coreblocks.params.dependencies import DependencyManager
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.params.isa import Funct3, ExceptionCause
 
 from transactron import *
diff --git a/coreblocks/fu/jumpbranch.py b/coreblocks/fu/jumpbranch.py
index 247ba9441..2b6c8aefa 100644
--- a/coreblocks/fu/jumpbranch.py
+++ b/coreblocks/fu/jumpbranch.py
@@ -7,6 +7,7 @@
 from transactron import *
 from transactron.core import def_method
 from transactron.lib import *
+from transactron.utils import DependencyManager
 
 from coreblocks.params import *
 from coreblocks.params.keys import AsyncInterruptInsertSignalKey
diff --git a/coreblocks/fu/priv.py b/coreblocks/fu/priv.py
index 48f2bbe9f..1dc31bfcc 100644
--- a/coreblocks/fu/priv.py
+++ b/coreblocks/fu/priv.py
@@ -6,6 +6,7 @@
 
 from transactron import *
 from transactron.lib import BasicFifo
+from transactron.utils import DependencyManager
 
 from coreblocks.params import *
 from coreblocks.params.keys import MretKey
diff --git a/coreblocks/lsu/dummyLsu.py b/coreblocks/lsu/dummyLsu.py
index a2099d873..cff013627 100644
--- a/coreblocks/lsu/dummyLsu.py
+++ b/coreblocks/lsu/dummyLsu.py
@@ -4,7 +4,7 @@
 from coreblocks.params import *
 from coreblocks.peripherals.wishbone import WishboneMaster
 from transactron.lib.connectors import Forwarder
-from transactron.utils import assign, ModuleLike
+from transactron.utils import assign, ModuleLike, DependencyManager
 from coreblocks.utils.protocols import FuncBlock
 
 from coreblocks.lsu.pma import PMAChecker
diff --git a/coreblocks/params/__init__.py b/coreblocks/params/__init__.py
index 8b2a588db..8e5e6dfc1 100644
--- a/coreblocks/params/__init__.py
+++ b/coreblocks/params/__init__.py
@@ -5,5 +5,4 @@
 from .fu_params import *  # noqa: F401
 from .keys import *  # noqa: F401
 from .icache_params import *  # noqa: F401
-from .dependencies import *  # noqa: F401
 from .instr import *  # noqa: F401
diff --git a/coreblocks/params/genparams.py b/coreblocks/params/genparams.py
index 113a90706..916832b49 100644
--- a/coreblocks/params/genparams.py
+++ b/coreblocks/params/genparams.py
@@ -1,13 +1,12 @@
 from __future__ import annotations
 
-from typing import TypeVar, Type, Any
 from amaranth.utils import log2_int
 
 from .isa import ISA, gen_isa_string
 from .icache_params import ICacheParameters
 from .fu_params import extensions_supported
 from ..peripherals.wishbone import WishboneParameters
-from transactron.utils import make_hashable
+from transactron.utils import DependentCache
 
 from typing import TYPE_CHECKING
 
@@ -16,45 +15,6 @@
 
 __all__ = ["GenParams"]
 
-T = TypeVar("T")
-
-
-class DependentCache:
-    """
-    Cache for classes, that depend on the `DependentCache` class itself.
-
-    Cached classes may accept one positional argument in the constructor, where this `DependentCache` class will
-    be passed. Classes may define any number keyword arguments in the constructor and separate cache entry will
-    be created for each set of the arguments.
-
-    Methods
-    -------
-    get: T, **kwargs -> T
-        Gets class `cls` from cache. Caches `cls` reference if this is the first call for it.
-        Optionally accepts `kwargs` for additional arguments in `cls` constructor.
-
-    """
-
-    def __init__(self):
-        self._depcache: dict[tuple[Type, Any], Type] = {}
-
-    def get(self, cls: Type[T], **kwargs) -> T:
-        cache_key = make_hashable(kwargs)
-        v = self._depcache.get((cls, cache_key), None)
-        if v is None:
-            positional_count = cls.__init__.__code__.co_argcount
-
-            # first positional arg is `self` field, second may be `DependentCache`
-            if positional_count > 2:
-                raise KeyError(f"Too many positional arguments in {cls!r} constructor")
-
-            if positional_count > 1:
-                v = cls(self, **kwargs)
-            else:
-                v = cls(**kwargs)
-            self._depcache[(cls, cache_key)] = v
-        return v
-
 
 class GenParams(DependentCache):
     def __init__(self, cfg: CoreConfiguration):
diff --git a/coreblocks/params/keys.py b/coreblocks/params/keys.py
index a3e2e0ead..6d4fee456 100644
--- a/coreblocks/params/keys.py
+++ b/coreblocks/params/keys.py
@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-from coreblocks.params.dependencies import SimpleKey, UnifierKey
+from transactron.lib.dependencies import SimpleKey, UnifierKey
 from transactron import Method
 from transactron.lib import MethodTryProduct, Collector
 from coreblocks.peripherals.wishbone import WishboneMaster
diff --git a/coreblocks/stages/func_blocks_unifier.py b/coreblocks/stages/func_blocks_unifier.py
index 8748bf5df..2888f9a81 100644
--- a/coreblocks/stages/func_blocks_unifier.py
+++ b/coreblocks/stages/func_blocks_unifier.py
@@ -2,11 +2,10 @@
 
 from amaranth import *
 
-from coreblocks.params import GenParams, BlockComponentParams, DependencyManager
-from coreblocks.params.dependencies import UnifierKey
+from coreblocks.params import GenParams, BlockComponentParams
+from transactron.lib.dependencies import UnifierKey, DependencyManager
 from transactron import Method, TModule
-from transactron.lib import MethodProduct, Collector
-from coreblocks.utils.protocols import Unifier
+from transactron.lib import MethodProduct, Collector, Unifier
 
 __all__ = ["FuncBlocksUnifier"]
 
diff --git a/coreblocks/stages/retirement.py b/coreblocks/stages/retirement.py
index 842ae3e7b..7a9a961e8 100644
--- a/coreblocks/stages/retirement.py
+++ b/coreblocks/stages/retirement.py
@@ -2,9 +2,9 @@
 
 from transactron.core import Method, Transaction, TModule
 from transactron.lib.simultaneous import condition
+from transactron.utils.dependencies import DependencyManager
 
 from coreblocks.params.genparams import GenParams
-from coreblocks.params.dependencies import DependencyManager
 from coreblocks.params.isa import ExceptionCause
 from coreblocks.params.keys import GenericCSRRegistersKey
 from coreblocks.structs_common.csr_generic import CSRAddress, DoubleCounterCSR
diff --git a/coreblocks/structs_common/csr.py b/coreblocks/structs_common/csr.py
index bebb5697b..9277e1c85 100644
--- a/coreblocks/structs_common/csr.py
+++ b/coreblocks/structs_common/csr.py
@@ -5,7 +5,7 @@
 from transactron import Method, def_method, Transaction, TModule
 from transactron.utils import assign, bits_from_int
 from coreblocks.params.genparams import GenParams
-from coreblocks.params.dependencies import DependencyManager, ListKey
+from transactron.utils.dependencies import DependencyManager, ListKey
 from coreblocks.params.fu_params import BlockComponentParams
 from coreblocks.params.layouts import FetchLayouts, FuncUnitLayouts, CSRLayouts
 from coreblocks.params.isa import Funct3, ExceptionCause
diff --git a/coreblocks/structs_common/exception.py b/coreblocks/structs_common/exception.py
index 4787a6f9d..3b44f7ca8 100644
--- a/coreblocks/structs_common/exception.py
+++ b/coreblocks/structs_common/exception.py
@@ -1,5 +1,5 @@
 from amaranth import *
-from coreblocks.params.dependencies import DependencyManager
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.params.genparams import GenParams
 
 from coreblocks.params.isa import ExceptionCause
diff --git a/coreblocks/structs_common/interrupt_controller.py b/coreblocks/structs_common/interrupt_controller.py
index 2fa9c932d..3b98f8a53 100644
--- a/coreblocks/structs_common/interrupt_controller.py
+++ b/coreblocks/structs_common/interrupt_controller.py
@@ -1,5 +1,5 @@
 from amaranth import *
-from coreblocks.params.dependencies import DependencyManager
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.params.genparams import GenParams
 from coreblocks.params.keys import AsyncInterruptInsertSignalKey, MretKey
 
diff --git a/coreblocks/utils/protocols.py b/coreblocks/utils/protocols.py
index 70f542bc2..9024b0d7c 100644
--- a/coreblocks/utils/protocols.py
+++ b/coreblocks/utils/protocols.py
@@ -3,14 +3,7 @@
 from transactron.utils._typing import HasElaborate
 
 
-__all__ = ["FuncUnit", "FuncBlock", "Unifier"]
-
-
-class Unifier(HasElaborate, Protocol):
-    method: Method
-
-    def __init__(self, targets: list[Method]):
-        ...
+__all__ = ["FuncUnit", "FuncBlock"]
 
 
 class FuncUnit(HasElaborate, Protocol):
diff --git a/test/fu/functional_common.py b/test/fu/functional_common.py
index 23c28c7ea..82d78beb0 100644
--- a/test/fu/functional_common.py
+++ b/test/fu/functional_common.py
@@ -9,7 +9,7 @@
 
 from coreblocks.params import GenParams
 from coreblocks.params.configurations import test_core_config
-from coreblocks.params.dependencies import DependencyManager
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.params.fu_params import FunctionalComponentParams
 from coreblocks.params.isa import Funct3, Funct7
 from coreblocks.params.keys import AsyncInterruptInsertSignalKey, ExceptionReportKey
diff --git a/test/lsu/test_dummylsu.py b/test/lsu/test_dummylsu.py
index 07372582c..58d554458 100644
--- a/test/lsu/test_dummylsu.py
+++ b/test/lsu/test_dummylsu.py
@@ -11,7 +11,7 @@
 from coreblocks.params.configurations import test_core_config
 from coreblocks.params.isa import *
 from coreblocks.params.keys import ExceptionReportKey
-from coreblocks.params.dependencies import DependencyManager
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.params.layouts import ExceptionRegisterLayouts
 from coreblocks.peripherals.wishbone import *
 from test.common import TestbenchIO, TestCaseWithSimulator, def_method_mock
diff --git a/test/lsu/test_pma.py b/test/lsu/test_pma.py
index 65dab095a..78869dd21 100644
--- a/test/lsu/test_pma.py
+++ b/test/lsu/test_pma.py
@@ -7,7 +7,7 @@
 from coreblocks.params.configurations import test_core_config
 from coreblocks.params.isa import *
 from coreblocks.params.keys import ExceptionReportKey
-from coreblocks.params.dependencies import DependencyManager
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.params.layouts import ExceptionRegisterLayouts
 from coreblocks.peripherals.wishbone import *
 from test.common import TestbenchIO, TestCaseWithSimulator, def_method_mock
diff --git a/test/structs_common/test_csr.py b/test/structs_common/test_csr.py
index 1d87d20ba..03b5a0e51 100644
--- a/test/structs_common/test_csr.py
+++ b/test/structs_common/test_csr.py
@@ -7,7 +7,7 @@
 from coreblocks.params.configurations import test_core_config
 from coreblocks.params.layouts import ExceptionRegisterLayouts
 from coreblocks.params.keys import AsyncInterruptInsertSignalKey, ExceptionReportKey
-from coreblocks.params.dependencies import DependencyManager
+from transactron.utils.dependencies import DependencyManager
 from coreblocks.frontend.decoder import OpType
 
 from ..common import *
diff --git a/transactron/lib/dependencies.py b/transactron/lib/dependencies.py
new file mode 100644
index 000000000..05ea3c42f
--- /dev/null
+++ b/transactron/lib/dependencies.py
@@ -0,0 +1,40 @@
+from collections.abc import Callable
+
+from .. import Method
+from .transformers import Unifier
+from ..utils.dependencies import *
+
+
+__all__ = [
+    "DependencyManager",
+    "DependencyKey",
+    "SimpleKey",
+    "ListKey",
+    "UnifierKey"
+]
+
+
+class UnifierKey(DependencyKey["Method", tuple["Method", dict[str, "Unifier"]]]):
+    """Base class for method unifier dependency keys.
+
+    Method unifier dependency keys are used to collect methods to be called by
+    some part of the core. As multiple modules may wish to be called, a method
+    unifier is used to present a single method interface to the caller, which
+    allows to customize the calling behavior.
+    """
+
+    unifier: Callable[[list["Method"]], "Unifier"]
+
+    def __init_subclass__(cls, unifier: Callable[[list["Method"]], "Unifier"], **kwargs) -> None:
+        cls.unifier = unifier
+        return super().__init_subclass__(**kwargs)
+
+    def combine(self, data: list["Method"]) -> tuple["Method", dict[str, "Unifier"]]:
+        if len(data) == 1:
+            return data[0], {}
+        else:
+            unifiers: dict[str, Unifier] = {}
+            unifier_inst = self.unifier(data)
+            unifiers[self.__class__.__name__ + "_unifier"] = unifier_inst
+            method = unifier_inst.method
+        return method, unifiers
diff --git a/transactron/lib/transformers.py b/transactron/lib/transformers.py
index 71c783679..c0034b67e 100644
--- a/transactron/lib/transformers.py
+++ b/transactron/lib/transformers.py
@@ -1,18 +1,18 @@
-from abc import ABC
 from amaranth import *
 
 from transactron.utils.transactron_helpers import get_src_loc
 from ..core import *
 from ..core import RecordDict
 from ..utils import SrcLoc
-from typing import Optional
+from typing import Optional, Protocol
 from collections.abc import Callable
-from transactron.utils import ValueLike, assign, AssignType, ModuleLike
+from transactron.utils import ValueLike, assign, AssignType, ModuleLike, HasElaborate
 from .connectors import Forwarder, ManyToOneConnectTrans, ConnectTrans
 from .simultaneous import condition
 
 __all__ = [
     "Transformer",
+    "Unifier",
     "MethodMap",
     "MethodFilter",
     "MethodProduct",
@@ -23,7 +23,7 @@
 ]
 
 
-class Transformer(ABC, Elaboratable):
+class Transformer(HasElaborate, Protocol):
     """Method transformer abstract class.
 
     Method transformers construct a new method which utilizes other methods.
@@ -49,7 +49,14 @@ def use(self, m: ModuleLike):
         return self.method
 
 
-class MethodMap(Transformer):
+class Unifier(Transformer, Protocol):
+    method: Method
+
+    def __init__(self, targets: list[Method]):
+        ...
+
+
+class MethodMap(Elaboratable, Transformer):
     """Bidirectional map for methods.
 
     Takes a target method and creates a transformed method which calls the
@@ -110,7 +117,7 @@ def _(arg):
         return m
 
 
-class MethodFilter(Transformer):
+class MethodFilter(Elaboratable, Transformer):
     """Method filter.
 
     Takes a target method and creates a method which calls the target method
@@ -192,7 +199,7 @@ def _(arg):
         return m
 
 
-class MethodProduct(Transformer):
+class MethodProduct(Elaboratable, Unifier):
     def __init__(
         self,
         targets: list[Method],
@@ -246,7 +253,7 @@ def _(arg):
         return m
 
 
-class MethodTryProduct(Transformer):
+class MethodTryProduct(Elaboratable, Unifier):
     def __init__(
         self,
         targets: list[Method],
@@ -304,7 +311,7 @@ def _(arg):
         return m
 
 
-class Collector(Transformer):
+class Collector(Elaboratable, Unifier):
     """Single result collector.
 
     Creates method that collects results of many methods with identical
diff --git a/transactron/utils/__init__.py b/transactron/utils/__init__.py
index 8b437c355..cfe28ff2a 100644
--- a/transactron/utils/__init__.py
+++ b/transactron/utils/__init__.py
@@ -4,3 +4,5 @@
 from .assign import *  # noqa: F401
 from .amaranth_ext import *  # noqa: F401
 from .transactron_helpers import *  # noqa: F401
+from .dependencies import *  # noqa: F401
+from .depcache import *  # noqa: F401
diff --git a/transactron/utils/depcache.py b/transactron/utils/depcache.py
new file mode 100644
index 000000000..0fbe356c3
--- /dev/null
+++ b/transactron/utils/depcache.py
@@ -0,0 +1,44 @@
+from typing import TypeVar, Type, Any
+
+from transactron.utils import make_hashable
+
+__all__ = ["DependentCache"]
+
+T = TypeVar("T")
+
+
+class DependentCache:
+    """
+    Cache for classes, that depend on the `DependentCache` class itself.
+
+    Cached classes may accept one positional argument in the constructor, where this `DependentCache` class will
+    be passed. Classes may define any number keyword arguments in the constructor and separate cache entry will
+    be created for each set of the arguments.
+
+    Methods
+    -------
+    get: T, **kwargs -> T
+        Gets class `cls` from cache. Caches `cls` reference if this is the first call for it.
+        Optionally accepts `kwargs` for additional arguments in `cls` constructor.
+
+    """
+
+    def __init__(self):
+        self._depcache: dict[tuple[Type, Any], Type] = {}
+
+    def get(self, cls: Type[T], **kwargs) -> T:
+        cache_key = make_hashable(kwargs)
+        v = self._depcache.get((cls, cache_key), None)
+        if v is None:
+            positional_count = cls.__init__.__code__.co_argcount
+
+            # first positional arg is `self` field, second may be `DependentCache`
+            if positional_count > 2:
+                raise KeyError(f"Too many positional arguments in {cls!r} constructor")
+
+            if positional_count > 1:
+                v = cls(self, **kwargs)
+            else:
+                v = cls(**kwargs)
+            self._depcache[(cls, cache_key)] = v
+        return v
diff --git a/coreblocks/params/dependencies.py b/transactron/utils/dependencies.py
similarity index 75%
rename from coreblocks/params/dependencies.py
rename to transactron/utils/dependencies.py
index 5b8b11cd5..2aa0c73df 100644
--- a/coreblocks/params/dependencies.py
+++ b/transactron/utils/dependencies.py
@@ -1,16 +1,14 @@
 from collections import defaultdict
 
 from abc import abstractmethod, ABC
-from collections.abc import Callable
 from typing import Any, Generic, TypeVar
 
-from transactron import Method
-from coreblocks.utils.protocols import Unifier
-
 
 __all__ = [
     "DependencyManager",
     "DependencyKey",
+    "SimpleKey",
+    "ListKey"
 ]
 
 T = TypeVar("T")
@@ -84,32 +82,6 @@ def combine(self, data: list[T]) -> list[T]:
         return data
 
 
-class UnifierKey(DependencyKey[Method, tuple[Method, dict[str, Unifier]]]):
-    """Base class for method unifier dependency keys.
-
-    Method unifier dependency keys are used to collect methods to be called by
-    some part of the core. As multiple modules may wish to be called, a method
-    unifier is used to present a single method interface to the caller, which
-    allows to customize the calling behavior.
-    """
-
-    unifier: Callable[[list[Method]], Unifier]
-
-    def __init_subclass__(cls, unifier: Callable[[list[Method]], Unifier], **kwargs) -> None:
-        cls.unifier = unifier
-        return super().__init_subclass__(**kwargs)
-
-    def combine(self, data: list[Method]) -> tuple[Method, dict[str, Unifier]]:
-        if len(data) == 1:
-            return data[0], {}
-        else:
-            unifiers: dict[str, Unifier] = {}
-            unifier_inst = self.unifier(data)
-            unifiers[self.__class__.__name__ + "_unifier"] = unifier_inst
-            method = unifier_inst.method
-        return method, unifiers
-
-
 class DependencyManager:
     """Dependency manager.