From a6833faff4bdf19562eb27f48b887ce56c4e166f Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Sat, 11 Nov 2023 13:37:00 +0100
Subject: [PATCH 1/7] Add Now() support

---
 test/common/_test/__init__.py            |  0
 test/common/_test/test_infrastructure.py | 29 +++++++++++++++++++
 test/common/infrastructure.py            | 37 ++++++++++++++++++++++++
 3 files changed, 66 insertions(+)
 create mode 100644 test/common/_test/__init__.py
 create mode 100644 test/common/_test/test_infrastructure.py

diff --git a/test/common/_test/__init__.py b/test/common/_test/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/test/common/_test/test_infrastructure.py b/test/common/_test/test_infrastructure.py
new file mode 100644
index 000000000..dda5083af
--- /dev/null
+++ b/test/common/_test/test_infrastructure.py
@@ -0,0 +1,29 @@
+from amaranth import *
+from test.common import *
+
+class EmptyCircuit(Elaboratable):
+    def __init__(self):
+        pass
+
+    def elaborate(self, platform):
+        m = Module()
+        return m
+
+class TestNow(TestCaseWithSimulator):
+    def setUp(self):
+        self.test_cycles=10
+        self.m = SimpleTestCircuit(EmptyCircuit())
+
+    def process(self):
+        for k in range(self.test_cycles):
+            now = yield Now()
+            assert k == now
+            # check if second call don't change the returned value
+            now = yield Now()
+            assert k == now
+
+            yield
+
+    def test_random(self):
+        with self.run_simulation(self.m, 50) as sim:
+            sim.add_sync_process(self.process)
diff --git a/test/common/infrastructure.py b/test/common/infrastructure.py
index fe0b337d4..2d7c6b870 100644
--- a/test/common/infrastructure.py
+++ b/test/common/infrastructure.py
@@ -98,6 +98,39 @@ def elaborate(self, platform) -> HasElaborate:
 
         return m
 
+class CoreblockCommand:
+    pass
+
+
+class Now(CoreblockCommand):
+    pass
+
+
+class SyncProcessWrapper:
+    def __init__(self, f):
+        self.org_process = f
+        self.current_cycle = 0
+
+    def _wrapping_function(self):
+        response = None
+        org_corutine = self.org_process()
+        try:
+            while True:
+                # call orginal test process and catch data yielded by it in `command` variable
+                command = org_corutine.send(response)
+                # If process wait for new cycle
+                if command is None:
+                    self.current_cycle += 1
+                    # forward to amaranth
+                    yield
+                elif isinstance(command, Now):
+                    response = self.current_cycle
+                # Pass everything else to amaranth simulator without modifications
+                else:
+                    response = yield command
+        except StopIteration:
+            pass
+
 
 class PysimSimulator(Simulator):
     def __init__(self, module: HasElaborate, max_cycles: float = 10e4, add_transaction_module=True, traces_file=None):
@@ -133,6 +166,10 @@ def __init__(self, module: HasElaborate, max_cycles: float = 10e4, add_transacti
 
         self.deadline = clk_period * max_cycles
 
+    def add_sync_process(self, f):
+        f_wrapped = SyncProcessWrapper(f)
+        super().add_sync_process(f_wrapped._wrapping_function)
+
     def run(self) -> bool:
         with self.ctx:
             self.run_until(self.deadline)

From d0deb69546938fab8e4c7cfeeda53320b19db048 Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Sat, 11 Nov 2023 13:42:37 +0100
Subject: [PATCH 2/7] Lint; typing

---
 test/common/_test/test_infrastructure.py | 4 +++-
 test/common/functions.py                 | 7 +------
 test/common/infrastructure.py            | 8 ++++++--
 3 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/test/common/_test/test_infrastructure.py b/test/common/_test/test_infrastructure.py
index dda5083af..ecf1c84d9 100644
--- a/test/common/_test/test_infrastructure.py
+++ b/test/common/_test/test_infrastructure.py
@@ -1,6 +1,7 @@
 from amaranth import *
 from test.common import *
 
+
 class EmptyCircuit(Elaboratable):
     def __init__(self):
         pass
@@ -9,9 +10,10 @@ def elaborate(self, platform):
         m = Module()
         return m
 
+
 class TestNow(TestCaseWithSimulator):
     def setUp(self):
-        self.test_cycles=10
+        self.test_cycles = 10
         self.m = SimpleTestCircuit(EmptyCircuit())
 
     def process(self):
diff --git a/test/common/functions.py b/test/common/functions.py
index c4ffc814a..1f36e69fe 100644
--- a/test/common/functions.py
+++ b/test/common/functions.py
@@ -1,11 +1,6 @@
 from amaranth import *
-from amaranth.hdl.ast import Statement
-from amaranth.sim.core import Command
-from typing import TypeVar, Any, Generator, TypeAlias
 from coreblocks.utils._typing import RecordValueDict, RecordIntDict
-
-T = TypeVar("T")
-TestGen: TypeAlias = Generator[Command | Value | Statement | None, Any, T]
+from .infrastructure import TestGen
 
 
 def set_inputs(values: RecordValueDict, field: Record) -> TestGen[None]:
diff --git a/test/common/infrastructure.py b/test/common/infrastructure.py
index 2d7c6b870..6c0a1a271 100644
--- a/test/common/infrastructure.py
+++ b/test/common/infrastructure.py
@@ -3,9 +3,11 @@
 import unittest
 import functools
 from contextlib import contextmanager, nullcontext
-from typing import TypeVar, Generic, Type, TypeGuard, Any, Union, Callable, cast
+from typing import TypeVar, Generic, Type, TypeGuard, Any, Union, Callable, cast, Generator, TypeAlias
 from amaranth import *
 from amaranth.sim import *
+from amaranth.hdl.ast import Statement
+from amaranth.sim.core import Command
 from .testbenchio import TestbenchIO
 from ..gtkw_extension import write_vcd_ext
 from transactron import Method
@@ -15,6 +17,7 @@
 
 T = TypeVar("T")
 _T_nested_collection = T | list["_T_nested_collection[T]"] | dict[str, "_T_nested_collection[T]"]
+TestGen: TypeAlias = Generator[Command | Value | Statement | "CoreblockCommand" | None, Any, T]
 
 
 def guard_nested_collection(cont: Any, t: Type[T]) -> TypeGuard[_T_nested_collection[T]]:
@@ -98,6 +101,7 @@ def elaborate(self, platform) -> HasElaborate:
 
         return m
 
+
 class CoreblockCommand:
     pass
 
@@ -166,7 +170,7 @@ def __init__(self, module: HasElaborate, max_cycles: float = 10e4, add_transacti
 
         self.deadline = clk_period * max_cycles
 
-    def add_sync_process(self, f):
+    def add_sync_process(self, f: Callable[[], TestGen]):
         f_wrapped = SyncProcessWrapper(f)
         super().add_sync_process(f_wrapped._wrapping_function)
 

From 0c6e50f680bf45e3df88734f9744c160a3c26d3e Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Sat, 11 Nov 2023 14:22:11 +0100
Subject: [PATCH 3/7] Fix imports.

---
 test/common/functions.py      | 15 ++++++++++++---
 test/common/infrastructure.py |  8 +++-----
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/test/common/functions.py b/test/common/functions.py
index 1f36e69fe..3c37aec61 100644
--- a/test/common/functions.py
+++ b/test/common/functions.py
@@ -1,9 +1,18 @@
 from amaranth import *
+from typing import TYPE_CHECKING, Generator, Any, TypeAlias, TypeVar, Union
 from coreblocks.utils._typing import RecordValueDict, RecordIntDict
-from .infrastructure import TestGen
+from amaranth.hdl.ast import Statement
+from amaranth.sim.core import Command
 
+if TYPE_CHECKING:
+    from .infrastructure import CoreblockCommand
 
-def set_inputs(values: RecordValueDict, field: Record) -> TestGen[None]:
+
+T = TypeVar("T")
+TestGen: TypeAlias = Generator[Union[Command, Value, Statement, "CoreblockCommand", None], Any, T]
+
+
+def set_inputs(values: RecordValueDict, field: Record) -> "TestGen[None]":
     for name, value in values.items():
         if isinstance(value, dict):
             yield from set_inputs(value, getattr(field, name))
@@ -11,7 +20,7 @@ def set_inputs(values: RecordValueDict, field: Record) -> TestGen[None]:
             yield getattr(field, name).eq(value)
 
 
-def get_outputs(field: Record) -> TestGen[RecordIntDict]:
+def get_outputs(field: Record) -> "TestGen[RecordIntDict]":
     # return dict of all signal values in a record because amaranth's simulator can't read all
     # values of a Record in a single yield - it can only read Values (Signals)
     result = {}
diff --git a/test/common/infrastructure.py b/test/common/infrastructure.py
index 6c0a1a271..4eb9e3eac 100644
--- a/test/common/infrastructure.py
+++ b/test/common/infrastructure.py
@@ -3,12 +3,11 @@
 import unittest
 import functools
 from contextlib import contextmanager, nullcontext
-from typing import TypeVar, Generic, Type, TypeGuard, Any, Union, Callable, cast, Generator, TypeAlias
+from typing import TypeVar, Generic, Type, TypeGuard, Any, Union, Callable, cast, TypeAlias
 from amaranth import *
 from amaranth.sim import *
-from amaranth.hdl.ast import Statement
-from amaranth.sim.core import Command
 from .testbenchio import TestbenchIO
+from .functions import TestGen
 from ..gtkw_extension import write_vcd_ext
 from transactron import Method
 from transactron.lib import AdapterTrans
@@ -16,8 +15,7 @@
 from coreblocks.utils import ModuleConnector, HasElaborate, auto_debug_signals, HasDebugSignals
 
 T = TypeVar("T")
-_T_nested_collection = T | list["_T_nested_collection[T]"] | dict[str, "_T_nested_collection[T]"]
-TestGen: TypeAlias = Generator[Command | Value | Statement | "CoreblockCommand" | None, Any, T]
+_T_nested_collection: TypeAlias = T | list["_T_nested_collection[T]"] | dict[str, "_T_nested_collection[T]"]
 
 
 def guard_nested_collection(cont: Any, t: Type[T]) -> TypeGuard[_T_nested_collection[T]]:

From a978de9757bda1ba971d5abff943f4462442d614 Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Sat, 11 Nov 2023 14:23:26 +0100
Subject: [PATCH 4/7] Improvements

---
 test/common/functions.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/common/functions.py b/test/common/functions.py
index 3c37aec61..0c72fe318 100644
--- a/test/common/functions.py
+++ b/test/common/functions.py
@@ -1,8 +1,8 @@
 from amaranth import *
-from typing import TYPE_CHECKING, Generator, Any, TypeAlias, TypeVar, Union
-from coreblocks.utils._typing import RecordValueDict, RecordIntDict
 from amaranth.hdl.ast import Statement
 from amaranth.sim.core import Command
+from typing import TYPE_CHECKING, Generator, Any, TypeAlias, TypeVar, Union
+from coreblocks.utils._typing import RecordValueDict, RecordIntDict
 
 if TYPE_CHECKING:
     from .infrastructure import CoreblockCommand
@@ -12,7 +12,7 @@
 TestGen: TypeAlias = Generator[Union[Command, Value, Statement, "CoreblockCommand", None], Any, T]
 
 
-def set_inputs(values: RecordValueDict, field: Record) -> "TestGen[None]":
+def set_inputs(values: RecordValueDict, field: Record) -> TestGen[None]:
     for name, value in values.items():
         if isinstance(value, dict):
             yield from set_inputs(value, getattr(field, name))
@@ -20,7 +20,7 @@ def set_inputs(values: RecordValueDict, field: Record) -> "TestGen[None]":
             yield getattr(field, name).eq(value)
 
 
-def get_outputs(field: Record) -> "TestGen[RecordIntDict]":
+def get_outputs(field: Record) -> TestGen[RecordIntDict]:
     # return dict of all signal values in a record because amaranth's simulator can't read all
     # values of a Record in a single yield - it can only read Values (Signals)
     result = {}

From 3e0296baaada906f2f82d5716fbcfb5b18ffc909 Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Sat, 11 Nov 2023 14:24:25 +0100
Subject: [PATCH 5/7] Format

---
 test/common/functions.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/common/functions.py b/test/common/functions.py
index 0c72fe318..bdb4a7474 100644
--- a/test/common/functions.py
+++ b/test/common/functions.py
@@ -1,7 +1,7 @@
 from amaranth import *
 from amaranth.hdl.ast import Statement
 from amaranth.sim.core import Command
-from typing import TYPE_CHECKING, Generator, Any, TypeAlias, TypeVar, Union
+from typing import TypeVar, Any, Generator, TypeAlias, TYPE_CHECKING, Union
 from coreblocks.utils._typing import RecordValueDict, RecordIntDict
 
 if TYPE_CHECKING:

From bdcb7469e8ff516d8d9030ef81e3355c57fe6e40 Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Mon, 13 Nov 2023 18:35:17 +0100
Subject: [PATCH 6/7] Fix comments from review

---
 test/common/infrastructure.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/test/common/infrastructure.py b/test/common/infrastructure.py
index 4eb9e3eac..fb90e5505 100644
--- a/test/common/infrastructure.py
+++ b/test/common/infrastructure.py
@@ -4,6 +4,7 @@
 import functools
 from contextlib import contextmanager, nullcontext
 from typing import TypeVar, Generic, Type, TypeGuard, Any, Union, Callable, cast, TypeAlias
+from abc import ABC
 from amaranth import *
 from amaranth.sim import *
 from .testbenchio import TestbenchIO
@@ -100,7 +101,7 @@ def elaborate(self, platform) -> HasElaborate:
         return m
 
 
-class CoreblockCommand:
+class CoreblockCommand(ABC):
     pass
 
 
@@ -115,11 +116,11 @@ def __init__(self, f):
 
     def _wrapping_function(self):
         response = None
-        org_corutine = self.org_process()
+        org_coroutine = self.org_process()
         try:
             while True:
                 # call orginal test process and catch data yielded by it in `command` variable
-                command = org_corutine.send(response)
+                command = org_coroutine.send(response)
                 # If process wait for new cycle
                 if command is None:
                     self.current_cycle += 1

From e914a21c53f12ec3b561f1be235d6c646fada39d Mon Sep 17 00:00:00 2001
From: Lekcyjna <309016@uwr.edu.pl>
Date: Mon, 13 Nov 2023 18:36:18 +0100
Subject: [PATCH 7/7] Fix comments from review

---
 test/common/functions.py      | 4 ++--
 test/common/infrastructure.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/common/functions.py b/test/common/functions.py
index bdb4a7474..7a412e48c 100644
--- a/test/common/functions.py
+++ b/test/common/functions.py
@@ -5,11 +5,11 @@
 from coreblocks.utils._typing import RecordValueDict, RecordIntDict
 
 if TYPE_CHECKING:
-    from .infrastructure import CoreblockCommand
+    from .infrastructure import CoreblocksCommand
 
 
 T = TypeVar("T")
-TestGen: TypeAlias = Generator[Union[Command, Value, Statement, "CoreblockCommand", None], Any, T]
+TestGen: TypeAlias = Generator[Union[Command, Value, Statement, "CoreblocksCommand", None], Any, T]
 
 
 def set_inputs(values: RecordValueDict, field: Record) -> TestGen[None]:
diff --git a/test/common/infrastructure.py b/test/common/infrastructure.py
index fb90e5505..7bed7bd85 100644
--- a/test/common/infrastructure.py
+++ b/test/common/infrastructure.py
@@ -101,11 +101,11 @@ def elaborate(self, platform) -> HasElaborate:
         return m
 
 
-class CoreblockCommand(ABC):
+class CoreblocksCommand(ABC):
     pass
 
 
-class Now(CoreblockCommand):
+class Now(CoreblocksCommand):
     pass