[otbn,dv] Correctly model timing for IMEM integrity errors

This was broken by 0ac448a (adding the prefetch stage) because the RTL
now caches IMEM contents for a cycle, meaning that injected IMEM
integrity errors arrive a cycle later.

We also slightly re-jig the timing for where INSN_CNT gets zeroed to
match the new behaviour.

Signed-off-by: Rupert Swarbrick <rswarbrick@lowrisc.org>
diff --git a/hw/ip/otbn/dv/otbnsim/sim/sim.py b/hw/ip/otbn/dv/otbnsim/sim/sim.py
index 39c29b1..20bac52 100644
--- a/hw/ip/otbn/dv/otbnsim/sim/sim.py
+++ b/hw/ip/otbn/dv/otbnsim/sim/sim.py
@@ -154,11 +154,7 @@
             self.state.ext_regs.write('INSN_CNT', 0, True)
             return (None, changes)
 
-        if self.state.fsm_state == FsmState.POST_EXEC:
-            return (None, self._on_stall(verbose, fetch_next=False))
-
-        if self.state.fsm_state == FsmState.LOCKING:
-            self.state.ext_regs.write('INSN_CNT', 0, True)
+        if self.state.fsm_state in [FsmState.POST_EXEC, FsmState.LOCKING]:
             return (None, self._on_stall(verbose, fetch_next=False))
 
         assert self.state.fsm_state == FsmState.EXEC
diff --git a/hw/ip/otbn/dv/otbnsim/sim/state.py b/hw/ip/otbn/dv/otbnsim/sim/state.py
index 7eb485c..b92bb7a 100644
--- a/hw/ip/otbn/dv/otbnsim/sim/state.py
+++ b/hw/ip/otbn/dv/otbnsim/sim/state.py
@@ -104,8 +104,13 @@
         self.urnd_256b = 4 * [0]
         self.urnd_64b = 0
 
-        # This flag is set to true if we've injected integrity errors, trashing
-        # the whole of IMEM. The next fetch should fail.
+        # To simulate injecting integrity errors, we set a flag to say that
+        # IMEM is no longer readable without getting an error. This can't take
+        # effect instantly because the RTL's prefetch stage (which we don't
+        # model except for matching its timing) has a copy of the next
+        # instruction. Make this a counter, decremented once per cycle. When we
+        # get to zero, we set the flag.
+        self._time_to_imem_invalidation = None  # type: Optional[int]
         self.invalidated_imem = False
 
     def get_next_pc(self) -> int:
@@ -265,6 +270,12 @@
         # We shouldn't be running commit() in IDLE or LOCKED mode
         assert self.running()
 
+        if self._time_to_imem_invalidation is not None:
+            self._time_to_imem_invalidation -= 1
+            if self._time_to_imem_invalidation == 0:
+                self.invalidated_imem = True
+                self._time_to_imem_invalidation = None
+
         # Check if we processed the RND data, if so set the register. This is
         # done seperately from the rnd_completed method because in other case
         # we are exiting stall caused by RND waiting by one cycle too early.
@@ -327,8 +338,11 @@
         # POST_EXEC to allow one more cycle.
         if self.pending_halt:
             self.ext_regs.commit()
-            self.fsm_state = (FsmState.LOCKING
-                              if self._err_bits >> 16 else FsmState.POST_EXEC)
+            if self._err_bits >> 16:
+                self.ext_regs.write('INSN_CNT', 0, True)
+                self.fsm_state = FsmState.LOCKING
+            else:
+                self.fsm_state = FsmState.POST_EXEC
             return
 
         # As pending_halt wasn't set, there shouldn't be any pending error bits
@@ -498,3 +512,6 @@
         '''
         self._err_bits |= err_bits
         self.pending_halt = True
+
+    def invalidate_imem(self) -> None:
+        self._time_to_imem_invalidation = 2
diff --git a/hw/ip/otbn/dv/otbnsim/stepped.py b/hw/ip/otbn/dv/otbnsim/stepped.py
index 5e8f0c3..208fcf8 100755
--- a/hw/ip/otbn/dv/otbnsim/stepped.py
+++ b/hw/ip/otbn/dv/otbnsim/stepped.py
@@ -291,7 +291,7 @@
 def on_invalidate_imem(sim: OTBNSim, args: List[str]) -> Optional[OTBNSim]:
     check_arg_count('invalidate_imem', 0, args)
 
-    sim.state.invalidated_imem = True
+    sim.state.invalidate_imem()
 
     return None