Clean up context imports

test/backend/test_annoucement.py

@@ -3,13 +3,12 @@
 from functools import reduce
 
 from amaranth import *
-from amaranth_types.types import TestbenchContext
 from transactron.lib import FIFO, AdapterTrans, Adapter, ManyToOneConnectTrans
 from coreblocks.backend.annoucement import ResultAnnouncement
 from coreblocks.interface.layouts import *
 from coreblocks.params import GenParams
 from coreblocks.params.configurations import test_core_config
-from transactron.testing import TestCaseWithSimulator, TestbenchIO
+from transactron.testing import TestCaseWithSimulator, TestbenchIO, TestbenchContext
 
 
 class BackendTestCircuit(Elaboratable):

test/cache/test_icache.py

@@ -1,6 +1,4 @@
 from collections import deque
-from amaranth_types import AnySimulatorContext
-from amaranth_types.types import TestbenchContext
 from parameterized import parameterized_class
 import random
 
@@ -16,7 +14,7 @@
 from coreblocks.params.configurations import test_core_config
 from coreblocks.cache.refiller import SimpleCommonBusCacheRefiller
 
-from transactron.testing import TestCaseWithSimulator, TestbenchIO, def_method_mock
+from transactron.testing import TestCaseWithSimulator, TestbenchIO, def_method_mock, TestbenchContext
 from transactron.testing.functions import MethodData
 from transactron.testing.method_mock import MethodMock
 from transactron.testing.testbenchio import CallTrigger
@@ -118,7 +116,7 @@ async def wishbone_slave(self, sim: TestbenchContext):
 
             await self.test_module.wb_ctrl.slave_respond(sim, data, err=err)
 
-    async def refiller_process(self, sim: AnySimulatorContext):
+    async def refiller_process(self, sim: TestbenchContext):
         while self.requests:
             req_addr = self.requests.pop()
             await self.test_module.start_refill.call(sim, addr=req_addr)
@@ -389,11 +387,11 @@ def add_bad_addr(self, addr: int):
         self.bad_addrs.add(addr)
         self.bad_cache_lines.add(addr & ~((1 << self.cp.offset_bits) - 1))
 
-    async def send_req(self, sim: AnySimulatorContext, addr: int):
+    async def send_req(self, sim: TestbenchContext, addr: int):
         self.issued_requests.append(addr)
         await self.m.issue_req.call(sim, addr=addr)
 
-    async def expect_resp(self, sim: AnySimulatorContext, wait=False):
+    async def expect_resp(self, sim: TestbenchContext, wait=False):
         if wait:
             *_, resp = await self.m.accept_res.sample_outputs_until_done(sim)
         else:
@@ -417,7 +415,7 @@ def assert_resp(self, resp: MethodData):
     def expect_refill(self, addr: int):
         assert self.refill_requests.popleft() == addr
 
-    async def call_cache(self, sim: AnySimulatorContext, addr: int):
+    async def call_cache(self, sim: TestbenchContext, addr: int):
         await self.send_req(sim, addr)
         self.m.accept_res.enable(sim)
         await self.expect_resp(sim, wait=True)

test/core_structs/test_rat.py

@@ -1,5 +1,4 @@
-from amaranth_types.types import TestbenchContext
-from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit
+from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit, TestbenchContext
 
 from coreblocks.core_structs.rat import FRAT, RRAT
 from coreblocks.params import GenParams

test/core_structs/test_reorder_buffer.py

@@ -1,5 +1,4 @@
-from amaranth_types.types import TestbenchContext
-from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit
+from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit, TestbenchContext
 
 from coreblocks.core_structs.rob import ReorderBuffer
 from coreblocks.params import GenParams

test/frontend/test_decode_stage.py

@@ -1,8 +1,7 @@
-from amaranth_types.types import TestbenchContext
 import pytest
 from transactron.lib import AdapterTrans, FIFO
 from transactron.utils.amaranth_ext.elaboratables import ModuleConnector
-from transactron.testing import TestCaseWithSimulator, TestbenchIO, SimpleTestCircuit
+from transactron.testing import TestCaseWithSimulator, TestbenchIO, SimpleTestCircuit, TestbenchContext
 
 from coreblocks.frontend.decoder.decode_stage import DecodeStage
 from coreblocks.params import GenParams

test/frontend/test_fetch.py

@@ -1,4 +1,3 @@
-from amaranth_types.types import ProcessContext, TestbenchContext
 import pytest
 from typing import Optional
 from collections import deque
@@ -18,6 +17,8 @@
     TestbenchIO,
     def_method_mock,
     SimpleTestCircuit,
+    TestbenchContext,
+    ProcessContext,
 )
 
 from coreblocks.frontend.fetch.fetch import FetchUnit, PredictionChecker

test/frontend/test_instr_decoder.py

@@ -1,7 +1,6 @@
 from amaranth.sim import *
-from amaranth_types.types import TestbenchContext
 
-from transactron.testing import TestCaseWithSimulator
+from transactron.testing import TestCaseWithSimulator, TestbenchContext
 
 from coreblocks.params import *
 from coreblocks.params.configurations import test_core_config

test/frontend/test_rvc.py

@@ -1,4 +1,3 @@
-from amaranth_types.types import TestbenchContext
 from parameterized import parameterized_class
 
 from amaranth import *
@@ -9,7 +8,7 @@
 from coreblocks.params.configurations import test_core_config
 from transactron.utils import ValueLike
 
-from transactron.testing import TestCaseWithSimulator
+from transactron.testing import TestCaseWithSimulator, TestbenchContext
 
 COMMON_TESTS = [
     # Illegal instruction

test/func_blocks/fu/common/test_rs.py

@@ -1,9 +1,8 @@
 import random
 from collections import deque
-from amaranth_types.types import TestbenchContext
 from parameterized import parameterized_class
 
-from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit
+from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit, TestbenchContext
 
 from coreblocks.func_blocks.fu.common.rs import RS, RSBase
 from coreblocks.func_blocks.fu.common.fifo_rs import FifoRS

test/func_blocks/fu/functional_common.py

@@ -1,7 +1,6 @@
 from dataclasses import asdict, dataclass
 from itertools import product
 import random
-from amaranth_types.types import ProcessContext, TestbenchContext
 import pytest
 from collections import deque
 from typing import Generic, TypeVar
@@ -24,6 +23,8 @@
     TestbenchIO,
     TestCaseWithSimulator,
     SimpleTestCircuit,
+    ProcessContext,
+    TestbenchContext,
 )
 from transactron.utils import ModuleConnector
 

test/func_blocks/fu/test_fu_decoder.py

@@ -1,8 +1,7 @@
 import random
 from collections.abc import Sequence
-from amaranth_types.types import TestbenchContext
 
-from transactron.testing import SimpleTestCircuit, TestCaseWithSimulator
+from transactron.testing import SimpleTestCircuit, TestCaseWithSimulator, TestbenchContext
 
 from coreblocks.func_blocks.fu.common.fu_decoder import DecoderManager, Decoder
 from coreblocks.arch import OpType, Funct3, Funct7

test/func_blocks/fu/test_pipelined_mul_unit.py

@@ -2,12 +2,11 @@
 import math
 from collections import deque
 
-from amaranth_types.types import TestbenchContext
 from parameterized import parameterized_class
 
 from coreblocks.func_blocks.fu.unsigned_multiplication.pipelined import PipelinedUnsignedMul
 
-from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit
+from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit, TestbenchContext
 
 from coreblocks.params import GenParams
 from coreblocks.params.configurations import test_core_config

test/func_blocks/fu/test_unsigned_mul_unit.py

@@ -1,7 +1,6 @@
 import random
 from collections import deque
 
-from amaranth_types.types import TestbenchContext
 from parameterized import parameterized_class
 
 from coreblocks.func_blocks.fu.unsigned_multiplication.common import MulBaseUnsigned
@@ -10,7 +9,7 @@
 from coreblocks.func_blocks.fu.unsigned_multiplication.shift import ShiftUnsignedMul
 from coreblocks.func_blocks.fu.unsigned_multiplication.pipelined import PipelinedUnsignedMul
 
-from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit
+from transactron.testing import TestCaseWithSimulator, SimpleTestCircuit, TestbenchContext
 
 from coreblocks.params import GenParams
 from coreblocks.params.configurations import test_core_config

test/func_blocks/lsu/test_dummylsu.py

@@ -1,8 +1,6 @@
 import random
 from collections import deque
 
-from amaranth_types.types import TestbenchContext
-
 from transactron.lib import Adapter
 from transactron.testing.method_mock import MethodMock
 from transactron.utils import int_to_signed, signed_to_int
@@ -14,7 +12,7 @@
 from transactron.utils.dependencies import DependencyContext
 from coreblocks.interface.layouts import ExceptionRegisterLayouts, RetirementLayouts
 from coreblocks.peripherals.wishbone import *
-from transactron.testing import TestbenchIO, TestCaseWithSimulator, def_method_mock
+from transactron.testing import TestbenchIO, TestCaseWithSimulator, def_method_mock, TestbenchContext
 from coreblocks.peripherals.bus_adapter import WishboneMasterAdapter
 from test.peripherals.test_wishbone import WishboneInterfaceWrapper
 

test/func_blocks/lsu/test_pma.py

@@ -1,5 +1,4 @@
 import random
-from amaranth_types.types import TestbenchContext
 from coreblocks.func_blocks.fu.lsu.pma import PMAChecker, PMARegion
 
 from transactron.lib import Adapter
@@ -12,7 +11,7 @@
 from transactron.utils.dependencies import DependencyContext
 from coreblocks.interface.layouts import ExceptionRegisterLayouts, RetirementLayouts
 from coreblocks.peripherals.wishbone import *
-from transactron.testing import TestbenchIO, TestCaseWithSimulator, def_method_mock
+from transactron.testing import TestbenchIO, TestCaseWithSimulator, def_method_mock, TestbenchContext
 from coreblocks.peripherals.bus_adapter import WishboneMasterAdapter
 from test.peripherals.test_wishbone import WishboneInterfaceWrapper
 

test/peripherals/test_wishbone.py

@@ -3,7 +3,7 @@
 from collections import deque
 
 from amaranth.lib.wiring import connect
-from amaranth_types import AnySimulatorContext, ValueLike
+from amaranth_types import ValueLike
 
 from coreblocks.peripherals.wishbone import *
 
@@ -16,19 +16,19 @@ class WishboneInterfaceWrapper:
     def __init__(self, wishbone_interface: WishboneInterface):
         self.wb = wishbone_interface
 
-    def master_set(self, sim: AnySimulatorContext, addr: int, data: int, we: int):
+    def master_set(self, sim: SimulatorContext, addr: int, data: int, we: int):
         sim.set(self.wb.dat_w, data)
         sim.set(self.wb.adr, addr)
         sim.set(self.wb.we, we)
         sim.set(self.wb.cyc, 1)
         sim.set(self.wb.stb, 1)
 
-    def master_release(self, sim: AnySimulatorContext, release_cyc: bool = True):
+    def master_release(self, sim: SimulatorContext, release_cyc: bool = True):
         sim.set(self.wb.stb, 0)
         if release_cyc:
             sim.set(self.wb.cyc, 0)
 
-    async def slave_wait(self, sim: AnySimulatorContext):
+    async def slave_wait(self, sim: SimulatorContext):
         *_, adr, we, sel, dat_w = (
             await sim.tick()
             .sample(self.wb.adr, self.wb.we, self.wb.sel, self.wb.dat_w)
@@ -37,7 +37,7 @@ async def slave_wait(self, sim: AnySimulatorContext):
         return adr, we, sel, dat_w
 
     async def slave_wait_and_verify(
-        self, sim: AnySimulatorContext, exp_addr: int, exp_data: int, exp_we: int, exp_sel: int = 0
+        self, sim: SimulatorContext, exp_addr: int, exp_data: int, exp_we: int, exp_sel: int = 0
     ):
         adr, we, sel, dat_w = await self.slave_wait(sim)
 
@@ -48,7 +48,7 @@ async def slave_wait_and_verify(
             assert dat_w == exp_data
 
     async def slave_tick_and_verify(
-        self, sim: AnySimulatorContext, exp_addr: int, exp_data: int, exp_we: int, exp_sel: int = 0
+        self, sim: SimulatorContext, exp_addr: int, exp_data: int, exp_we: int, exp_sel: int = 0
     ):
         *_, adr, we, sel, dat_w, stb, cyc = await sim.tick().sample(
             self.wb.adr, self.wb.we, self.wb.sel, self.wb.dat_w, self.wb.stb, self.wb.cyc
@@ -63,7 +63,7 @@ async def slave_tick_and_verify(
 
     async def slave_respond(
         self,
-        sim: AnySimulatorContext,
+        sim: SimulatorContext,
         data: int,
         ack: int = 1,
         err: int = 0,
@@ -83,12 +83,12 @@ async def slave_respond(
         return ret
 
     async def slave_respond_master_verify(
-        self, sim: AnySimulatorContext, master: WishboneInterface, data: int, ack: int = 1, err: int = 0, rty: int = 0
+        self, sim: SimulatorContext, master: WishboneInterface, data: int, ack: int = 1, err: int = 0, rty: int = 0
     ):
         *_, ack, dat_r = await self.slave_respond(sim, data, ack, err, rty, sample=[master.ack, master.dat_r])
         assert ack and dat_r == data
 
-    async def wait_ack(self, sim: AnySimulatorContext):
+    async def wait_ack(self, sim: SimulatorContext):
         await sim.tick().until(self.wb.stb & self.wb.cyc & self.wb.ack)
 
 

test/regression/pysim.py

@@ -4,7 +4,6 @@
 
 from amaranth.utils import exact_log2
 from amaranth import *
-from amaranth_types.types import TestbenchContext
 
 from transactron.core.keys import TransactionManagerKey
 from transactron.profiler import Profile
@@ -18,6 +17,7 @@
     profiler_process,
     make_logging_process,
     parse_logging_level,
+    TestbenchContext,
 )
 from transactron.utils.dependencies import DependencyContext, DependencyManager
 from transactron.lib.metrics import HardwareMetricsManager

test/scheduler/test_rs_selection.py

@@ -2,7 +2,6 @@
 import random
 
 from amaranth import *
-from amaranth_types.types import TestbenchContext
 
 from coreblocks.params import GenParams
 from coreblocks.interface.layouts import RSLayouts, SchedulerLayouts
@@ -11,7 +10,7 @@
 from coreblocks.params.configurations import test_core_config
 from coreblocks.scheduler.scheduler import RSSelection
 from transactron.lib import FIFO, Adapter, AdapterTrans
-from transactron.testing import TestCaseWithSimulator, TestbenchIO
+from transactron.testing import TestCaseWithSimulator, TestbenchIO, TestbenchContext
 from transactron.testing.functions import data_const_to_dict
 from transactron.testing.method_mock import MethodMock, def_method_mock
 

test/scheduler/test_scheduler.py

@@ -3,7 +3,6 @@
 from collections import namedtuple, deque
 from typing import Callable, Optional, Iterable
 from amaranth import *
-from amaranth_types.types import TestbenchContext
 from parameterized import parameterized_class
 from coreblocks.interface.keys import CoreStateKey
 from coreblocks.interface.layouts import RetirementLayouts
@@ -23,7 +22,7 @@
 from coreblocks.params.configurations import test_core_config
 from coreblocks.core_structs.rob import ReorderBuffer
 from coreblocks.func_blocks.interface.func_protocols import FuncBlock
-from transactron.testing import TestCaseWithSimulator, TestbenchIO, def_method_mock
+from transactron.testing import TestCaseWithSimulator, TestbenchIO, def_method_mock, TestbenchContext
 
 
 class SchedulerTestCircuit(Elaboratable):

test/scheduler/test_wakeup_select.py

@@ -7,8 +7,6 @@
 from inspect import isclass
 import random
 
-from amaranth_types.types import TestbenchContext
-
 from coreblocks.params import GenParams
 from coreblocks.interface.layouts import RSLayouts
 from coreblocks.params.configurations import test_core_config
@@ -17,7 +15,7 @@
 from transactron.lib import Adapter
 from coreblocks.scheduler.wakeup_select import *
 
-from transactron.testing import RecordIntDict, TestCaseWithSimulator, TestbenchIO
+from transactron.testing import RecordIntDict, TestCaseWithSimulator, TestbenchIO, TestbenchContext
 from transactron.testing.functions import data_const_to_dict
 
 

test/test_core.py

@@ -3,12 +3,11 @@
 from amaranth import *
 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
 
-from transactron.testing import TestCaseWithSimulator
+from transactron.testing import TestCaseWithSimulator, ProcessContext, TestbenchContext
 
 from coreblocks.arch.isa_consts import PrivilegeLevel
 from coreblocks.core import Core

test/transactron/lib/test_fifo.py

@@ -1,9 +1,8 @@
 from amaranth import *
-from amaranth_types.types import TestbenchContext
 
 from transactron.lib import AdapterTrans, BasicFifo
 
-from transactron.testing import TestCaseWithSimulator, TestbenchIO, data_layout
+from transactron.testing import TestCaseWithSimulator, TestbenchIO, data_layout, TestbenchContext
 from collections import deque
 from parameterized import parameterized_class
 import random