From 3a98e42b966cd01c6fde6122467d4be3eca3844f Mon Sep 17 00:00:00 2001
From: Marek Materzok <tilk@tilk.eu>
Date: Wed, 13 Nov 2024 14:35:29 +0100
Subject: [PATCH] Use infrastructure cycle counter

---
 test/test_core.py                | 16 +++++++---------
 test/transactron/test_metrics.py | 23 ++++++++---------------
 2 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/test/test_core.py b/test/test_core.py
index 620f09679..3cd0a61dd 100644
--- a/test/test_core.py
+++ b/test/test_core.py
@@ -4,6 +4,7 @@
 from amaranth.lib.wiring import connect
 from amaranth_types import ValueLike
 from amaranth_types.types import ProcessContext, TestbenchContext
+from transactron.testing.tick_count import TicksKey
 
 from transactron.utils import align_to_power_of_two
 
@@ -21,6 +22,8 @@
 import tempfile
 from parameterized import parameterized_class
 
+from transactron.utils.dependencies import DependencyContext
+
 
 class CoreTestElaboratable(Elaboratable):
     def __init__(self, gen_params: GenParams, instr_mem: list[int] = [0], data_mem: list[int] = []):
@@ -308,23 +311,19 @@ def setup_method(self):
         self.gen_params = GenParams(self.configuration)
         random.seed(161453)
 
-    # TODO: global cycle counter?
-    async def cycle_counter_process(self, sim: ProcessContext):
-        self.cycles = 0
-        async for _ in sim.tick():
-            self.cycles += 1
-
     async def run_with_interrupt_process(self, sim: TestbenchContext):
+        ticks = DependencyContext.get().get_dependency(TicksKey())
+
         # wait for interrupt enable
         async def wait_or_timeout(cond: ValueLike, pred: Callable[[Any], bool]):
             async for *_, value in sim.tick().sample(cond):
-                if pred(value) or self.cycles >= self.cycle_count:
+                if pred(value) or sim.get(ticks) >= self.cycle_count:
                     break
 
         await wait_or_timeout(self.m.core.interrupt_controller.mie.value, lambda value: value != 0)
         await self.random_wait(sim, 5)
 
-        while self.cycles < self.cycle_count:
+        while sim.get(ticks) < self.cycle_count:
             sim.set(self.m.interrupt_level, 1)
 
             if self.always_mmode:  # if test happens only in m_mode, just enable fixed interrupt
@@ -353,5 +352,4 @@ def test_interrupted_prog(self):
         self.m = CoreTestElaboratable(self.gen_params, instr_mem=bin_src["text"], data_mem=bin_src["data"])
 
         with self.run_simulation(self.m) as sim:
-            sim.add_process(self.cycle_counter_process)
             sim.add_testbench(self.run_with_interrupt_process)
diff --git a/test/transactron/test_metrics.py b/test/transactron/test_metrics.py
index 3278022fc..e7c5c21d4 100644
--- a/test/transactron/test_metrics.py
+++ b/test/transactron/test_metrics.py
@@ -3,6 +3,7 @@
 import queue
 from typing import Type
 from enum import IntFlag, IntEnum, auto, Enum
+from amaranth_types.types import TestbenchContext
 
 from parameterized import parameterized_class
 
@@ -337,39 +338,31 @@ def test_latency_measurer(self):
 
         finish = False
 
-        # TODO: timer in testing infrastructure
-        time = 0
-
-        async def timer(sim):
-            nonlocal time
-
-            while True:
-                await sim.tick()
-                time += 1
-
-        async def producer(sim):
+        async def producer(sim: TestbenchContext):
             nonlocal finish
+            ticks = DependencyContext.get().get_dependency(TicksKey())
 
             for _ in range(200):
                 await m._start.call(sim)
 
-                event_queue.put(time)
+                event_queue.put(sim.get(ticks))
                 await self.random_wait_geom(sim, 0.8)
 
             finish = True
 
-        async def consumer(sim):
+        async def consumer(sim: TestbenchContext):
+            ticks = DependencyContext.get().get_dependency(TicksKey())
+
             while not finish:
                 await m._stop.call(sim)
 
-                latencies.append(time - event_queue.get())
+                latencies.append(sim.get(ticks) - event_queue.get())
 
                 await self.random_wait_geom(sim, 1.0 / self.expected_consumer_wait)
 
             self.check_latencies(sim, m, latencies)
 
         with self.run_simulation(m) as sim:
-            sim.add_process(timer)
             sim.add_testbench(producer)
             sim.add_testbench(consumer)