[otbnsim] Add more functionality to testutil

Add more functionality from the statistics test to the shared testutil
module, making it available for other tests to use as well.

Signed-off-by: Philipp Wagner <phw@lowrisc.org>

hw/ip/otbn/dv/otbnsim/test/simple_test.py

diff --git a/hw/ip/otbn/dv/otbnsim/test/simple_test.py b/hw/ip/otbn/dv/otbnsim/test/simple_test.py
index 6f8f47e..1bf7711 100644
--- a/hw/ip/otbn/dv/otbnsim/test/simple_test.py
+++ b/hw/ip/otbn/dv/otbnsim/test/simple_test.py
@@ -142,7 +142,7 @@
                asm_file: str,
                expected_file: str) -> None:
     # Start by assembling and linking the input file
-    elf_file = asm_and_link_one_file(asm_file, str(tmpdir))
+    elf_file = asm_and_link_one_file(asm_file, tmpdir)
 
     # Run the simulation. We can just pass a list of commands to stdin, and
     # don't need to do anything clever to track what's going on.

hw/ip/otbn/dv/otbnsim/test/stats_test.py

diff --git a/hw/ip/otbn/dv/otbnsim/test/stats_test.py b/hw/ip/otbn/dv/otbnsim/test/stats_test.py
index 99fa087..c820a31 100644
--- a/hw/ip/otbn/dv/otbnsim/test/stats_test.py
+++ b/hw/ip/otbn/dv/otbnsim/test/stats_test.py
@@ -7,31 +7,34 @@
 
 from sim.sim import OTBNSim
 from sim.stats import ExecutionStats
-from sim.elf import load_elf
 import testutil
 
 
-def _run_sim_for_stats(asm_file: str, tmpdir: str) -> ExecutionStats:
-    """ Run the OTBN simulator, collect statistics, and return them. """
-
-    assert os.path.exists(asm_file)
-    elf_file = testutil.asm_and_link_one_file(asm_file, tmpdir)
-
-    sim = OTBNSim()
-    load_elf(sim, elf_file)
-
-    sim.start(0)
+def _run_sim_for_stats(sim: OTBNSim) -> ExecutionStats:
     sim.run(verbose=False, collect_stats=True)
-    assert sim.stats
 
+    assert sim.stats
     return sim.stats
 
+
+def _simulate_asm_file(asm_file: str, tmpdir: py.path.local) -> ExecutionStats:
+    '''Run the OTBN simulator, collect statistics, and return them.'''
+
+    sim = testutil.prepare_sim_for_asm_file(asm_file, tmpdir, start_addr=0)
+    return _run_sim_for_stats(sim)
+
+
+def _simulate_asm_str(assembly: str, tmpdir: py.path.local) -> ExecutionStats:
+    sim = testutil.prepare_sim_for_asm_str(assembly, tmpdir, start_addr=0)
+    return _run_sim_for_stats(sim)
+
+
 def test_general_and_loop(tmpdir: py.path.local) -> None:
     '''Test the collection of general statistics as well as loop stats.'''
 
     asm_file = os.path.join(os.path.dirname(__file__),
                             'simple', 'loops', 'loops.s')
-    stats = _run_sim_for_stats(asm_file, str(tmpdir))
+    stats = _simulate_asm_file(asm_file, tmpdir)
 
     # General statistics
     assert stats.stall_count == 2
@@ -58,7 +61,7 @@
 
     asm_file = os.path.join(os.path.dirname(__file__),
                             'simple', 'subroutines', 'direct-call.s')
-    stats = _run_sim_for_stats(asm_file, str(tmpdir))
+    stats = _simulate_asm_file(asm_file, tmpdir)
 
     exp = [{'call_site': 4, 'callee_func': 12, 'caller_func': 0}]
     assert stats.func_calls == exp
@@ -69,7 +72,7 @@
 
     asm_file = os.path.join(os.path.dirname(__file__),
                             'simple', 'subroutines', 'indirect-call.s')
-    stats = _run_sim_for_stats(asm_file, str(tmpdir))
+    stats = _simulate_asm_file(asm_file, tmpdir)
 
     exp = [{'call_site': 8, 'callee_func': 16, 'caller_func': 0}]
     assert stats.func_calls == exp

hw/ip/otbn/dv/otbnsim/test/testutil.py

diff --git a/hw/ip/otbn/dv/otbnsim/test/testutil.py b/hw/ip/otbn/dv/otbnsim/test/testutil.py
index 7900e89..c37c01f 100644
--- a/hw/ip/otbn/dv/otbnsim/test/testutil.py
+++ b/hw/ip/otbn/dv/otbnsim/test/testutil.py
@@ -2,8 +2,13 @@
 # Licensed under the Apache License, Version 2.0, see LICENSE for details.
 # SPDX-License-Identifier: Apache-2.0
 
+import py
 import os
 import subprocess
+import tempfile
+
+from sim.elf import load_elf
+from sim.sim import OTBNSim
 
 
 OTBN_DIR = os.path.join(os.path.dirname(__file__), '../../..')
@@ -11,7 +16,7 @@
 SIM_DIR = os.path.join(os.path.dirname(__file__), '..')
 
 
-def asm_and_link_one_file(asm_path: str, work_dir: str) -> str:
+def asm_and_link_one_file(asm_path: str, work_dir: py.path.local) -> str:
     '''Assemble and link file at asm_path in work_dir.
 
     Returns the path to the resulting ELF
@@ -25,3 +30,33 @@
     subprocess.run([otbn_as, '-o', obj_path, asm_path], check=True)
     subprocess.run([otbn_ld, '-o', elf_path, obj_path], check=True)
     return elf_path
+
+
+def prepare_sim_for_asm_file(asm_file: str, tmpdir: py.path.local, start_addr: int = 0) -> OTBNSim:
+    '''Set up the simulation of a single assembly file.
+
+    The returned simulation is ready to be run through the run() method.
+
+    '''
+    assert os.path.exists(asm_file)
+    elf_file = asm_and_link_one_file(asm_file, tmpdir)
+
+    sim = OTBNSim()
+    load_elf(sim, elf_file)
+
+    sim.state.ext_regs.write('START_ADDR', start_addr, False)
+    sim.state.ext_regs.commit()
+    sim.start()
+    return sim
+
+
+def prepare_sim_for_asm_str(assembly: str, tmpdir: py.path.local, start_addr: int = 0) -> OTBNSim:
+    '''Set up the simulation for an assembly snippet passed as string.
+
+    The returned simulation is ready to be run through the run() method.
+
+    '''
+    with tempfile.NamedTemporaryFile('w', dir=tmpdir) as fp:
+        fp.write(assembly)
+        fp.flush()
+        return prepare_sim_for_asm_file(fp.name, tmpdir, start_addr)