diff --git a/src/definitions.v b/src/definitions.v
index 1e934bc811da49b755961764c1d668b0dbe69ee9..788c278cc6c7a926727067dcf64389ce5d56eba7 100644
--- a/src/definitions.v
+++ b/src/definitions.v
@@ -2,7 +2,7 @@
    various related proofs. *)
 (* They may or may not be useful… *)
 
-Require Import Bool List Nat.
+Require Import Bool List Nat PeanoNat.
 Import ListNotations.
 
 (* Basic definitions: opcodes, stages, store buffer state. *)
@@ -119,9 +119,21 @@ Definition state := (stage * nat)%type.
 (* Comparison between two states. *)
 Definition state_leb st1 st2 :=
   match st1, st2 with
-  | (s1, n1), (s2, n2) => leb s1 s2 && Nat.leb n2 n1
+  | (s1, n1), (s2, n2) =>
+      if stage_beq s1 s2 then
+             Nat.leb n2 n1
+      else
+             leb s1 s2
   end.
 
+Lemma state_leb_eq_true : forall st, state_leb st st = true.
+Proof.
+  intro st.
+  destruct st.
+  simpl.
+  now rewrite stage_beq_eq_true, Nat.leb_le.
+Qed.
+
 (* Type definitions etc. *)
 Definition instr_kind := (opcode * state)%type.
 Definition trace_kind := list instr_kind.
@@ -135,6 +147,16 @@ Definition compare_two (e1 e2 : option instr_kind) :=
   | Some (_, st1), Some (_, st2) => state_leb st1 st2
   end.
 
+Lemma compare_two_eq_true : forall e, e <> None -> compare_two e e = true.
+Proof.
+  intros e Hn.
+  destruct e.
+  - destruct i.
+    simpl.
+    now rewrite state_leb_eq_true.
+  - contradiction.
+Qed.
+
 Definition get_opc (trace : trace_kind) :=
   fst (List.split trace).