From 7686c295d9b2337bbe7db5295e620a0ae2dc0621 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Fri, 30 Sep 2022 02:05:51 +0200
Subject: [PATCH 01/23] full replay only when necessary

---
 src/tracedAtomic.ml | 40 +++++++++++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 0180d64..627e5c3 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -165,10 +165,15 @@ let num_runs = ref 0
 (* we stash the current state in case a check fails and we need to log it *)
 let schedule_for_checks = ref []
 
-let do_run init_func init_schedule =
-  init_func (); (*set up run *)
-  tracing := true;
+let setup_run func init_schedule =
+  atomics_counter := 1;
+  CCVector.clear processes;
   schedule_for_checks := init_schedule;
+  func () ;
+  num_runs := !num_runs + 1;
+  finished_processes := 0
+
+let do_run init_schedule =
   (* cache the number of processes in case it's expensive*)
   let num_processes = CCVector.length processes in
   (* current number of ops we are through the current run *)
@@ -197,9 +202,7 @@ let do_run init_func init_schedule =
   in
   tracing := true;
   run_trace init_schedule ();
-  finished_processes := 0;
   tracing := false;
-  num_runs := !num_runs + 1;
   if !num_runs mod 1000 == 0 then
     Printf.printf "run: %d\n" !num_runs;
   let procs = CCVector.mapi (fun i p -> { proc_id = i; op = p.next_op; obj_ptr = p.next_repr }) processes |> CCVector.to_list in
@@ -208,8 +211,6 @@ let do_run init_func init_schedule =
                         |> Seq.filter (fun (_,proc) -> not proc.finished)
                         |> Seq.map (fun (id,_) -> id)
                         |> IntSet.of_seq in
-  CCVector.clear processes;
-  atomics_counter := 1;
   match last_element init_schedule with
   | (run_proc, run_op, run_ptr) ->
     { procs; enabled = current_enabled; run_proc; run_op; run_ptr; backtrack = IntSet.empty }
@@ -231,12 +232,26 @@ let rec explore func state clock last_access =
     let p = IntSet.min_elt s.enabled in
     let dones = ref IntSet.empty in
     s.backtrack <- IntSet.singleton p;
+    let is_backtracking = ref false in
     while IntSet.(cardinal (diff s.backtrack !dones)) > 0 do
       let j = IntSet.min_elt (IntSet.diff s.backtrack !dones) in
       dones := IntSet.add j !dones;
       let j_proc = List.nth s.procs j in
-      let schedule = (List.map (fun s -> (s.run_proc, s.run_op, s.run_ptr)) state) @ [(j, j_proc.op, j_proc.obj_ptr)] in
-      let statedash = state @ [do_run func schedule] in
+      let new_step = [(j, j_proc.op, j_proc.obj_ptr)] in
+      let full_schedule = List.map (fun s -> (s.run_proc, s.run_op, s.run_ptr)) state @ new_step in
+      let schedule =
+        if !is_backtracking
+        then begin
+          setup_run func full_schedule ;
+          full_schedule
+        end
+        else begin
+          is_backtracking := true ;
+          schedule_for_checks := full_schedule;
+          new_step
+        end
+      in
+      let statedash = state @ [do_run schedule] in
       let state_time = (List.length statedash)-1 in
       let new_last_access = match j_proc.obj_ptr with Some(ptr) -> IntMap.add ptr state_time last_access | None -> last_access in
       let new_clock = IntMap.add j state_time clock in
@@ -269,7 +284,10 @@ let check f =
 
 
 let trace func =
-  let empty_state = do_run func [(0, Start, None)] :: [] in
+  let schedule = [(0, Start, None)] in
+  setup_run func schedule ;
+  let empty_state = do_run schedule :: [] in
   let empty_clock = IntMap.empty in
   let empty_last_access = IntMap.empty in
-  explore func empty_state empty_clock empty_last_access
+  explore func empty_state empty_clock empty_last_access ;
+  Printf.printf "Finished after %i runs.\n%!" !num_runs

From 89a80aaa2be3674def4c708b673c4d00a39b487a Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Fri, 30 Sep 2022 02:15:13 +0200
Subject: [PATCH 02/23] list optims

---
 src/tracedAtomic.ml | 47 ++++++++++++++++++++-------------------------
 1 file changed, 21 insertions(+), 26 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 627e5c3..fbedb61 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -151,12 +151,6 @@ let spawn f =
   CCVector.push processes
     { next_op = Start; next_repr = None; resume_func = fiber_f; finished = false }
 
-let rec last_element l =
-  match l with
-  | h :: [] -> h
-  | [] -> assert(false)
-  | _ :: tl -> last_element tl
-
 type proc_rec = { proc_id: int; op: atomic_op; obj_ptr : int option }
 type state_cell = { procs: proc_rec list; run_proc: int; run_op: atomic_op; run_ptr: int option; enabled : IntSet.t; mutable backtrack : IntSet.t }
 
@@ -201,7 +195,7 @@ let do_run init_schedule =
       end
   in
   tracing := true;
-  run_trace init_schedule ();
+  run_trace (List.rev init_schedule) ();
   tracing := false;
   if !num_runs mod 1000 == 0 then
     Printf.printf "run: %d\n" !num_runs;
@@ -211,17 +205,16 @@ let do_run init_schedule =
                         |> Seq.filter (fun (_,proc) -> not proc.finished)
                         |> Seq.map (fun (id,_) -> id)
                         |> IntSet.of_seq in
-  match last_element init_schedule with
-  | (run_proc, run_op, run_ptr) ->
-    { procs; enabled = current_enabled; run_proc; run_op; run_ptr; backtrack = IntSet.empty }
+  let (run_proc, run_op, run_ptr) = List.hd init_schedule in
+  { procs; enabled = current_enabled; run_proc; run_op; run_ptr; backtrack = IntSet.empty }
 
-let rec explore func state clock last_access =
-  let s = last_element state in
+let rec explore func time state current_schedule clock last_access =
+  let s = List.hd state in
   List.iter (fun proc ->
       let j = proc.proc_id in
       let i = Option.bind proc.obj_ptr (fun ptr -> IntMap.find_opt ptr last_access) |> Option.value ~default:0 in
       if i != 0 then begin
-        let pre_s = List.nth state (i-1) in
+        let pre_s = List.nth state (time - i + 1) in
         if IntSet.mem j pre_s.enabled then
           pre_s.backtrack <- IntSet.add j pre_s.backtrack
         else
@@ -237,8 +230,8 @@ let rec explore func state clock last_access =
       let j = IntSet.min_elt (IntSet.diff s.backtrack !dones) in
       dones := IntSet.add j !dones;
       let j_proc = List.nth s.procs j in
-      let new_step = [(j, j_proc.op, j_proc.obj_ptr)] in
-      let full_schedule = List.map (fun s -> (s.run_proc, s.run_op, s.run_ptr)) state @ new_step in
+      let new_step = (j, j_proc.op, j_proc.obj_ptr) in
+      let full_schedule = new_step :: current_schedule in
       let schedule =
         if !is_backtracking
         then begin
@@ -248,14 +241,16 @@ let rec explore func state clock last_access =
         else begin
           is_backtracking := true ;
           schedule_for_checks := full_schedule;
-          new_step
+          [new_step]
         end
       in
-      let statedash = state @ [do_run schedule] in
-      let state_time = (List.length statedash)-1 in
-      let new_last_access = match j_proc.obj_ptr with Some(ptr) -> IntMap.add ptr state_time last_access | None -> last_access in
-      let new_clock = IntMap.add j state_time clock in
-      explore func statedash new_clock new_last_access
+      let s = do_run schedule in
+      let new_state = s :: state in
+      let new_schedule = (s.run_proc, s.run_op, s.run_ptr) :: current_schedule in
+      let new_time = time + 1 in
+      let new_last_access = match j_proc.obj_ptr with Some(ptr) -> IntMap.add ptr new_time last_access | None -> last_access in
+      let new_clock = IntMap.add j new_time clock in
+      explore func new_time new_state new_schedule new_clock new_last_access
     done
   end
 
@@ -277,17 +272,17 @@ let check f =
             Printf.printf "Process %d: %s %s\n" last_run_proc (atomic_op_str last_run_op) last_run_ptr
         end;
       end;
-    ) !schedule_for_checks;
+    ) (List.rev !schedule_for_checks) ;
       assert(false)
   end;
   tracing := tracing_at_start
 
 
 let trace func =
-  let schedule = [(0, Start, None)] in
-  setup_run func schedule ;
-  let empty_state = do_run schedule :: [] in
+  let empty_schedule = [(0, Start, None)] in
+  setup_run func empty_schedule ;
+  let empty_state = do_run empty_schedule :: [] in
   let empty_clock = IntMap.empty in
   let empty_last_access = IntMap.empty in
-  explore func empty_state empty_clock empty_last_access ;
+  explore func 1 empty_state empty_schedule empty_clock empty_last_access ;
   Printf.printf "Finished after %i runs.\n%!" !num_runs

From 235553191862279cbc9048999e95f83b700d3da4 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Fri, 30 Sep 2022 03:11:55 +0200
Subject: [PATCH 03/23] more precise backtracking

---
 src/tracedAtomic.ml | 56 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 47 insertions(+), 9 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index fbedb61..9f94461 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -208,18 +208,42 @@ let do_run init_schedule =
   let (run_proc, run_op, run_ptr) = List.hd init_schedule in
   { procs; enabled = current_enabled; run_proc; run_op; run_ptr; backtrack = IntSet.empty }
 
-let rec explore func time state current_schedule clock last_access =
+type category =
+  | Ignore
+  | Read
+  | Write
+  | Read_write
+
+let categorize = function
+  | Start | Make -> Ignore
+  | Get -> Read
+  | Set | Exchange -> Write
+  | CompareAndSwap | FetchAndAdd -> Read_write
+
+let rec explore func time state current_schedule clock (last_read, last_write) =
   let s = List.hd state in
   List.iter (fun proc ->
       let j = proc.proc_id in
-      let i = Option.bind proc.obj_ptr (fun ptr -> IntMap.find_opt ptr last_access) |> Option.value ~default:0 in
-      if i != 0 then begin
+      let find ptr map = match IntMap.find_opt ptr map with
+        | None -> None
+        | Some lst ->
+            List.find_opt (fun (_, proc_id) -> proc_id <> j) lst
+      in
+      let i = match categorize proc.op, proc.obj_ptr with
+        | Ignore, _ -> None
+        | Read, Some ptr -> find ptr last_write
+        | Write, Some ptr -> find ptr last_read
+        | Read_write, Some ptr -> max (find ptr last_read) (find ptr last_write)
+        | _ -> assert false
+      in
+      match i with
+      | None -> ()
+      | Some (i, _) ->
         let pre_s = List.nth state (time - i + 1) in
         if IntSet.mem j pre_s.enabled then
           pre_s.backtrack <- IntSet.add j pre_s.backtrack
         else
           pre_s.backtrack <- IntSet.union pre_s.backtrack pre_s.enabled
-      end
     ) s.procs;
   if IntSet.cardinal s.enabled > 0 then begin
     let p = IntSet.min_elt s.enabled in
@@ -244,11 +268,25 @@ let rec explore func time state current_schedule clock last_access =
           [new_step]
         end
       in
-      let s = do_run schedule in
-      let new_state = s :: state in
-      let new_schedule = (s.run_proc, s.run_op, s.run_ptr) :: current_schedule in
+      let step = do_run schedule in
+      let new_state = step :: state in
+      let new_schedule = (step.run_proc, step.run_op, step.run_ptr) :: current_schedule in
       let new_time = time + 1 in
-      let new_last_access = match j_proc.obj_ptr with Some(ptr) -> IntMap.add ptr new_time last_access | None -> last_access in
+      let add ptr map =
+        IntMap.update
+          ptr
+          (function None -> Some [time, step.run_proc]
+           | Some steps -> Some ((time, step.run_proc) :: steps))
+          map
+      in
+      let new_last_access =
+        match categorize step.run_op, step.run_ptr with
+        | Ignore, _ -> last_read, last_write
+        | Read, Some ptr -> add ptr last_read, last_write
+        | Write, Some ptr -> last_read, add ptr last_write
+        | Read_write, Some ptr -> add ptr last_read, add ptr last_write
+        | _ -> assert false
+      in
       let new_clock = IntMap.add j new_time clock in
       explore func new_time new_state new_schedule new_clock new_last_access
     done
@@ -283,6 +321,6 @@ let trace func =
   setup_run func empty_schedule ;
   let empty_state = do_run empty_schedule :: [] in
   let empty_clock = IntMap.empty in
-  let empty_last_access = IntMap.empty in
+  let empty_last_access = IntMap.empty, IntMap.empty in
   explore func 1 empty_state empty_schedule empty_clock empty_last_access ;
   Printf.printf "Finished after %i runs.\n%!" !num_runs

From 9222f07faa8db4d8ed51f6bea6360a09e43c8086 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Fri, 30 Sep 2022 11:53:46 +0200
Subject: [PATCH 04/23] fixup num_runs

---
 src/tracedAtomic.ml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 9f94461..c69ca64 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -164,8 +164,10 @@ let setup_run func init_schedule =
   CCVector.clear processes;
   schedule_for_checks := init_schedule;
   func () ;
+  finished_processes := 0;
   num_runs := !num_runs + 1;
-  finished_processes := 0
+  if !num_runs mod 1000 == 0 then
+    Printf.printf "run: %d\n" !num_runs
 
 let do_run init_schedule =
   (* cache the number of processes in case it's expensive*)
@@ -197,8 +199,6 @@ let do_run init_schedule =
   tracing := true;
   run_trace (List.rev init_schedule) ();
   tracing := false;
-  if !num_runs mod 1000 == 0 then
-    Printf.printf "run: %d\n" !num_runs;
   let procs = CCVector.mapi (fun i p -> { proc_id = i; op = p.next_op; obj_ptr = p.next_repr }) processes |> CCVector.to_list in
   let current_enabled = CCVector.to_seq processes
                         |> OSeq.zip_index

From 9a6f00c79d9ab0a99fcb90ab47d37cf8cf2e7e14 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Fri, 30 Sep 2022 11:57:21 +0200
Subject: [PATCH 05/23] fixup backtrack

---
 src/tracedAtomic.ml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index c69ca64..021d99c 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -275,8 +275,8 @@ let rec explore func time state current_schedule clock (last_read, last_write) =
       let add ptr map =
         IntMap.update
           ptr
-          (function None -> Some [time, step.run_proc]
-           | Some steps -> Some ((time, step.run_proc) :: steps))
+          (function None -> Some [new_time, step.run_proc]
+           | Some steps -> Some ((new_time, step.run_proc) :: steps))
           map
       in
       let new_last_access =

From 4036a90374b1857ced5f478514b024b9fd5e3b39 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Fri, 30 Sep 2022 11:59:59 +0200
Subject: [PATCH 06/23] test: only one bad apple

---
 tests/test_list.ml | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/tests/test_list.ml b/tests/test_list.ml
index 38bb645..adc77f3 100644
--- a/tests/test_list.ml
+++ b/tests/test_list.ml
@@ -4,17 +4,19 @@ module Atomic = Dscheck.TracedAtomic
 
 type conc_list = { value: int; next: conc_list option }
 
-let rec add_node list_head n =
+let rec add_node ~bug list_head n =
   (* try to add a new node to head *)
   let old_head = Atomic.get list_head in
   let new_node = { value = n ; next = (Some old_head) } in
     (* introduce bug *)
-    if Atomic.get list_head = old_head then begin
+    if bug && Atomic.get list_head = old_head then begin
       Atomic.set list_head new_node;
       true
     end
+    else if Atomic.compare_and_set list_head old_head new_node
+    then true
     else
-      add_node list_head n
+      add_node ~bug list_head n
 
 let check_node list_head n =
   let rec check_from_node node =
@@ -28,14 +30,14 @@ let check_node list_head n =
   (* try to find the node *)
     check_from_node (Atomic.get list_head)
 
-let add_and_check list_head n () =
-  assert(add_node list_head n);
+let add_and_check ~bug list_head n () =
+  assert(add_node ~bug list_head n);
   assert(check_node list_head n)
 
-let create_test upto () =
+let create_test ~buggy upto () =
   let list_head = Atomic.make { value = 0 ; next = None } in
   for x = 1 to upto do
-    Atomic.spawn (add_and_check list_head x);
+    Atomic.spawn (add_and_check ~bug:(x = buggy) list_head x);
   done;
   Atomic.final (fun () ->
     for x = 1 to upto do
@@ -44,4 +46,4 @@ let create_test upto () =
   )
 
 let () =
-  Atomic.trace (create_test 8)
+  Atomic.trace (create_test ~buggy:2 4)

From abdcc0b7bd8c52833b017a812ce04d2aa235a0bb Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Fri, 30 Sep 2022 12:32:20 +0200
Subject: [PATCH 07/23] backtracking: mark only last operation

---
 src/tracedAtomic.ml | 86 +++++++++++++++++++++++----------------------
 1 file changed, 44 insertions(+), 42 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 021d99c..7275356 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -152,7 +152,7 @@ let spawn f =
     { next_op = Start; next_repr = None; resume_func = fiber_f; finished = false }
 
 type proc_rec = { proc_id: int; op: atomic_op; obj_ptr : int option }
-type state_cell = { procs: proc_rec list; run_proc: int; run_op: atomic_op; run_ptr: int option; enabled : IntSet.t; mutable backtrack : IntSet.t }
+type state_cell = { procs: proc_rec list; run: proc_rec; enabled : IntSet.t; mutable backtrack : IntSet.t }
 
 let num_runs = ref 0
 
@@ -183,7 +183,7 @@ let do_run init_schedule =
         !final_func ();
         tracing := true
       end
-    | (process_id_to_run, next_op, next_ptr) :: schedule -> begin
+    | { proc_id = process_id_to_run ; op = next_op ; obj_ptr = next_ptr } :: schedule -> begin
         if !finished_processes == num_processes then
           (* this should never happen *)
           failwith("no enabled processes")
@@ -205,8 +205,8 @@ let do_run init_schedule =
                         |> Seq.filter (fun (_,proc) -> not proc.finished)
                         |> Seq.map (fun (id,_) -> id)
                         |> IntSet.of_seq in
-  let (run_proc, run_op, run_ptr) = List.hd init_schedule in
-  { procs; enabled = current_enabled; run_proc; run_op; run_ptr; backtrack = IntSet.empty }
+  let last_run = List.hd init_schedule in
+  { procs; enabled = current_enabled; run = last_run ; backtrack = IntSet.empty }
 
 type category =
   | Ignore
@@ -220,31 +220,32 @@ let categorize = function
   | Set | Exchange -> Write
   | CompareAndSwap | FetchAndAdd -> Read_write
 
+let mark_backtrack proc time state (last_read, last_write) =
+  let j = proc.proc_id in
+  let find ptr map = match IntMap.find_opt ptr map with
+    | None -> None
+    | Some lst ->
+        List.find_opt (fun (_, proc_id) -> proc_id <> j) lst
+  in
+  let i = match categorize proc.op, proc.obj_ptr with
+    | Ignore, _ -> None
+    | Read, Some ptr -> find ptr last_write
+    | Write, Some ptr -> find ptr last_read
+    | Read_write, Some ptr -> max (find ptr last_read) (find ptr last_write)
+    | _ -> assert false
+  in
+  match i with
+  | None -> ()
+  | Some (i, _) ->
+    assert (List.length state = time) ;
+    let pre_s = List.nth state (time - i) in
+    if IntSet.mem j pre_s.enabled then
+      pre_s.backtrack <- IntSet.add j pre_s.backtrack
+    else
+      pre_s.backtrack <- IntSet.union pre_s.backtrack pre_s.enabled
+
 let rec explore func time state current_schedule clock (last_read, last_write) =
   let s = List.hd state in
-  List.iter (fun proc ->
-      let j = proc.proc_id in
-      let find ptr map = match IntMap.find_opt ptr map with
-        | None -> None
-        | Some lst ->
-            List.find_opt (fun (_, proc_id) -> proc_id <> j) lst
-      in
-      let i = match categorize proc.op, proc.obj_ptr with
-        | Ignore, _ -> None
-        | Read, Some ptr -> find ptr last_write
-        | Write, Some ptr -> find ptr last_read
-        | Read_write, Some ptr -> max (find ptr last_read) (find ptr last_write)
-        | _ -> assert false
-      in
-      match i with
-      | None -> ()
-      | Some (i, _) ->
-        let pre_s = List.nth state (time - i + 1) in
-        if IntSet.mem j pre_s.enabled then
-          pre_s.backtrack <- IntSet.add j pre_s.backtrack
-        else
-          pre_s.backtrack <- IntSet.union pre_s.backtrack pre_s.enabled
-    ) s.procs;
   if IntSet.cardinal s.enabled > 0 then begin
     let p = IntSet.min_elt s.enabled in
     let dones = ref IntSet.empty in
@@ -254,7 +255,7 @@ let rec explore func time state current_schedule clock (last_read, last_write) =
       let j = IntSet.min_elt (IntSet.diff s.backtrack !dones) in
       dones := IntSet.add j !dones;
       let j_proc = List.nth s.procs j in
-      let new_step = (j, j_proc.op, j_proc.obj_ptr) in
+      let new_step = j_proc in
       let full_schedule = new_step :: current_schedule in
       let schedule =
         if !is_backtracking
@@ -269,18 +270,19 @@ let rec explore func time state current_schedule clock (last_read, last_write) =
         end
       in
       let step = do_run schedule in
+      mark_backtrack step.run time state (last_read, last_write);
       let new_state = step :: state in
-      let new_schedule = (step.run_proc, step.run_op, step.run_ptr) :: current_schedule in
+      let new_schedule = step.run :: current_schedule in
       let new_time = time + 1 in
       let add ptr map =
         IntMap.update
           ptr
-          (function None -> Some [new_time, step.run_proc]
-           | Some steps -> Some ((new_time, step.run_proc) :: steps))
+          (function None -> Some [time, step.run.proc_id]
+           | Some steps -> Some ((time, step.run.proc_id) :: steps))
           map
       in
       let new_last_access =
-        match categorize step.run_op, step.run_ptr with
+        match categorize step.run.op, step.run.obj_ptr with
         | Ignore, _ -> last_read, last_write
         | Read, Some ptr -> add ptr last_read, last_write
         | Write, Some ptr -> last_read, add ptr last_write
@@ -303,13 +305,9 @@ let check f =
   tracing := false;
   if not (f ()) then begin
     Printf.printf "Found assertion violation at run %d:\n" !num_runs;
-    List.iter (fun s ->
-      begin match s with
-      | (last_run_proc, last_run_op, last_run_ptr) -> begin
-          let last_run_ptr = Option.map string_of_int last_run_ptr |> Option.value ~default:"" in
-            Printf.printf "Process %d: %s %s\n" last_run_proc (atomic_op_str last_run_op) last_run_ptr
-        end;
-      end;
+    List.iter (fun { proc_id ; op ; obj_ptr } ->
+      let last_run_ptr = Option.map string_of_int obj_ptr |> Option.value ~default:"" in
+      Printf.printf "Process %d: %s %s\n" proc_id (atomic_op_str op) last_run_ptr
     ) (List.rev !schedule_for_checks) ;
       assert(false)
   end;
@@ -317,10 +315,14 @@ let check f =
 
 
 let trace func =
-  let empty_schedule = [(0, Start, None)] in
+  let empty_schedule = [{ proc_id = 0 ; op = Start ; obj_ptr = None }] in
   setup_run func empty_schedule ;
   let empty_state = do_run empty_schedule :: [] in
   let empty_clock = IntMap.empty in
   let empty_last_access = IntMap.empty, IntMap.empty in
-  explore func 1 empty_state empty_schedule empty_clock empty_last_access ;
-  Printf.printf "Finished after %i runs.\n%!" !num_runs
+  explore func 1 empty_state empty_schedule empty_clock empty_last_access
+
+let trace func =
+  Fun.protect
+    (fun () -> trace func)
+    ~finally:(fun () -> Printf.printf "Finished after %i runs.\n%!" !num_runs)

From 18e0bb0633c95a2102f22e667bd14f98bb191d5c Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 1 Oct 2022 11:51:51 +0200
Subject: [PATCH 08/23] big bad backtracking steps

---
 src/tracedAtomic.ml | 76 ++++++++++++++++++++++++++++-----------------
 1 file changed, 48 insertions(+), 28 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 7275356..4faf7b2 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -152,7 +152,12 @@ let spawn f =
     { next_op = Start; next_repr = None; resume_func = fiber_f; finished = false }
 
 type proc_rec = { proc_id: int; op: atomic_op; obj_ptr : int option }
-type state_cell = { procs: proc_rec list; run: proc_rec; enabled : IntSet.t; mutable backtrack : IntSet.t }
+type state_cell = {
+  procs : proc_rec list;
+  run : proc_rec;
+  enabled : IntSet.t;
+  mutable backtrack : proc_rec list IntMap.t;
+}
 
 let num_runs = ref 0
 
@@ -167,7 +172,7 @@ let setup_run func init_schedule =
   finished_processes := 0;
   num_runs := !num_runs + 1;
   if !num_runs mod 1000 == 0 then
-    Printf.printf "run: %d\n" !num_runs
+    Printf.printf "run: %d\n%!" !num_runs
 
 let do_run init_schedule =
   (* cache the number of processes in case it's expensive*)
@@ -206,7 +211,7 @@ let do_run init_schedule =
                         |> Seq.map (fun (id,_) -> id)
                         |> IntSet.of_seq in
   let last_run = List.hd init_schedule in
-  { procs; enabled = current_enabled; run = last_run ; backtrack = IntSet.empty }
+  { procs; enabled = current_enabled; run = last_run ; backtrack = IntMap.empty }
 
 type category =
   | Ignore
@@ -220,65 +225,80 @@ let categorize = function
   | Set | Exchange -> Write
   | CompareAndSwap | FetchAndAdd -> Read_write
 
-let mark_backtrack proc time state (last_read, last_write) =
+let mark_backtrack ~is_last proc time state (last_read, last_write) =
   let j = proc.proc_id in
   let find ptr map = match IntMap.find_opt ptr map with
     | None -> None
     | Some lst ->
         List.find_opt (fun (_, proc_id) -> proc_id <> j) lst
   in
-  let i = match categorize proc.op, proc.obj_ptr with
+  let find_loc ~is_last proc =
+    match categorize proc.op, proc.obj_ptr with
     | Ignore, _ -> None
     | Read, Some ptr -> find ptr last_write
-    | Write, Some ptr -> find ptr last_read
-    | Read_write, Some ptr -> max (find ptr last_read) (find ptr last_write)
+    | Write, Some ptr when not is_last -> find ptr last_read
+    | (Write | Read_write), Some ptr -> max (find ptr last_read) (find ptr last_write)
     | _ -> assert false
   in
-  match i with
+  match find_loc ~is_last proc with
   | None -> ()
   | Some (i, _) ->
     assert (List.length state = time) ;
-    let pre_s = List.nth state (time - i) in
-    if IntSet.mem j pre_s.enabled then
-      pre_s.backtrack <- IntSet.add j pre_s.backtrack
+    let pre_s = List.nth state (time - i + 1) in
+    if IntSet.mem j pre_s.enabled then begin
+      let replay_steps = List.filteri (fun k s -> k <= time - i && s.run.proc_id = j) state in
+      let replay_steps = List.map (fun s -> s.run) replay_steps in
+      let todo =
+        match IntMap.find_opt j pre_s.backtrack with
+        | None -> true
+        | Some lst -> List.length lst > List.length replay_steps
+      in
+      let causal p = match find_loc ~is_last:false p with None -> true | Some (k, _) -> k <= i in
+      if todo && List.for_all causal replay_steps
+      then pre_s.backtrack <- IntMap.add j replay_steps pre_s.backtrack
+    end
     else
-      pre_s.backtrack <- IntSet.union pre_s.backtrack pre_s.enabled
+      failwith "TODO: currently untested"
+
+let map_diff_set map set =
+  IntMap.filter (fun key _ -> not (IntSet.mem key set)) map
 
 let rec explore func time state current_schedule clock (last_read, last_write) =
   let s = List.hd state in
   if IntSet.cardinal s.enabled > 0 then begin
     let p = IntSet.min_elt s.enabled in
+    let init_step = List.nth s.procs p in
     let dones = ref IntSet.empty in
-    s.backtrack <- IntSet.singleton p;
+    s.backtrack <- IntMap.singleton p [init_step] ;
     let is_backtracking = ref false in
-    while IntSet.(cardinal (diff s.backtrack !dones)) > 0 do
-      let j = IntSet.min_elt (IntSet.diff s.backtrack !dones) in
+    while IntMap.(cardinal (map_diff_set s.backtrack !dones)) > 0 do
+      let j, new_step = IntMap.min_binding (map_diff_set s.backtrack !dones) in
       dones := IntSet.add j !dones;
-      let j_proc = List.nth s.procs j in
-      let new_step = j_proc in
-      let full_schedule = new_step :: current_schedule in
+      let new_schedule = List.rev_append (List.rev new_step) current_schedule in
       let schedule =
         if !is_backtracking
         then begin
-          setup_run func full_schedule ;
-          full_schedule
+          setup_run func new_schedule ;
+          new_schedule
         end
         else begin
           is_backtracking := true ;
-          schedule_for_checks := full_schedule;
-          [new_step]
+          schedule_for_checks := new_schedule;
+          new_step
         end
       in
       let step = do_run schedule in
-      mark_backtrack step.run time state (last_read, last_write);
-      let new_state = step :: state in
-      let new_schedule = step.run :: current_schedule in
-      let new_time = time + 1 in
+      let new_state =
+        List.map (fun run -> { step with run; backtrack = IntMap.empty }) new_step
+        @ state
+      in
+      let new_time = time + List.length new_step in
+      mark_backtrack ~is_last:(IntSet.is_empty step.enabled) step.run new_time new_state (last_read, last_write);
       let add ptr map =
         IntMap.update
           ptr
-          (function None -> Some [time, step.run.proc_id]
-           | Some steps -> Some ((time, step.run.proc_id) :: steps))
+          (function None -> Some [new_time, step.run.proc_id]
+           | Some steps -> Some ((new_time, step.run.proc_id) :: steps))
           map
       in
       let new_last_access =

From ecbff9b6e5eeb10e0bedf254effa67c50929057d Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 1 Oct 2022 11:59:28 +0200
Subject: [PATCH 09/23] add simple test

---
 tests/dune           |  5 +++++
 tests/test_simple.ml | 16 ++++++++++++++++
 2 files changed, 21 insertions(+)
 create mode 100644 tests/test_simple.ml

diff --git a/tests/dune b/tests/dune
index 2c7a92f..a3216e8 100644
--- a/tests/dune
+++ b/tests/dune
@@ -2,3 +2,8 @@
  (names test_list test_naive_counter)
  (modules test_list test_naive_counter)
  (libraries dscheck))
+
+(executable
+ (name test_simple)
+ (modules test_simple)
+ (libraries dscheck))
diff --git a/tests/test_simple.ml b/tests/test_simple.ml
new file mode 100644
index 0000000..7689270
--- /dev/null
+++ b/tests/test_simple.ml
@@ -0,0 +1,16 @@
+module Atomic = Dscheck.TracedAtomic
+
+(* Example from ocaml-parrafuzz presentation at the OCaml Workshop 2021:
+   https://www.youtube.com/watch?v=GZsUoSaIpIs *)
+
+let test i =
+  let x = Atomic.make i in
+  let y = Atomic.make 0 in
+  Atomic.spawn (fun () -> if Atomic.get x = 10 then Atomic.set y 2) ;
+  Atomic.spawn (fun () -> Atomic.set x 0 ; Atomic.set y 1) ;
+  Atomic.final (fun () -> Atomic.check (fun () -> Atomic.get y <> 2))
+
+let () =
+  Atomic.trace (fun () -> test 0) ;
+  Printf.printf "\n-----------------------------\n\n%!" ;
+  Atomic.trace (fun () -> test 10)

From 36bb3ec2f164b2a9c2856275b29223dafe823fd6 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 1 Oct 2022 15:26:47 +0200
Subject: [PATCH 10/23] fix: allow backtracking inside big steps

---
 src/tracedAtomic.ml | 131 +++++++++++++++++++++++++-------------------
 1 file changed, 74 insertions(+), 57 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 4faf7b2..47a098c 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -172,7 +172,7 @@ let setup_run func init_schedule =
   finished_processes := 0;
   num_runs := !num_runs + 1;
   if !num_runs mod 1000 == 0 then
-    Printf.printf "run: %d\n%!" !num_runs
+    Format.printf "run: %d@." !num_runs
 
 let do_run init_schedule =
   (* cache the number of processes in case it's expensive*)
@@ -211,7 +211,7 @@ let do_run init_schedule =
                         |> Seq.map (fun (id,_) -> id)
                         |> IntSet.of_seq in
   let last_run = List.hd init_schedule in
-  { procs; enabled = current_enabled; run = last_run ; backtrack = IntMap.empty }
+  { procs; enabled = current_enabled; run = last_run; backtrack = IntMap.empty }
 
 type category =
   | Ignore
@@ -253,9 +253,11 @@ let mark_backtrack ~is_last proc time state (last_read, last_write) =
         | None -> true
         | Some lst -> List.length lst > List.length replay_steps
       in
-      let causal p = match find_loc ~is_last:false p with None -> true | Some (k, _) -> k <= i in
+      let causal p = match find_loc ~is_last:false p with
+        | None -> true
+        | Some (k, _) -> k <= i in
       if todo && List.for_all causal replay_steps
-      then pre_s.backtrack <- IntMap.add j replay_steps pre_s.backtrack
+      then pre_s.backtrack <- IntMap.add j (List.rev replay_steps) pre_s.backtrack
     end
     else
       failwith "TODO: currently untested"
@@ -263,56 +265,70 @@ let mark_backtrack ~is_last proc time state (last_read, last_write) =
 let map_diff_set map set =
   IntMap.filter (fun key _ -> not (IntSet.mem key set)) map
 
-let rec explore func time state current_schedule clock (last_read, last_write) =
+let rec explore func time state (explored, state_planned) current_schedule clock (last_read, last_write) =
   let s = List.hd state in
-  if IntSet.cardinal s.enabled > 0 then begin
-    let p = IntSet.min_elt s.enabled in
-    let init_step = List.nth s.procs p in
-    let dones = ref IntSet.empty in
-    s.backtrack <- IntMap.singleton p [init_step] ;
-    let is_backtracking = ref false in
-    while IntMap.(cardinal (map_diff_set s.backtrack !dones)) > 0 do
-      let j, new_step = IntMap.min_binding (map_diff_set s.backtrack !dones) in
-      dones := IntSet.add j !dones;
-      let new_schedule = List.rev_append (List.rev new_step) current_schedule in
-      let schedule =
-        if !is_backtracking
-        then begin
-          setup_run func new_schedule ;
-          new_schedule
-        end
-        else begin
-          is_backtracking := true ;
-          schedule_for_checks := new_schedule;
-          new_step
-        end
-      in
-      let step = do_run schedule in
-      let new_state =
-        List.map (fun run -> { step with run; backtrack = IntMap.empty }) new_step
-        @ state
-      in
-      let new_time = time + List.length new_step in
-      mark_backtrack ~is_last:(IntSet.is_empty step.enabled) step.run new_time new_state (last_read, last_write);
-      let add ptr map =
-        IntMap.update
-          ptr
-          (function None -> Some [new_time, step.run.proc_id]
-           | Some steps -> Some ((new_time, step.run.proc_id) :: steps))
-          map
-      in
-      let new_last_access =
-        match categorize step.run.op, step.run.obj_ptr with
-        | Ignore, _ -> last_read, last_write
-        | Read, Some ptr -> add ptr last_read, last_write
-        | Write, Some ptr -> last_read, add ptr last_write
-        | Read_write, Some ptr -> add ptr last_read, add ptr last_write
-        | _ -> assert false
-      in
-      let new_clock = IntMap.add j new_time clock in
-      explore func new_time new_state new_schedule new_clock new_last_access
-    done
-  end
+  assert (IntMap.is_empty s.backtrack) ;
+  let dones = ref IntSet.empty in
+  begin match state_planned with
+  | [] ->
+    if IntSet.cardinal s.enabled > 0 then begin
+      let p = IntSet.min_elt s.enabled in
+      let init_step = List.nth s.procs p in
+      s.backtrack <- IntMap.singleton p [init_step] ;
+    end
+  | { proc_id ; _ } :: _ ->
+      dones := explored ;
+      s.backtrack <- IntMap.singleton proc_id state_planned
+  end ;
+  let is_backtracking = ref false in
+  while IntMap.(cardinal (map_diff_set s.backtrack !dones)) > 0 do
+    let j, new_steps = IntMap.min_binding (map_diff_set s.backtrack !dones) in
+    let new_explored =
+      if !is_backtracking || state_planned <> [] then !dones else IntSet.empty in
+    dones := IntSet.add j !dones;
+    let new_step, new_state_planned =
+      match new_steps with
+      | [] -> assert false
+      | next :: future -> next, future
+    in
+    let new_schedule = new_step :: current_schedule in
+    let schedule =
+      if !is_backtracking
+      then begin
+        setup_run func new_schedule ;
+        new_schedule
+      end
+      else begin
+        is_backtracking := true ;
+        schedule_for_checks := new_schedule;
+        [new_step]
+      end
+    in
+    let step = do_run schedule in
+    let new_state =
+      { step with run = new_step ; backtrack = IntMap.empty }
+      :: state
+    in
+    let new_time = time + 1 in
+    mark_backtrack ~is_last:(IntSet.is_empty step.enabled) step.run new_time new_state (last_read, last_write);
+    let add ptr map =
+      IntMap.update
+        ptr
+        (function None -> Some [new_time, step.run.proc_id]
+         | Some steps -> Some ((new_time, step.run.proc_id) :: steps))
+        map
+    in
+    let new_last_access =
+      match categorize step.run.op, step.run.obj_ptr with
+      | Ignore, _ -> last_read, last_write
+      | Read, Some ptr -> add ptr last_read, last_write
+      | Write, Some ptr -> last_read, add ptr last_write
+      | Read_write, Some ptr -> add ptr last_read, add ptr last_write
+      | _ -> assert false
+    in
+    let new_clock = IntMap.add j new_time clock in
+    explore func new_time new_state (new_explored, new_state_planned) new_schedule new_clock new_last_access
+  done
 
 let every f =
   every_func := f
@@ -324,10 +340,10 @@ let check f =
   let tracing_at_start = !tracing in
   tracing := false;
   if not (f ()) then begin
-    Printf.printf "Found assertion violation at run %d:\n" !num_runs;
+    Format.printf "@.@.Found assertion violation at run %d:@." !num_runs;
     List.iter (fun { proc_id ; op ; obj_ptr } ->
       let last_run_ptr = Option.map string_of_int obj_ptr |> Option.value ~default:"" in
-      Printf.printf "Process %d: %s %s\n" proc_id (atomic_op_str op) last_run_ptr
+      Format.printf "  Process %d: %s %s@." proc_id (atomic_op_str op) last_run_ptr
     ) (List.rev !schedule_for_checks) ;
       assert(false)
   end;
@@ -338,11 +354,12 @@ let trace func =
   let empty_schedule = [{ proc_id = 0 ; op = Start ; obj_ptr = None }] in
   setup_run func empty_schedule ;
   let empty_state = do_run empty_schedule :: [] in
+  let empty_state_planned = (IntSet.empty, []) in
   let empty_clock = IntMap.empty in
   let empty_last_access = IntMap.empty, IntMap.empty in
-  explore func 1 empty_state empty_schedule empty_clock empty_last_access
+  explore func 1 empty_state empty_state_planned empty_schedule empty_clock empty_last_access
 
 let trace func =
   Fun.protect
     (fun () -> trace func)
-    ~finally:(fun () -> Printf.printf "Finished after %i runs.\n%!" !num_runs)
+    ~finally:(fun () -> Format.printf "@.Finished after %i runs.@." !num_runs)

From 6b2195fdc347b8cdc63a96e00849b99c5b2ef7fb Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 1 Oct 2022 15:32:49 +0200
Subject: [PATCH 11/23] fix: print replay trace for all unhandled errors

---
 src/tracedAtomic.ml | 21 ++++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 47a098c..e652156 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -336,16 +336,18 @@ let every f =
 let final f =
   final_func := f
 
+let print_trace () =
+  List.iter (fun { proc_id ; op ; obj_ptr } ->
+    let last_run_ptr = Option.map string_of_int obj_ptr |> Option.value ~default:"" in
+    Format.printf "  Process %d: %s %s@." proc_id (atomic_op_str op) last_run_ptr
+  ) (List.rev !schedule_for_checks)
+
 let check f =
   let tracing_at_start = !tracing in
   tracing := false;
   if not (f ()) then begin
-    Format.printf "@.@.Found assertion violation at run %d:@." !num_runs;
-    List.iter (fun { proc_id ; op ; obj_ptr } ->
-      let last_run_ptr = Option.map string_of_int obj_ptr |> Option.value ~default:"" in
-      Format.printf "  Process %d: %s %s@." proc_id (atomic_op_str op) last_run_ptr
-    ) (List.rev !schedule_for_checks) ;
-      assert(false)
+    Format.printf "Check: found assertion violation!@." ;
+    assert false
   end;
   tracing := tracing_at_start
 
@@ -357,7 +359,12 @@ let trace func =
   let empty_state_planned = (IntSet.empty, []) in
   let empty_clock = IntMap.empty in
   let empty_last_access = IntMap.empty, IntMap.empty in
-  explore func 1 empty_state empty_state_planned empty_schedule empty_clock empty_last_access
+  try explore func 1 empty_state empty_state_planned empty_schedule empty_clock empty_last_access
+  with exn ->
+    Format.printf "Found error at run %d:@." !num_runs;
+    print_trace () ;
+    Format.printf "Unhandled exception: %s@." (Printexc.to_string exn) ;
+    Printexc.print_backtrace stdout
 
 let trace func =
   Fun.protect

From 708a2ab614d91e81fecac171e39697614928c680 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 1 Oct 2022 16:04:30 +0200
Subject: [PATCH 12/23] stable domains/atomics identifiers

---
 dscheck.opam         |   2 -
 dune-project         |   2 +-
 src/dune             |   3 +-
 src/tracedAtomic.ml  | 206 +++++++++++++++++++++++--------------------
 tests/test_simple.ml |   3 +-
 5 files changed, 112 insertions(+), 104 deletions(-)

diff --git a/dscheck.opam b/dscheck.opam
index c90f6fd..8ec0706 100644
--- a/dscheck.opam
+++ b/dscheck.opam
@@ -8,8 +8,6 @@ bug-reports: "https://github.com/ocaml-multicore/dscheck/issues"
 depends: [
   "ocaml" {>= "5.0.0"}
   "dune" {>= "2.9"}
-  "containers"
-  "oseq"
   "odoc" {with-doc}
 ]
 build: [
diff --git a/dune-project b/dune-project
index c39a362..1343b55 100644
--- a/dune-project
+++ b/dune-project
@@ -10,5 +10,5 @@
 (package
  (name dscheck)
  (synopsis "Traced Atomics")
- (depends (ocaml (>= 5.0.0)) dune containers oseq))
+ (depends (ocaml (>= 5.0.0)) dune))
 
diff --git a/src/dune b/src/dune
index 82e41aa..eef9f13 100644
--- a/src/dune
+++ b/src/dune
@@ -1,3 +1,2 @@
 (library
-  (public_name dscheck)
-  (libraries containers oseq))
+  (public_name dscheck))
diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index e652156..833551a 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -1,7 +1,16 @@
 open Effect
 open Effect.Shallow
 
-type 'a t = 'a Atomic.t * int
+module Uid = struct
+  type t = int list
+  let compare = List.compare Int.compare
+  let to_string t = String.concat "," (List.map string_of_int t)
+end
+
+module IdSet = Set.Make (Uid)
+module IdMap = Map.Make (Uid)
+
+type 'a t = 'a Atomic.t * Uid.t
 
 type _ Effect.t +=
   | Make : 'a -> 'a t Effect.t
@@ -10,24 +19,9 @@ type _ Effect.t +=
   | Exchange : ('a t * 'a) -> 'a Effect.t
   | CompareAndSwap : ('a t * 'a * 'a) -> bool Effect.t
   | FetchAndAdd : (int t * int) -> int Effect.t
+  | Spawn : (unit -> unit) -> unit Effect.t
 
-module IntSet = Set.Make(
-  struct
-    let compare = Stdlib.compare
-    type t = int
-  end )
-
-module IntMap = Map.Make(
-  struct
-    type t = int
-    let compare = Int.compare
-  end
-  )
-
-let _string_of_set s =
-  IntSet.fold (fun y x -> (string_of_int y) ^ "," ^ x) s ""
-
-type atomic_op = Start | Make | Get | Set | Exchange | CompareAndSwap | FetchAndAdd
+type atomic_op = Start | Make | Get | Set | Exchange | CompareAndSwap | FetchAndAdd | Spawn
 
 let atomic_op_str x =
   match x with
@@ -38,14 +32,18 @@ let atomic_op_str x =
   | Exchange -> "exchange"
   | CompareAndSwap -> "compare_and_swap"
   | FetchAndAdd -> "fetch_and_add"
+  | Spawn -> "spawn"
 
 let tracing = ref false
 
 let finished_processes = ref 0
 
 type process_data = {
+  uid : Uid.t;
+  mutable domain_generator : int;
+  mutable atomic_generator : int;
   mutable next_op: atomic_op;
-  mutable next_repr: int option;
+  mutable next_repr: Uid.t option;
   mutable resume_func : (unit, unit) handler -> unit;
   mutable finished : bool;
 }
@@ -54,13 +52,13 @@ let every_func = ref (fun () -> ())
 let final_func = ref (fun () -> ())
 
 (* Atomics implementation *)
-let atomics_counter = ref 1
+let atomics_counter = ref (-1)
 
 let make v = if !tracing then perform (Make v) else
     begin
       let i = !atomics_counter in
-      atomics_counter := !atomics_counter + 1;
-      (Atomic.make v, i)
+      atomics_counter := !atomics_counter - 1;
+      (Atomic.make v, [i])
     end
 
 let get r = if !tracing then perform (Get r) else match r with | (v,_) -> Atomic.get v
@@ -82,25 +80,30 @@ let incr r = ignore (fetch_and_add r 1)
 let decr r = ignore (fetch_and_add r (-1))
 
 (* Tracing infrastructure *)
-let processes = CCVector.create ()
+let processes = ref IdMap.empty
+let push_process p =
+  assert (not (IdMap.mem p.uid !processes)) ;
+  processes := IdMap.add p.uid p !processes
+
+let get_process id = IdMap.find id !processes
 
 let update_process_data process_id f op repr =
-  let process_rec = CCVector.get processes process_id in
+  let process_rec = get_process process_id in
   process_rec.resume_func <- f;
   process_rec.next_repr <- repr;
   process_rec.next_op <- op
 
 let finish_process process_id =
-  let process_rec = CCVector.get processes process_id in
+  let process_rec = get_process process_id in
   process_rec.finished <- true;
   finished_processes := !finished_processes + 1
 
-let handler current_process_id runner =
+let handler current_process runner =
   {
     retc =
       (fun _ ->
          (
-           finish_process current_process_id;
+           finish_process current_process.uid;
            runner ()));
     exnc = (fun s -> raise s);
     effc =
@@ -109,54 +112,62 @@ let handler current_process_id runner =
          | Make v ->
            Some
              (fun (k : (a, _) continuation) ->
-                let i = !atomics_counter in
+                let i = current_process.atomic_generator :: current_process.uid in
+                current_process.atomic_generator <- current_process.atomic_generator + 1 ;
                 let m = (Atomic.make v, i) in
-                atomics_counter := !atomics_counter + 1;
-                update_process_data current_process_id (fun h -> continue_with k m h) Make (Some i);
+                update_process_data current_process.uid (fun h -> continue_with k m h) Make (Some i);
                 runner ())
          | Get (v,i) ->
            Some
              (fun (k : (a, _) continuation) ->
-                update_process_data current_process_id (fun h -> continue_with k (Atomic.get v) h) Get (Some i);
+                update_process_data current_process.uid (fun h -> continue_with k (Atomic.get v) h) Get (Some i);
                 runner ())
          | Set ((r,i), v) ->
            Some
              (fun (k : (a, _) continuation) ->
-                update_process_data current_process_id (fun h -> continue_with k (Atomic.set r v) h) Set (Some i);
+                update_process_data current_process.uid (fun h -> continue_with k (Atomic.set r v) h) Set (Some i);
                 runner ())
          | Exchange ((a,i), b) ->
            Some
              (fun (k : (a, _) continuation) ->
-                update_process_data current_process_id (fun h -> continue_with k (Atomic.exchange a b) h) Exchange (Some i);
+                update_process_data current_process.uid (fun h -> continue_with k (Atomic.exchange a b) h) Exchange (Some i);
                 runner ())
          | CompareAndSwap ((x,i), s, v) ->
            Some
              (fun (k : (a, _) continuation) ->
-                update_process_data current_process_id (fun h ->
+                update_process_data current_process.uid (fun h ->
                     continue_with k (Atomic.compare_and_set x s v) h) CompareAndSwap (Some i);
                 runner ())
          | FetchAndAdd ((v,i), x) ->
            Some
              (fun (k : (a, _) continuation) ->
-                update_process_data current_process_id (fun h ->
+                update_process_data current_process.uid (fun h ->
                     continue_with k (Atomic.fetch_and_add v x) h) FetchAndAdd (Some i);
                 runner ())
+         | Spawn f ->
+            Some
+              (fun (k : (a, _) continuation) ->
+                 let fiber_f h = continue_with (fiber f) () h in
+                 let uid = current_process.domain_generator :: current_process.uid in
+                 current_process.domain_generator <- current_process.domain_generator + 1 ;
+                 push_process
+                   { uid; domain_generator = 0; atomic_generator = 0; next_op = Start; next_repr = None; resume_func = fiber_f; finished = false } ;
+                 update_process_data current_process.uid (fun h ->
+                     continue_with k () h) Spawn (Some uid);
+                 runner ())
          | _ ->
            None);
   }
 
 let spawn f =
-  let fiber_f h =
-    continue_with (fiber f) () h in
-  CCVector.push processes
-    { next_op = Start; next_repr = None; resume_func = fiber_f; finished = false }
+  if !tracing then perform (Spawn f) else failwith "spawn outside tracing"
 
-type proc_rec = { proc_id: int; op: atomic_op; obj_ptr : int option }
+type proc_rec = { proc_id: Uid.t; op: atomic_op; obj_ptr : Uid.t option }
 type state_cell = {
-  procs : proc_rec list;
+  procs : proc_rec IdMap.t;
   run : proc_rec;
-  enabled : IntSet.t;
-  mutable backtrack : proc_rec list IntMap.t;
+  enabled : IdSet.t;
+  mutable backtrack : proc_rec list IdMap.t;
 }
 
 let num_runs = ref 0
@@ -165,53 +176,51 @@ let num_runs = ref 0
 let schedule_for_checks = ref []
 
 let setup_run func init_schedule =
-  atomics_counter := 1;
-  CCVector.clear processes;
-  schedule_for_checks := init_schedule;
-  func () ;
+  processes := IdMap.empty ;
   finished_processes := 0;
+  schedule_for_checks := init_schedule;
+  tracing := true;
+  let uid = [] in
+  let fiber_f h = continue_with (fiber func) () h in
+  push_process
+    { uid; domain_generator = 0; atomic_generator = 0; next_op = Start; next_repr = None; resume_func = fiber_f; finished = false } ;
+  tracing := false;
   num_runs := !num_runs + 1;
   if !num_runs mod 1000 == 0 then
     Format.printf "run: %d@." !num_runs
 
 let do_run init_schedule =
-  (* cache the number of processes in case it's expensive*)
-  let num_processes = CCVector.length processes in
-  (* current number of ops we are through the current run *)
   let rec run_trace s () =
     tracing := false;
     !every_func ();
     tracing := true;
     match s with
-    | [] -> if !finished_processes == num_processes then begin
+    | [] -> if !finished_processes == IdMap.cardinal !processes then begin
         tracing := false;
         !final_func ();
         tracing := true
       end
-    | { proc_id = process_id_to_run ; op = next_op ; obj_ptr = next_ptr } :: schedule -> begin
-        if !finished_processes == num_processes then
-          (* this should never happen *)
-          failwith("no enabled processes")
-        else
-          begin
-            let process_to_run = CCVector.get processes process_id_to_run in
-            assert(process_to_run.next_op = next_op);
-            assert(process_to_run.next_repr = next_ptr);
-            process_to_run.resume_func (handler process_id_to_run (run_trace schedule))
-          end
-      end
+    | { proc_id = process_id_to_run ; op = next_op ; obj_ptr = next_ptr } :: schedule ->
+      let process_to_run = get_process process_id_to_run in
+      assert(not process_to_run.finished);
+      assert(process_to_run.next_op = next_op);
+      assert(process_to_run.next_repr = next_ptr);
+      process_to_run.resume_func (handler process_to_run (run_trace schedule))
   in
   tracing := true;
   run_trace (List.rev init_schedule) ();
   tracing := false;
-  let procs = CCVector.mapi (fun i p -> { proc_id = i; op = p.next_op; obj_ptr = p.next_repr }) processes |> CCVector.to_list in
-  let current_enabled = CCVector.to_seq processes
-                        |> OSeq.zip_index
-                        |> Seq.filter (fun (_,proc) -> not proc.finished)
-                        |> Seq.map (fun (id,_) -> id)
-                        |> IntSet.of_seq in
+  let procs =
+    IdMap.mapi (fun i p -> { proc_id = i; op = p.next_op; obj_ptr = p.next_repr }) !processes
+  in
+  let current_enabled =
+    IdMap.to_seq !processes
+    |> Seq.filter (fun (_, p) -> not p.finished)
+    |> Seq.map (fun (id, _) -> id)
+    |> IdSet.of_seq
+  in
   let last_run = List.hd init_schedule in
-  { procs; enabled = current_enabled; run = last_run; backtrack = IntMap.empty }
+  { procs; enabled = current_enabled; run = last_run; backtrack = IdMap.empty }
 
 type category =
   | Ignore
@@ -220,14 +229,14 @@ type category =
   | Read_write
 
 let categorize = function
-  | Start | Make -> Ignore
+  | Spawn | Start | Make -> Ignore
   | Get -> Read
   | Set | Exchange -> Write
   | CompareAndSwap | FetchAndAdd -> Read_write
 
 let mark_backtrack ~is_last proc time state (last_read, last_write) =
   let j = proc.proc_id in
-  let find ptr map = match IntMap.find_opt ptr map with
+  let find ptr map = match IdMap.find_opt ptr map with
     | None -> None
     | Some lst ->
         List.find_opt (fun (_, proc_id) -> proc_id <> j) lst
@@ -245,11 +254,11 @@ let mark_backtrack ~is_last proc time state (last_read, last_write) =
   | Some (i, _) ->
     assert (List.length state = time) ;
     let pre_s = List.nth state (time - i + 1) in
-    if IntSet.mem j pre_s.enabled then begin
+    if IdSet.mem j pre_s.enabled then begin
       let replay_steps = List.filteri (fun k s -> k <= time - i && s.run.proc_id = j) state in
       let replay_steps = List.map (fun s -> s.run) replay_steps in
       let todo =
-        match IntMap.find_opt j pre_s.backtrack with
+        match IdMap.find_opt j pre_s.backtrack with
         | None -> true
         | Some lst -> List.length lst > List.length replay_steps
       in
@@ -257,35 +266,35 @@ let mark_backtrack ~is_last proc time state (last_read, last_write) =
         | None -> true
         | Some (k, _) -> k <= i in
       if todo && List.for_all causal replay_steps
-      then pre_s.backtrack <- IntMap.add j (List.rev replay_steps) pre_s.backtrack
+      then pre_s.backtrack <- IdMap.add j (List.rev replay_steps) pre_s.backtrack
     end
     else
-      failwith "TODO: currently untested"
+      () (* failwith "TODO: currently untested" *)
 
 let map_diff_set map set =
-  IntMap.filter (fun key _ -> not (IntSet.mem key set)) map
+  IdMap.filter (fun key _ -> not (IdSet.mem key set)) map
 
 let rec explore func time state (explored, state_planned) current_schedule clock (last_read, last_write) =
   let s = List.hd state in
-  assert (IntMap.is_empty s.backtrack) ;
-  let dones = ref IntSet.empty in
+  assert (IdMap.is_empty s.backtrack) ;
+  let dones = ref IdSet.empty in
   begin match state_planned with
   | [] ->
-    if IntSet.cardinal s.enabled > 0 then begin
-      let p = IntSet.min_elt s.enabled in
-      let init_step = List.nth s.procs p in
-      s.backtrack <- IntMap.singleton p [init_step] ;
+    if IdSet.cardinal s.enabled > 0 then begin
+      let p = IdSet.min_elt s.enabled in
+      let init_step = IdMap.find p s.procs in
+      s.backtrack <- IdMap.singleton p [init_step] ;
     end
   | { proc_id ; _ } :: _ ->
       dones := explored ;
-      s.backtrack <- IntMap.singleton proc_id state_planned
+      s.backtrack <- IdMap.singleton proc_id state_planned
   end ;
   let is_backtracking = ref false in
-  while IntMap.(cardinal (map_diff_set s.backtrack !dones)) > 0 do
-    let j, new_steps = IntMap.min_binding (map_diff_set s.backtrack !dones) in
+  while IdMap.(cardinal (map_diff_set s.backtrack !dones)) > 0 do
+    let j, new_steps = IdMap.min_binding (map_diff_set s.backtrack !dones) in
     let new_explored =
-      if !is_backtracking || state_planned <> [] then !dones else IntSet.empty in
-    dones := IntSet.add j !dones;
+      if !is_backtracking || state_planned <> [] then !dones else IdSet.empty in
+    dones := IdSet.add j !dones;
     let new_step, new_state_planned =
       match new_steps with
       | [] -> assert false
@@ -306,13 +315,13 @@ let rec explore func time state (explored, state_planned) current_schedule clock
     in
     let step = do_run schedule in
     let new_state =
-      { step with run = new_step ; backtrack = IntMap.empty }
+      { step with run = new_step ; backtrack = IdMap.empty }
       :: state
     in
     let new_time = time + 1 in
-    mark_backtrack ~is_last:(IntSet.is_empty step.enabled) step.run new_time new_state (last_read, last_write);
+    mark_backtrack ~is_last:(IdSet.is_empty step.enabled) step.run new_time new_state (last_read, last_write);
     let add ptr map =
-      IntMap.update
+      IdMap.update
         ptr
         (function None -> Some [new_time, step.run.proc_id]
          | Some steps -> Some ((new_time, step.run.proc_id) :: steps))
@@ -326,7 +335,7 @@ let rec explore func time state (explored, state_planned) current_schedule clock
       | Read_write, Some ptr -> add ptr last_read, add ptr last_write
       | _ -> assert false
     in
-    let new_clock = IntMap.add j new_time clock in
+    let new_clock = IdMap.add j new_time clock in
     explore func new_time new_state (new_explored, new_state_planned) new_schedule new_clock new_last_access
   done
 
@@ -338,8 +347,8 @@ let final f =
 
 let print_trace () =
   List.iter (fun { proc_id ; op ; obj_ptr } ->
-    let last_run_ptr = Option.map string_of_int obj_ptr |> Option.value ~default:"" in
-    Format.printf "  Process %d: %s %s@." proc_id (atomic_op_str op) last_run_ptr
+    let last_run_ptr = Option.map Uid.to_string obj_ptr |> Option.value ~default:"" in
+    Format.printf "  Process %s: %s %s@." (Uid.to_string proc_id) (atomic_op_str op) last_run_ptr
   ) (List.rev !schedule_for_checks)
 
 let check f =
@@ -353,12 +362,13 @@ let check f =
 
 
 let trace func =
-  let empty_schedule = [{ proc_id = 0 ; op = Start ; obj_ptr = None }] in
+  num_runs := 0 ;
+  let empty_schedule = [{ proc_id = [] ; op = Start ; obj_ptr = None }] in
   setup_run func empty_schedule ;
   let empty_state = do_run empty_schedule :: [] in
-  let empty_state_planned = (IntSet.empty, []) in
-  let empty_clock = IntMap.empty in
-  let empty_last_access = IntMap.empty, IntMap.empty in
+  let empty_state_planned = (IdSet.empty, []) in
+  let empty_clock = IdMap.empty in
+  let empty_last_access = IdMap.empty, IdMap.empty in
   try explore func 1 empty_state empty_state_planned empty_schedule empty_clock empty_last_access
   with exn ->
     Format.printf "Found error at run %d:@." !num_runs;
diff --git a/tests/test_simple.ml b/tests/test_simple.ml
index 7689270..9b5dc4b 100644
--- a/tests/test_simple.ml
+++ b/tests/test_simple.ml
@@ -7,7 +7,8 @@ let test i =
   let x = Atomic.make i in
   let y = Atomic.make 0 in
   Atomic.spawn (fun () -> if Atomic.get x = 10 then Atomic.set y 2) ;
-  Atomic.spawn (fun () -> Atomic.set x 0 ; Atomic.set y 1) ;
+  Atomic.set x 0 ;
+  Atomic.set y 1 ;
   Atomic.final (fun () -> Atomic.check (fun () -> Atomic.get y <> 2))
 
 let () =

From a43282ff4c453b5f1aaedd8aff7c7c74c758772c Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 1 Oct 2022 21:14:33 +0200
Subject: [PATCH 13/23] fix: backtracking nested spawns

---
 src/tracedAtomic.ml  | 47 +++++++++++++++++++++++++++++++-------------
 tests/test_simple.ml |  3 ++-
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 833551a..a642293 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -234,6 +234,13 @@ let categorize = function
   | Set | Exchange -> Write
   | CompareAndSwap | FetchAndAdd -> Read_write
 
+let rec list_findi predicate lst i = match lst with
+  | [] -> None
+  | x::_ when predicate i x -> Some (i, x)
+  | _::xs -> list_findi predicate xs (i + 1)
+
+let list_findi predicate lst = list_findi predicate lst 0
+
 let mark_backtrack ~is_last proc time state (last_read, last_write) =
   let j = proc.proc_id in
   let find ptr map = match IdMap.find_opt ptr map with
@@ -249,27 +256,39 @@ let mark_backtrack ~is_last proc time state (last_read, last_write) =
     | (Write | Read_write), Some ptr -> max (find ptr last_read) (find ptr last_write)
     | _ -> assert false
   in
+  let rec find_replay_trace ~lower ~upper proc_id =
+    let pre_s = List.nth state (time - upper + 1) in
+    let replay_steps = List.filteri (fun k s -> k >= lower && k <= time - upper && s.run.proc_id = proc_id) state in
+    let replay_steps = List.rev_map (fun s -> s.run) replay_steps in
+    let causal p = match find_loc ~is_last:false p with
+      | None -> true
+      | Some (k, _) -> k <= upper in
+    if List.for_all causal replay_steps
+    then if IdSet.mem proc_id pre_s.enabled
+         then Some replay_steps
+         else let is_parent k s = k > lower && k < time - upper && s.run.op = Spawn && s.run.obj_ptr = Some proc_id in
+              match list_findi is_parent state with
+              | None -> None
+              | Some (parent_i, spawn) ->
+                  assert (parent_i > lower) ;
+                  begin match find_replay_trace ~lower:parent_i ~upper spawn.run.proc_id with
+                  | None -> None
+                  | Some spawn_steps -> Some (spawn_steps @ replay_steps)
+                  end
+    else None
+  in
   match find_loc ~is_last proc with
   | None -> ()
   | Some (i, _) ->
     assert (List.length state = time) ;
     let pre_s = List.nth state (time - i + 1) in
-    if IdSet.mem j pre_s.enabled then begin
-      let replay_steps = List.filteri (fun k s -> k <= time - i && s.run.proc_id = j) state in
-      let replay_steps = List.map (fun s -> s.run) replay_steps in
-      let todo =
-        match IdMap.find_opt j pre_s.backtrack with
+    match find_replay_trace ~lower:0 ~upper:i proc.proc_id with
+    | None -> ()
+    | Some replay_steps ->
+        if match IdMap.find_opt j pre_s.backtrack with
         | None -> true
         | Some lst -> List.length lst > List.length replay_steps
-      in
-      let causal p = match find_loc ~is_last:false p with
-        | None -> true
-        | Some (k, _) -> k <= i in
-      if todo && List.for_all causal replay_steps
-      then pre_s.backtrack <- IdMap.add j (List.rev replay_steps) pre_s.backtrack
-    end
-    else
-      () (* failwith "TODO: currently untested" *)
+        then pre_s.backtrack <- IdMap.add j replay_steps pre_s.backtrack
 
 let map_diff_set map set =
   IdMap.filter (fun key _ -> not (IdSet.mem key set)) map
diff --git a/tests/test_simple.ml b/tests/test_simple.ml
index 9b5dc4b..c2179aa 100644
--- a/tests/test_simple.ml
+++ b/tests/test_simple.ml
@@ -6,7 +6,8 @@ module Atomic = Dscheck.TracedAtomic
 let test i =
   let x = Atomic.make i in
   let y = Atomic.make 0 in
-  Atomic.spawn (fun () -> if Atomic.get x = 10 then Atomic.set y 2) ;
+  Atomic.spawn (fun () ->
+    Atomic.spawn (fun () -> if Atomic.get x = 10 then Atomic.set y 2)) ;
   Atomic.set x 0 ;
   Atomic.set y 1 ;
   Atomic.final (fun () -> Atomic.check (fun () -> Atomic.get y <> 2))

From 61482a126dffcb3d80dd70667bafb114e47115fa Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sun, 2 Oct 2022 12:57:01 +0200
Subject: [PATCH 14/23] fix: atomic ops categorization

---
 src/tracedAtomic.ml | 45 +++++++++++++++++++++++++++++++++------------
 1 file changed, 33 insertions(+), 12 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index a642293..5d42a33 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -21,7 +21,15 @@ type _ Effect.t +=
   | FetchAndAdd : (int t * int) -> int Effect.t
   | Spawn : (unit -> unit) -> unit Effect.t
 
-type atomic_op = Start | Make | Get | Set | Exchange | CompareAndSwap | FetchAndAdd | Spawn
+type cas_result = Unknown | Success | Failed
+
+type atomic_op =
+  | Start | Make | Get | Set | Exchange | FetchAndAdd | Spawn
+  | CompareAndSwap of cas_result ref
+
+let atomic_op_equal a b = match a, b with
+  | CompareAndSwap _, CompareAndSwap _ -> true
+  | _ -> a = b
 
 let atomic_op_str x =
   match x with
@@ -30,9 +38,14 @@ let atomic_op_str x =
   | Get -> "get"
   | Set -> "set"
   | Exchange -> "exchange"
-  | CompareAndSwap -> "compare_and_swap"
   | FetchAndAdd -> "fetch_and_add"
   | Spawn -> "spawn"
+  | CompareAndSwap cas ->
+      begin match !cas with
+      | Unknown -> "compare_and_swap?"
+      | Success -> "compare_and_swap"
+      | Failed -> "compare_and_no_swap"
+      end
 
 let tracing = ref false
 
@@ -135,8 +148,11 @@ let handler current_process runner =
          | CompareAndSwap ((x,i), s, v) ->
            Some
              (fun (k : (a, _) continuation) ->
+                let res = ref Unknown in
                 update_process_data current_process.uid (fun h ->
-                    continue_with k (Atomic.compare_and_set x s v) h) CompareAndSwap (Some i);
+                  let ok = Atomic.compare_and_set x s v in
+                  res := if ok then Success else Failed ;
+                  continue_with k ok h) (CompareAndSwap res) (Some i);
                 runner ())
          | FetchAndAdd ((v,i), x) ->
            Some
@@ -203,7 +219,7 @@ let do_run init_schedule =
     | { proc_id = process_id_to_run ; op = next_op ; obj_ptr = next_ptr } :: schedule ->
       let process_to_run = get_process process_id_to_run in
       assert(not process_to_run.finished);
-      assert(process_to_run.next_op = next_op);
+      assert(atomic_op_equal process_to_run.next_op next_op);
       assert(process_to_run.next_repr = next_ptr);
       process_to_run.resume_func (handler process_to_run (run_trace schedule))
   in
@@ -231,8 +247,14 @@ type category =
 let categorize = function
   | Spawn | Start | Make -> Ignore
   | Get -> Read
-  | Set | Exchange -> Write
-  | CompareAndSwap | FetchAndAdd -> Read_write
+  | Set -> Write
+  | Exchange | FetchAndAdd -> Read_write
+  | CompareAndSwap res ->
+      begin match !res with
+      | Unknown -> failwith "CAS unknown outcome" (* should not happen *)
+      | Success -> Read_write
+      | Failed -> Read_write
+      end
 
 let rec list_findi predicate lst i = match lst with
   | [] -> None
@@ -241,18 +263,17 @@ let rec list_findi predicate lst i = match lst with
 
 let list_findi predicate lst = list_findi predicate lst 0
 
-let mark_backtrack ~is_last proc time state (last_read, last_write) =
+let mark_backtrack proc time state (last_read, last_write) =
   let j = proc.proc_id in
   let find ptr map = match IdMap.find_opt ptr map with
     | None -> None
     | Some lst ->
         List.find_opt (fun (_, proc_id) -> proc_id <> j) lst
   in
-  let find_loc ~is_last proc =
+  let find_loc proc =
     match categorize proc.op, proc.obj_ptr with
     | Ignore, _ -> None
     | Read, Some ptr -> find ptr last_write
-    | Write, Some ptr when not is_last -> find ptr last_read
     | (Write | Read_write), Some ptr -> max (find ptr last_read) (find ptr last_write)
     | _ -> assert false
   in
@@ -260,7 +281,7 @@ let mark_backtrack ~is_last proc time state (last_read, last_write) =
     let pre_s = List.nth state (time - upper + 1) in
     let replay_steps = List.filteri (fun k s -> k >= lower && k <= time - upper && s.run.proc_id = proc_id) state in
     let replay_steps = List.rev_map (fun s -> s.run) replay_steps in
-    let causal p = match find_loc ~is_last:false p with
+    let causal p = match find_loc p with
       | None -> true
       | Some (k, _) -> k <= upper in
     if List.for_all causal replay_steps
@@ -277,7 +298,7 @@ let mark_backtrack ~is_last proc time state (last_read, last_write) =
                   end
     else None
   in
-  match find_loc ~is_last proc with
+  match find_loc proc with
   | None -> ()
   | Some (i, _) ->
     assert (List.length state = time) ;
@@ -338,7 +359,7 @@ let rec explore func time state (explored, state_planned) current_schedule clock
       :: state
     in
     let new_time = time + 1 in
-    mark_backtrack ~is_last:(IdSet.is_empty step.enabled) step.run new_time new_state (last_read, last_write);
+    mark_backtrack step.run new_time new_state (last_read, last_write);
     let add ptr map =
       IdMap.update
         ptr

From 36a06f1eacdde161693262739b3b8c9714d8fd41 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 19 Nov 2022 08:16:16 +0100
Subject: [PATCH 15/23] fix: atomic uid generation

---
 src/tracedAtomic.ml | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 5d42a33..dd8690c 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -65,12 +65,11 @@ let every_func = ref (fun () -> ())
 let final_func = ref (fun () -> ())
 
 (* Atomics implementation *)
-let atomics_counter = ref (-1)
+let atomics_counter = Atomic.make (-1)
 
 let make v = if !tracing then perform (Make v) else
     begin
-      let i = !atomics_counter in
-      atomics_counter := !atomics_counter - 1;
+      let i = Atomic.fetch_and_add atomics_counter (- 1) in
       (Atomic.make v, [i])
     end
 

From 16c866390df03776a4b244f0acd15e0b738749a2 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 19 Nov 2022 08:44:59 +0100
Subject: [PATCH 16/23] simplify dead code

---
 src/tracedAtomic.ml | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index dd8690c..0303250 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -310,10 +310,10 @@ let mark_backtrack proc time state (last_read, last_write) =
         | Some lst -> List.length lst > List.length replay_steps
         then pre_s.backtrack <- IdMap.add j replay_steps pre_s.backtrack
 
-let map_diff_set map set =
+let map_subtract_set map set =
   IdMap.filter (fun key _ -> not (IdSet.mem key set)) map
 
-let rec explore func time state (explored, state_planned) current_schedule clock (last_read, last_write) =
+let rec explore func time state (explored, state_planned) current_schedule (last_read, last_write) =
   let s = List.hd state in
   assert (IdMap.is_empty s.backtrack) ;
   let dones = ref IdSet.empty in
@@ -329,11 +329,10 @@ let rec explore func time state (explored, state_planned) current_schedule clock
       s.backtrack <- IdMap.singleton proc_id state_planned
   end ;
   let is_backtracking = ref false in
-  while IdMap.(cardinal (map_diff_set s.backtrack !dones)) > 0 do
-    let j, new_steps = IdMap.min_binding (map_diff_set s.backtrack !dones) in
-    let new_explored =
-      if !is_backtracking || state_planned <> [] then !dones else IdSet.empty in
-    dones := IdSet.add j !dones;
+  while IdMap.(cardinal (map_subtract_set s.backtrack !dones)) > 0 do
+    let j, new_steps = IdMap.min_binding (map_subtract_set s.backtrack !dones) in
+    let new_explored = !dones in
+    dones := IdSet.add j new_explored;
     let new_step, new_state_planned =
       match new_steps with
       | [] -> assert false
@@ -374,8 +373,7 @@ let rec explore func time state (explored, state_planned) current_schedule clock
       | Read_write, Some ptr -> add ptr last_read, add ptr last_write
       | _ -> assert false
     in
-    let new_clock = IdMap.add j new_time clock in
-    explore func new_time new_state (new_explored, new_state_planned) new_schedule new_clock new_last_access
+    explore func new_time new_state (new_explored, new_state_planned) new_schedule new_last_access
   done
 
 let every f =
@@ -406,9 +404,8 @@ let trace func =
   setup_run func empty_schedule ;
   let empty_state = do_run empty_schedule :: [] in
   let empty_state_planned = (IdSet.empty, []) in
-  let empty_clock = IdMap.empty in
   let empty_last_access = IdMap.empty, IdMap.empty in
-  try explore func 1 empty_state empty_state_planned empty_schedule empty_clock empty_last_access
+  try explore func 1 empty_state empty_state_planned empty_schedule empty_last_access
   with exn ->
     Format.printf "Found error at run %d:@." !num_runs;
     print_trace () ;

From 63aad4d6de431858d4f7a2248acf6e70fab94faa Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Mon, 21 Nov 2022 13:24:43 +0100
Subject: [PATCH 17/23] fix dune-project url

---
 dune-project | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dune-project b/dune-project
index 1343b55..09ee8d1 100644
--- a/dune-project
+++ b/dune-project
@@ -3,7 +3,7 @@
 (generate_opam_files true)
 
 (name dscheck)
-(source (github sadiqj/dscheck))
+(source (github ocaml-multicore/dscheck))
 (authors "Sadiq Jaffer")
 (maintainers "Sadiq Jaffer")
 

From ec3bfc1b9b009df83286d41153b2f43221d2417f Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 1 Oct 2022 16:17:05 +0200
Subject: [PATCH 18/23] round-robin scheduling

---
 src/tracedAtomic.ml | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 0303250..b2e59ff 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -313,14 +313,21 @@ let mark_backtrack proc time state (last_read, last_write) =
 let map_subtract_set map set =
   IdMap.filter (fun key _ -> not (IdSet.mem key set)) map
 
+let round_robin previous =
+  let id = previous.run.proc_id in
+  match IdSet.find_first_opt (fun j -> j > id) previous.enabled with
+  | None -> IdSet.min_elt_opt previous.enabled
+  | found -> found
+
 let rec explore func time state (explored, state_planned) current_schedule (last_read, last_write) =
   let s = List.hd state in
   assert (IdMap.is_empty s.backtrack) ;
   let dones = ref IdSet.empty in
   begin match state_planned with
   | [] ->
-    if IdSet.cardinal s.enabled > 0 then begin
-      let p = IdSet.min_elt s.enabled in
+    begin match round_robin s with
+    | None -> assert (IdSet.is_empty s.enabled)
+    | Some p ->
       let init_step = IdMap.find p s.procs in
       s.backtrack <- IdMap.singleton p [init_step] ;
     end

From bae4ce4d00eaf6dcb48e3990ab7d0a6383d976a9 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 1 Oct 2022 16:17:25 +0200
Subject: [PATCH 19/23] add spinlock test

---
 tests/dune             |  5 +++++
 tests/test_spinlock.ml | 10 ++++++++++
 2 files changed, 15 insertions(+)
 create mode 100644 tests/test_spinlock.ml

diff --git a/tests/dune b/tests/dune
index a3216e8..899b715 100644
--- a/tests/dune
+++ b/tests/dune
@@ -7,3 +7,8 @@
  (name test_simple)
  (modules test_simple)
  (libraries dscheck))
+
+(executable
+ (name test_spinlock)
+ (modules test_spinlock)
+ (libraries dscheck))
diff --git a/tests/test_spinlock.ml b/tests/test_spinlock.ml
new file mode 100644
index 0000000..c703d84
--- /dev/null
+++ b/tests/test_spinlock.ml
@@ -0,0 +1,10 @@
+module Atomic = Dscheck.TracedAtomic
+
+let test () =
+  let lock = Atomic.make true in
+  let result = Atomic.make 0 in
+  Atomic.spawn (fun () -> while Atomic.get lock do Atomic.incr result done) ;
+  Atomic.spawn (fun () -> Atomic.set lock false) ;
+  Atomic.final (fun () -> Atomic.check (fun () -> Atomic.get result < 10))
+
+let () = Atomic.trace test

From 238a7e6cf5d6b066f8f01d400aa30a0df39a3c32 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sun, 9 Oct 2022 22:15:02 +0200
Subject: [PATCH 20/23] use trie for backtracking fairness

---
 src/tracedAtomic.ml | 311 ++++++++++++++++++++++++++++--------------
 src/trie.ml         | 325 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 534 insertions(+), 102 deletions(-)
 create mode 100644 src/trie.ml

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index b2e59ff..db06b31 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -5,6 +5,19 @@ module Uid = struct
   type t = int list
   let compare = List.compare Int.compare
   let to_string t = String.concat "," (List.map string_of_int t)
+
+  let rec alpha i =
+    let head = Char.chr (Char.code 'A' + i mod 26) in
+    let i = i / 26 in
+    head :: if i > 0 then alpha i else []
+
+  let pretty lst =
+    match List.concat (List.map alpha lst) with
+    | [] -> "_"
+    | [x] -> String.make 1 x
+    | lst ->
+      let arr = Array.of_list lst in
+      "(" ^ String.init (Array.length arr) (Array.get arr) ^ ")"
 end
 
 module IdSet = Set.Make (Uid)
@@ -160,14 +173,14 @@ let handler current_process runner =
                     continue_with k (Atomic.fetch_and_add v x) h) FetchAndAdd (Some i);
                 runner ())
          | Spawn f ->
+            let uid = current_process.domain_generator :: current_process.uid in
+            current_process.domain_generator <- current_process.domain_generator + 1 ;
             Some
               (fun (k : (a, _) continuation) ->
-                 let fiber_f h = continue_with (fiber f) () h in
-                 let uid = current_process.domain_generator :: current_process.uid in
-                 current_process.domain_generator <- current_process.domain_generator + 1 ;
-                 push_process
-                   { uid; domain_generator = 0; atomic_generator = 0; next_op = Start; next_repr = None; resume_func = fiber_f; finished = false } ;
                  update_process_data current_process.uid (fun h ->
+                     let fiber_f h = continue_with (fiber f) () h in
+                     push_process
+                       { uid; domain_generator = 0; atomic_generator = 0; next_op = Start; next_repr = None; resume_func = fiber_f; finished = false } ;
                      continue_with k () h) Spawn (Some uid);
                  runner ())
          | _ ->
@@ -185,24 +198,82 @@ type state_cell = {
   mutable backtrack : proc_rec list IdMap.t;
 }
 
+module Step = struct
+  type t = proc_rec
+  let compare a b = Uid.compare a.proc_id b.proc_id
+  let to_string t =
+    Printf.sprintf "%s %s(%s)"
+      (Uid.pretty t.proc_id)
+      (atomic_op_str t.op)
+      (Option.value (Option.map Uid.to_string t.obj_ptr) ~default:"")
+end
+
+module T = Trie.Make (Step)
+
 let num_runs = ref 0
 
 (* we stash the current state in case a check fails and we need to log it *)
 let schedule_for_checks = ref []
 
-let setup_run func init_schedule =
+let max_depth = ref 0
+let depth = ref 0
+let incr_depth () =
+  depth := !depth + 1 ;
+  if !depth > !max_depth then max_depth := !depth
+
+let group_by fn = function
+  | [] -> []
+  | first :: rest ->
+      let rec go previous previouses = function
+        | [] -> [previous::previouses]
+        | x::xs when fn x previous -> go x (previous::previouses) xs
+        | x::xs -> (previous::previouses) :: go x [] xs
+      in
+      List.rev (go first [] rest)
+
+let print_header () =
+  Format.printf "%5s %6s %5s/%-5s %s@."
+    "run" "todo" "depth" "max" "latest trace"
+
+let print_stats trie =
+  Format.printf "%#4dk %#6d %#5d/%-#5d "
+    (!num_runs / 1000)
+    (T.nb_todos trie)
+    (Option.value (T.min_depth trie) ~default:0)
+    !max_depth ;
+  List.iter
+    (function
+      | steps when List.compare_length_with steps 3 <= 0 ->
+          List.iter (fun step -> Format.printf "%s" (Uid.pretty step.proc_id)) steps
+      | (step :: _) as steps ->
+          Format.printf "(%s%i)" (Uid.pretty step.proc_id) (List.length steps)
+      | _ -> assert false)
+    (group_by (fun a b -> a.proc_id = b.proc_id) !schedule_for_checks) ;
+  Format.printf "%!"
+
+let print_trace () =
+  List.iter (fun { proc_id ; op ; obj_ptr } ->
+    let last_run_ptr = Option.map Uid.to_string obj_ptr |> Option.value ~default:"" in
+    Format.printf "  Process %s: %s %s@." (Uid.pretty proc_id) (atomic_op_str op) last_run_ptr
+  ) (List.rev !schedule_for_checks)
+
+
+let setup_run func init_schedule trie =
   processes := IdMap.empty ;
   finished_processes := 0;
+  num_runs := !num_runs + 1;
+  if !num_runs mod 1000 == 0 then begin
+    Format.printf "%c[2K\r%!" (Char.chr 0x1b) ;
+    print_stats trie
+  end ;
   schedule_for_checks := init_schedule;
+  depth := 0 ;
   tracing := true;
   let uid = [] in
   let fiber_f h = continue_with (fiber func) () h in
   push_process
     { uid; domain_generator = 0; atomic_generator = 0; next_op = Start; next_repr = None; resume_func = fiber_f; finished = false } ;
-  tracing := false;
-  num_runs := !num_runs + 1;
-  if !num_runs mod 1000 == 0 then
-    Format.printf "run: %d@." !num_runs
+  tracing := false
 
 let do_run init_schedule =
   let rec run_trace s () =
@@ -216,6 +287,7 @@ let do_run init_schedule =
         tracing := true
       end
     | { proc_id = process_id_to_run ; op = next_op ; obj_ptr = next_ptr } :: schedule ->
+      incr_depth () ;
       let process_to_run = get_process process_id_to_run in
       assert(not process_to_run.finished);
       assert(atomic_op_equal process_to_run.next_op next_op);
@@ -276,8 +348,7 @@ let mark_backtrack proc time state (last_read, last_write) =
     | (Write | Read_write), Some ptr -> max (find ptr last_read) (find ptr last_write)
     | _ -> assert false
   in
-  let rec find_replay_trace ~lower ~upper proc_id =
-    let pre_s = List.nth state (time - upper + 1) in
+  let rec find_replay_trace ~pre_s ~lower ~upper proc_id =
     let replay_steps = List.filteri (fun k s -> k >= lower && k <= time - upper && s.run.proc_id = proc_id) state in
     let replay_steps = List.rev_map (fun s -> s.run) replay_steps in
     let causal p = match find_loc p with
@@ -291,7 +362,9 @@ let mark_backtrack proc time state (last_read, last_write) =
               | None -> None
               | Some (parent_i, spawn) ->
                   assert (parent_i > lower) ;
-                  begin match find_replay_trace ~lower:parent_i ~upper spawn.run.proc_id with
+                  if pre_s.run.proc_id = spawn.run.proc_id
+                  then None
+                  else begin match find_replay_trace ~pre_s ~lower:parent_i ~upper spawn.run.proc_id with
                   | None -> None
                   | Some spawn_steps -> Some (spawn_steps @ replay_steps)
                   end
@@ -300,88 +373,118 @@ let mark_backtrack proc time state (last_read, last_write) =
   match find_loc proc with
   | None -> ()
   | Some (i, _) ->
-    assert (List.length state = time) ;
+    assert (i < time) ;
     let pre_s = List.nth state (time - i + 1) in
-    match find_replay_trace ~lower:0 ~upper:i proc.proc_id with
+    match find_replay_trace ~pre_s ~lower:0 ~upper:i proc.proc_id with
     | None -> ()
-    | Some replay_steps ->
+    | Some [] -> failwith "empty replay steps"
+    | Some ((first_step :: _) as replay_steps) ->
+        let j = first_step.proc_id in
         if match IdMap.find_opt j pre_s.backtrack with
         | None -> true
-        | Some lst -> List.length lst > List.length replay_steps
+        | Some lst ->
+            assert (List.length lst <= List.length replay_steps) ;
+            false
         then pre_s.backtrack <- IdMap.add j replay_steps pre_s.backtrack
 
-let map_subtract_set map set =
-  IdMap.filter (fun key _ -> not (IdSet.mem key set)) map
-
 let round_robin previous =
-  let id = previous.run.proc_id in
-  match IdSet.find_first_opt (fun j -> j > id) previous.enabled with
-  | None -> IdSet.min_elt_opt previous.enabled
-  | found -> found
+  let urgent =
+    IdMap.fold
+      (fun proc_id step acc ->
+        assert (proc_id = step.proc_id) ;
+        if IdSet.mem proc_id previous.enabled
+        then match step.op with
+        | Start | Spawn | Make -> proc_id :: acc
+        | _ -> acc
+        else acc)
+    previous.procs
+    []
+  in
+  match urgent with
+  | [] ->
+      let id = previous.run.proc_id in
+      begin match IdSet.find_first_opt (fun j -> j > id) previous.enabled with
+      | None -> IdSet.min_elt_opt previous.enabled
+      | found -> found
+      end
+  | [single] -> Some single
+  | first :: _ -> Some first (* which one is best? *)
 
-let rec explore func time state (explored, state_planned) current_schedule (last_read, last_write) =
+let rec explore func time state trie current_schedule (last_read, last_write) =
   let s = List.hd state in
   assert (IdMap.is_empty s.backtrack) ;
-  let dones = ref IdSet.empty in
-  begin match state_planned with
-  | [] ->
-    begin match round_robin s with
-    | None -> assert (IdSet.is_empty s.enabled)
-    | Some p ->
-      let init_step = IdMap.find p s.procs in
-      s.backtrack <- IdMap.singleton p [init_step] ;
+  let existed, todo_next =
+    begin match T.next trie with
+    | None ->
+        begin match round_robin s with
+        | None ->
+            assert (IdSet.is_empty s.enabled) ;
+            false, None
+        | Some p ->
+            let init_step = IdMap.find p s.procs in
+            false, Some (init_step, T.todo)
+        end
+    | Some (step, new_trie) ->
+        true, Some (step, new_trie)
     end
-  | { proc_id ; _ } :: _ ->
-      dones := explored ;
-      s.backtrack <- IdMap.singleton proc_id state_planned
-  end ;
-  let is_backtracking = ref false in
-  while IdMap.(cardinal (map_subtract_set s.backtrack !dones)) > 0 do
-    let j, new_steps = IdMap.min_binding (map_subtract_set s.backtrack !dones) in
-    let new_explored = !dones in
-    dones := IdSet.add j new_explored;
-    let new_step, new_state_planned =
-      match new_steps with
-      | [] -> assert false
-      | next :: future -> next, future
-    in
-    let new_schedule = new_step :: current_schedule in
-    let schedule =
-      if !is_backtracking
-      then begin
-        setup_run func new_schedule ;
-        new_schedule
-      end
-      else begin
-        is_backtracking := true ;
+  in
+  match todo_next with
+  | None -> false, T.ok 1
+  | Some (new_step, new_trie) ->
+    begin try
+      let new_schedule = new_step :: current_schedule in
+      let schedule =
         schedule_for_checks := new_schedule;
         [new_step]
-      end
-    in
-    let step = do_run schedule in
-    let new_state =
-      { step with run = new_step ; backtrack = IdMap.empty }
-      :: state
-    in
-    let new_time = time + 1 in
-    mark_backtrack step.run new_time new_state (last_read, last_write);
-    let add ptr map =
-      IdMap.update
-        ptr
-        (function None -> Some [new_time, step.run.proc_id]
-         | Some steps -> Some ((new_time, step.run.proc_id) :: steps))
-        map
-    in
-    let new_last_access =
-      match categorize step.run.op, step.run.obj_ptr with
-      | Ignore, _ -> last_read, last_write
-      | Read, Some ptr -> add ptr last_read, last_write
-      | Write, Some ptr -> last_read, add ptr last_write
-      | Read_write, Some ptr -> add ptr last_read, add ptr last_write
-      | _ -> assert false
-    in
-    explore func new_time new_state (new_explored, new_state_planned) new_schedule new_last_access
-  done
+      in
+      let step = do_run schedule in
+      let new_state =
+        { step with run = new_step ; backtrack = IdMap.empty }
+        :: state
+      in
+      let new_time = time + 1 in
+      mark_backtrack step.run new_time new_state (last_read, last_write);
+      let add ptr map =
+        IdMap.update
+          ptr
+          (function None -> Some [new_time, step.run.proc_id]
+           | Some steps -> Some ((new_time, step.run.proc_id) :: steps))
+          map
+      in
+      let new_last_access =
+        match categorize step.run.op, step.run.obj_ptr with
+        | Ignore, _ -> last_read, last_write
+        | Read, Some ptr -> add ptr last_read, last_write
+        | Write, Some ptr -> last_read, add ptr last_write
+        | Read_write, Some ptr -> add ptr last_read, add ptr last_write
+        | _ -> assert false
+      in
+      let has_error, child =
+        explore func new_time new_state new_trie new_schedule new_last_access
+      in
+      let trie =
+        IdMap.fold
+          (fun _ backstep trie ->
+            T.insert_todos ~skip_edge:new_step backstep trie)
+          s.backtrack
+          trie
+      in
+      let trie = T.add new_step child trie in
+      let trie = T.simplify trie in
+      has_error, trie
+  with exn ->
+    let msg = Printexc.to_string exn in
+    let trie = T.add new_step (T.error msg) trie in
+    print_header () ;
+    print_stats trie ;
+
+    Format.printf "@.@.ERROR: %s@." msg ;
+    Printexc.print_backtrace stdout ;
+    print_trace () ;
+
+    assert (not existed) ;
+    true, trie
+  end
 
 let every f =
   every_func := f
@@ -389,37 +492,41 @@ let every f =
 let final f =
   final_func := f
 
-let print_trace () =
-  List.iter (fun { proc_id ; op ; obj_ptr } ->
-    let last_run_ptr = Option.map Uid.to_string obj_ptr |> Option.value ~default:"" in
-    Format.printf "  Process %s: %s %s@." (Uid.to_string proc_id) (atomic_op_str op) last_run_ptr
-  ) (List.rev !schedule_for_checks)
-
 let check f =
   let tracing_at_start = !tracing in
   tracing := false;
   if not (f ()) then begin
     Format.printf "Check: found assertion violation!@." ;
-    assert false
+    failwith "invalid check"
   end;
   tracing := tracing_at_start
 
+let error_count = ref 0
 
-let trace func =
-  num_runs := 0 ;
+let rec explore_all func trie =
   let empty_schedule = [{ proc_id = [] ; op = Start ; obj_ptr = None }] in
-  setup_run func empty_schedule ;
+  setup_run func empty_schedule trie ;
   let empty_state = do_run empty_schedule :: [] in
-  let empty_state_planned = (IdSet.empty, []) in
   let empty_last_access = IdMap.empty, IdMap.empty in
-  try explore func 1 empty_state empty_state_planned empty_schedule empty_last_access
-  with exn ->
-    Format.printf "Found error at run %d:@." !num_runs;
-    print_trace () ;
-    Format.printf "Unhandled exception: %s@." (Printexc.to_string exn) ;
-    Printexc.print_backtrace stdout
+  let has_error, trie =
+    explore func 1 empty_state trie empty_schedule empty_last_access
+  in
+  match T.min_depth trie with
+  | None ->
+      assert (T.nb_todos trie = 0) ;
+      T.graphviz ~filename:"/tmp/dscheck.dot" trie ;
+      trie
+  | _ when has_error ->
+      error_count := !error_count + 1 ;
+      T.graphviz ~filename:"/tmp/dscheck.dot" trie ;
+      if !error_count >= 5
+      then trie
+      else explore_all func trie
+  | Some _ -> explore_all func trie
 
 let trace func =
-  Fun.protect
-    (fun () -> trace func)
-    ~finally:(fun () -> Format.printf "@.Finished after %i runs.@." !num_runs)
+  print_header () ;
+  num_runs := 0 ;
+  let _ = explore_all func T.todo in
+  Format.printf "@.Found %i errors after %i runs.@." !error_count !num_runs ;
+  ()
diff --git a/src/trie.ml b/src/trie.ml
new file mode 100644
index 0000000..e0f9781
--- /dev/null
+++ b/src/trie.ml
@@ -0,0 +1,325 @@
+module type EDGE = sig
+  type t
+  val compare : t -> t -> int
+  val to_string : t -> string
+end
+
+module Make (Edge : EDGE) = struct
+
+  module M = Map.Make (Edge)
+
+  type edge = Edge.t
+
+  type t =
+    { error : bool
+    ; is_todo : bool
+    ; todo : int option
+    ; nb : int
+    ; nb_ok : int
+    ; s : s
+    }
+
+  and s =
+    | Ok
+    | Skip
+    | Error of string
+    | Branch of t M.t
+    | Todo
+
+  let min_depth t = t.todo
+  let nb_todos t = t.nb
+  let has_error t = t.error
+
+  let min_todo a b = match a, b with
+    | None, t | t, None -> t
+    | Some a, Some b -> Some (min a b)
+
+  let incr_todo = function None -> None | Some t -> Some (t + 1)
+
+  let ok nb_ok = { error = false ; is_todo = false ; todo = None ; nb = 0 ; nb_ok ; s = Ok }
+  let skip = { (ok 0) with s = Skip }
+  let error msg = { (ok 0) with error = true ; s = Error msg }
+  let todo = { error = false ; is_todo = true ; todo = Some 0 ; nb = 1 ; nb_ok = 0 ; s = Todo }
+
+  let branch children =
+    let error, is_todo, todo, nb, nb_ok =
+      M.fold
+        (fun _ child (acc_error, acc_is_todo, acc_todo, acc_nb, acc_nb_ok) ->
+          child.error || acc_error,
+          (child.s = Skip || child.is_todo) && acc_is_todo,
+          min_todo child.todo acc_todo,
+          acc_nb + child.nb,
+          acc_nb_ok + child.nb_ok)
+        children
+        (false, true, None, 0, 0)
+    in
+    { error ; is_todo ; todo = incr_todo todo ; nb ; nb_ok ; s = Branch children }
+
+  let simplify t =
+    match t.s with
+    | Branch children when M.for_all (fun _ c -> c.s = Ok) children ->
+        assert (t.nb_ok > 0) ;
+        ok t.nb_ok
+    | _ -> t
+
+  let add edge child t =
+    match t.s with
+    | Ok | Skip | Error _ -> assert false
+    | Todo -> branch (M.singleton edge child)
+    | Branch m -> branch (M.add edge child m)
+
+  let todolist ~skip_edge edges =
+    List.fold_right
+      (fun edge t -> branch (M.add skip_edge skip (M.singleton edge t)))
+      edges
+      todo
+
+  let rec insert_todos ~skip_edge todos t =
+    match t.s, todos with
+    | Todo, []
+    | Skip, _
+    | (Ok | Error _), _
+    | Branch _, [] ->
+        t
+    | Todo, todos -> todolist ~skip_edge todos
+    | Branch m, edge::todos ->
+        assert (Edge.compare skip_edge edge <> 0) ;
+        let m = match M.find_opt edge m with
+          | None ->
+              let t = todolist ~skip_edge todos in
+              M.add edge t m
+          | Some child ->
+              let child = insert_todos ~skip_edge todos child in
+              M.add edge child m
+        in
+        let m =
+          if M.mem skip_edge m
+          then m
+          else M.add skip_edge skip m
+        in
+        assert (M.mem skip_edge m) ;
+        branch m
+
+  let next t = match t.todo, t.s with
+    | None, _ -> None
+    | Some _, Todo -> None
+    | Some todo_distance, Branch m ->
+        let candidates =
+          M.fold
+            (fun edge child acc ->
+              match child.todo with
+              | Some distance (* when distance = todo_distance - 1 *) ->
+                  let weight = 1.0 /. float (100 * child.nb + distance - todo_distance + 1) in
+                  (weight, edge, child) :: acc
+              | _ -> acc)
+            m
+            []
+        in
+        let candidates =
+          if not t.error
+          then candidates
+          else match List.filter (fun (_, _, child) -> child.error || child.is_todo) candidates with
+          | [] -> candidates
+          | cs -> cs
+        in
+        begin match candidates with
+        | [(_, edge, child)] -> Some (edge, child)
+        | _ ->
+          let total = List.fold_left (fun acc (w, _, _) -> acc +. w) 0.0 candidates in
+          let selection = Random.float total in
+          let rec go selection = function
+            | [] -> failwith "Trie: candidate not found"
+            | [(_, edge, child)] -> Some (edge, child)
+            | (weight, edge, child) :: _ when selection <= weight -> Some (edge, child)
+            | (weight, _, _) :: rest -> go (selection -. weight) rest
+          in
+          go selection candidates
+        end
+    | _ -> assert false
+
+
+  let graphviz ~filename t =
+    let h = open_out filename in
+    let gen = ref 0 in
+    let fresh () =
+      let u = !gen in
+      gen := u + 1 ;
+      u
+    in
+    let refcounts = Hashtbl.create 16 in
+    let refcount uid =
+      match Hashtbl.find_opt refcounts uid with
+      | None -> 0
+      | Some count -> count
+    in
+    let incr_refcount uid =
+      let rc = refcount uid + 1 in
+      assert (rc >= 1) ;
+      Hashtbl.replace refcounts uid rc
+    in
+    let decr_refcount uid =
+      let rc = refcount uid - 1 in
+      assert (rc >= 1) ;
+      Hashtbl.replace refcounts uid rc
+    in
+    let cache = Hashtbl.create 16 in
+    let id_magic = Char.chr 0 in
+    let rec go ~print ~rc t =
+      let s = t.s in
+      let label =
+        match s with
+        | Ok -> "Ok"
+        | Skip -> "Skip"
+        | Todo -> "TODO"
+        | Error msg -> "ERROR: " ^ msg
+        | Branch _ -> ""
+      in
+      let has_error =
+        if t.error
+        then "style=filled fillcolor=black fontcolor=white"
+        else if t.todo = None
+        then "style=filled fillcolor=green"
+        else if t.is_todo
+        then "style=filled fillcolor=grey"
+        else ""
+      in
+      let style = if label = "" then "shape=point" else Printf.sprintf "label=%S" label in
+      let buf = Buffer.create 16 in
+      Printf.bprintf buf "%c [%s %s] ;\n" id_magic style has_error ;
+      let ids_children = ref [] in
+      begin match s with
+      | _ when not t.error -> ()
+      | Branch children ->
+          M.iter
+            (fun edge child ->
+              if child.s = Skip
+              then ()
+              else begin
+                let id_edge = edge_to_string ~print ~rc edge child in
+                let arrow_weight =
+                  if t.error && child.error
+                  then "[weight=1000 penwidth=4] "
+                  else ""
+                in
+                Printf.bprintf buf "%c -> e%i %s ;\n" id_magic id_edge arrow_weight ;
+                ids_children := id_edge :: !ids_children ;
+              end
+            )
+            children
+      | _ -> ()
+      end ;
+      let self = Buffer.contents buf in
+      match Hashtbl.find cache self with
+      | (uid, printed) ->
+          if not print
+          then (if rc then (incr_refcount uid ; List.iter decr_refcount !ids_children))
+          else if not printed
+          then begin
+            Hashtbl.replace cache self (uid, true) ;
+            let self = String.concat ("n" ^ string_of_int uid) (String.split_on_char id_magic self) in
+            Printf.fprintf h "%s" self ;
+          end ;
+          uid
+      | exception Not_found ->
+          let uid = fresh () in
+          Hashtbl.add cache self (uid, print) ;
+          if not print
+          then (if rc then incr_refcount uid)
+          else begin
+            let self = String.concat ("n" ^ string_of_int uid) (String.split_on_char id_magic self) in
+            Printf.fprintf h "%s" self ;
+          end ;
+          uid
+
+    and skip_forward acc t =
+      if refcount (go ~print:false ~rc:false t) > 1
+      then List.rev acc, t
+      else
+      let skip = match t.s with
+        | Branch children ->
+            M.fold
+              (fun edge child acc -> if child.s = Skip then acc else (edge, child) :: acc)
+              children
+              []
+        | _ -> []
+      in
+      match skip with
+      | [edge, single] -> skip_forward (edge :: acc) single
+      | _ -> (List.rev acc, t)
+
+    and edge_to_string ~print ~rc edge child =
+      let edges, child =
+        if print
+        then skip_forward [edge] child
+        else [edge], child
+      in
+      let buf = Buffer.create 16 in
+      let label = String.concat "\\l" (List.map Edge.to_string edges) in
+      let has_error =
+        if child.error
+        then "style=filled fillcolor=black fontcolor=white"
+        else if child.todo = None
+        then "style=filled fillcolor=lightgreen"
+        else if child.is_todo
+        then "style=filled fillcolor=grey"
+        else ""
+      in
+      let arrow_weight =
+        if child.error
+        then "[weight=1000 penwidth=4] "
+        else ""
+      in
+      Printf.bprintf buf "%c [shape=rectangle label=\"%s\\l\" %s] ;\n" id_magic label has_error ;
+      let id_child =
+        if not child.error
+        then begin
+          if child.nb > 1 then begin
+            Printf.bprintf buf "%c_todo [label=%S fillcolor=grey style=filled];\n"
+              id_magic
+              (Printf.sprintf "TODO %#i" child.nb) ;
+            Printf.bprintf buf "%c -> %c_todo;\n" id_magic id_magic ;
+          end ;
+          if child.nb_ok > 0 then begin
+            Printf.bprintf buf "%c_ok [label=%S fillcolor=lightgreen style=filled];\n"
+              id_magic
+              (Printf.sprintf "OK %#i" child.nb_ok) ;
+            Printf.bprintf buf "%c -> %c_ok;\n" id_magic id_magic ;
+          end ;
+          None
+        end
+        else begin
+          let id_child = go ~print ~rc child in
+          Printf.bprintf buf "%c -> n%i %s;\n" id_magic id_child arrow_weight ;
+          Some id_child
+        end
+      in
+      let self = Buffer.contents buf in
+      match Hashtbl.find cache self with
+      | (uid, printed) ->
+          if not print
+          then (if rc then ( incr_refcount uid ; Option.iter decr_refcount id_child ) )
+          else if not printed
+          then begin
+            Hashtbl.replace cache self (uid, true) ;
+            let self = String.concat ("e" ^ string_of_int uid) (String.split_on_char id_magic self) in
+            Printf.fprintf h "%s" self ;
+          end ;
+          uid
+      | exception Not_found ->
+          let uid = fresh () in
+          Hashtbl.add cache self (uid, print) ;
+          if not print
+          then (if rc then incr_refcount uid)
+          else begin
+            let self = String.concat ("e" ^ string_of_int uid) (String.split_on_char id_magic self) in
+            Printf.fprintf h "%s" self ;
+          end ;
+          uid
+    in
+    Printf.fprintf h "digraph dscheck {\n" ;
+    Printf.fprintf h "node [fontname=%S] ;\n" "Courier New" ;
+    let _ : int = go ~print:false ~rc:true t in
+    let _ : int = go ~print:true ~rc:false t in
+    Printf.fprintf h "}\n" ;
+    close_out h
+end

From 5375983e70d67f60ec583b41941526c8f0458e48 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sun, 16 Oct 2022 00:21:52 +0200
Subject: [PATCH 21/23] cut trie branches to minimize memory consumption

---
 src/tracedAtomic.ml | 17 ++++++++--------
 src/trie.ml         | 48 +++++++++++++++++++++++++--------------------
 2 files changed, 35 insertions(+), 30 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index db06b31..8b1d4af 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -413,20 +413,19 @@ let round_robin previous =
 let rec explore func time state trie current_schedule (last_read, last_write) =
   let s = List.hd state in
   assert (IdMap.is_empty s.backtrack) ;
-  let existed, todo_next =
-    begin match T.next trie with
+  let todo_next =
+    match T.next trie with
     | None ->
         begin match round_robin s with
         | None ->
             assert (IdSet.is_empty s.enabled) ;
-            false, None
+            None
         | Some p ->
             let init_step = IdMap.find p s.procs in
-            false, Some (init_step, T.todo)
+            Some (init_step, T.todo)
         end
     | Some (step, new_trie) ->
-        true, Some (step, new_trie)
-    end
+        Some (step, new_trie)
   in
   match todo_next with
   | None -> false, T.ok 1
@@ -443,7 +442,8 @@ let rec explore func time state trie current_schedule (last_read, last_write) =
         :: state
       in
       let new_time = time + 1 in
-      mark_backtrack step.run new_time new_state (last_read, last_write);
+      if T.nb_oks new_trie = 0 && not (T.has_error new_trie)
+      then mark_backtrack step.run new_time new_state (last_read, last_write);
       let add ptr map =
         IdMap.update
           ptr
@@ -482,7 +482,6 @@ let rec explore func time state trie current_schedule (last_read, last_write) =
     Printexc.print_backtrace stdout ;
     print_trace () ;
 
-    assert (not existed) ;
     true, trie
   end
 
@@ -528,5 +527,5 @@ let trace func =
   print_header () ;
   num_runs := 0 ;
   let _ = explore_all func T.todo in
-  Format.printf "@.Found %i errors after %i runs.@." !error_count !num_runs ;
+  Format.printf "@.Found %#i errors after %#i runs.@." !error_count !num_runs ;
   ()
diff --git a/src/trie.ml b/src/trie.ml
index e0f9781..24bfb34 100644
--- a/src/trie.ml
+++ b/src/trie.ml
@@ -22,12 +22,13 @@ module Make (Edge : EDGE) = struct
   and s =
     | Ok
     | Skip
+    | Todo
     | Error of string
     | Branch of t M.t
-    | Todo
 
   let min_depth t = t.todo
   let nb_todos t = t.nb
+  let nb_oks t = t.nb_ok
   let has_error t = t.error
 
   let min_todo a b = match a, b with
@@ -57,9 +58,22 @@ module Make (Edge : EDGE) = struct
 
   let simplify t =
     match t.s with
-    | Branch children when M.for_all (fun _ c -> c.s = Ok) children ->
-        assert (t.nb_ok > 0) ;
-        ok t.nb_ok
+    | _ when t.error || t.todo <> None -> t
+    | Branch children ->
+        let ok_children, has_skip = ref 0, ref false in
+        M.iter
+          (fun _ c -> match c.s with
+            | Ok -> incr ok_children
+            | Skip -> has_skip := true
+            | Branch _ ->
+                assert (c.todo = None) ;
+                assert (c.error = false) ;
+                assert (c.nb_ok > 0)
+            | _ -> assert false)
+          children ;
+        if !ok_children = 1 && !has_skip
+        then t
+        else ok t.nb_ok
     | _ -> t
 
   let add edge child t =
@@ -74,7 +88,7 @@ module Make (Edge : EDGE) = struct
       edges
       todo
 
-  let rec insert_todos ~skip_edge todos t =
+  let insert_todos ~skip_edge todos t =
     match t.s, todos with
     | Todo, []
     | Skip, _
@@ -83,22 +97,14 @@ module Make (Edge : EDGE) = struct
         t
     | Todo, todos -> todolist ~skip_edge todos
     | Branch m, edge::todos ->
-        assert (Edge.compare skip_edge edge <> 0) ;
-        let m = match M.find_opt edge m with
-          | None ->
-              let t = todolist ~skip_edge todos in
-              M.add edge t m
-          | Some child ->
-              let child = insert_todos ~skip_edge todos child in
-              M.add edge child m
-        in
-        let m =
-          if M.mem skip_edge m
-          then m
-          else M.add skip_edge skip m
-        in
-        assert (M.mem skip_edge m) ;
-        branch m
+        if M.mem edge m
+        then t
+        else begin
+          let t = todolist ~skip_edge todos in
+          let m = M.add edge t m in
+          assert (M.mem skip_edge m) ;
+          branch m
+        end
 
   let next t = match t.todo, t.s with
     | None, _ -> None

From 8cf65e81f4651669f16519284a304d3c53b1b979 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 19 Nov 2022 10:43:43 +0100
Subject: [PATCH 22/23] prettier trace

---
 src/tracedAtomic.ml | 50 ++++++++++++++++++++++++---------------------
 1 file changed, 27 insertions(+), 23 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 8b1d4af..95d7170 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -44,22 +44,6 @@ let atomic_op_equal a b = match a, b with
   | CompareAndSwap _, CompareAndSwap _ -> true
   | _ -> a = b
 
-let atomic_op_str x =
-  match x with
-  | Start -> "start"
-  | Make -> "make"
-  | Get -> "get"
-  | Set -> "set"
-  | Exchange -> "exchange"
-  | FetchAndAdd -> "fetch_and_add"
-  | Spawn -> "spawn"
-  | CompareAndSwap cas ->
-      begin match !cas with
-      | Unknown -> "compare_and_swap?"
-      | Success -> "compare_and_swap"
-      | Failed -> "compare_and_no_swap"
-      end
-
 let tracing = ref false
 
 let finished_processes = ref 0
@@ -201,11 +185,32 @@ type state_cell = {
 module Step = struct
   type t = proc_rec
   let compare a b = Uid.compare a.proc_id b.proc_id
+
+  let atomic_op_str op arg =
+    let arg = match op, arg with
+      | _, None -> "?"
+      | Spawn, Some domain_id -> Uid.pretty domain_id
+      | _, Some ptr -> Uid.to_string ptr
+    in
+    match op with
+    | Start -> "start"
+    | Spawn -> "spawn() = " ^ arg
+    | Make -> "make() = " ^ arg
+    | Get -> "get(" ^ arg ^ ")"
+    | Set -> "set(" ^ arg ^ ")"
+    | Exchange -> "exchange(" ^ arg ^ ")"
+    | FetchAndAdd -> "fetch_and_add(" ^ arg ^ ")"
+    | CompareAndSwap cas ->
+        begin match !cas with
+        | Unknown -> "compare_and_swap(" ^ arg ^ ")"
+        | Success -> "compare_and_swap(" ^ arg ^ ")"
+        | Failed -> "compare_and_no_swap(" ^ arg ^ ")"
+        end
+
   let to_string t =
-    Printf.sprintf "%s %s(%s)"
+    Printf.sprintf "%s %s"
       (Uid.pretty t.proc_id)
-      (atomic_op_str t.op)
-      (Option.value (Option.map Uid.to_string t.obj_ptr) ~default:"")
+      (atomic_op_str t.op t.obj_ptr)
 end
 
 module T = Trie.Make (Step)
@@ -252,10 +257,9 @@ let print_stats trie =
   Format.printf "%!"
 
 let print_trace () =
-  List.iter (fun { proc_id ; op ; obj_ptr } ->
-    let last_run_ptr = Option.map Uid.to_string obj_ptr |> Option.value ~default:"" in
-    Format.printf "  Process %s: %s %s@." (Uid.pretty proc_id) (atomic_op_str op) last_run_ptr
-  ) (List.rev !schedule_for_checks)
+  List.iter
+    (fun step -> Format.printf "  %s@." (Step.to_string step))
+    (List.rev !schedule_for_checks)
 
 
 let setup_run func init_schedule trie =

From 4be4164502b9e7142c623711beed4819f2480f11 Mon Sep 17 00:00:00 2001
From: ArthurW <arthur@tarides.com>
Date: Sat, 19 Nov 2022 11:12:12 +0100
Subject: [PATCH 23/23] expose trace options

---
 src/tracedAtomic.ml  | 41 ++++++++++++++++++++++-------------------
 src/tracedAtomic.mli | 11 ++++++++---
 2 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/src/tracedAtomic.ml b/src/tracedAtomic.ml
index 95d7170..e7e7e63 100644
--- a/src/tracedAtomic.ml
+++ b/src/tracedAtomic.ml
@@ -506,30 +506,33 @@ let check f =
 
 let error_count = ref 0
 
-let rec explore_all func trie =
+let explore_one func trie =
   let empty_schedule = [{ proc_id = [] ; op = Start ; obj_ptr = None }] in
   setup_run func empty_schedule trie ;
   let empty_state = do_run empty_schedule :: [] in
   let empty_last_access = IdMap.empty, IdMap.empty in
-  let has_error, trie =
-    explore func 1 empty_state trie empty_schedule empty_last_access
-  in
-  match T.min_depth trie with
-  | None ->
-      assert (T.nb_todos trie = 0) ;
-      T.graphviz ~filename:"/tmp/dscheck.dot" trie ;
-      trie
-  | _ when has_error ->
-      error_count := !error_count + 1 ;
-      T.graphviz ~filename:"/tmp/dscheck.dot" trie ;
-      if !error_count >= 5
-      then trie
-      else explore_all func trie
-  | Some _ -> explore_all func trie
-
-let trace func =
+  explore func 1 empty_state trie empty_schedule empty_last_access
+
+let rec explore_all ~count ~errors func trie =
+  if !error_count >= errors
+  || !num_runs >= count
+  || T.nb_todos trie = 0
+  then trie (* graphviz_output ?graphviz trie *)
+  else begin
+    let has_error, trie = explore_one func trie in
+    if has_error then error_count := !error_count + 1 ;
+    explore_all ~count ~errors func trie
+  end
+
+let graphviz_output ?graphviz trie =
+  match graphviz with
+  | None -> ()
+  | Some filename -> T.graphviz ~filename trie
+
+let trace ?(count = max_int) ?(errors = 1) ?graphviz func =
   print_header () ;
   num_runs := 0 ;
-  let _ = explore_all func T.todo in
+  let trie = explore_all ~count ~errors func T.todo in
+  graphviz_output ?graphviz trie ;
   Format.printf "@.Found %#i errors after %#i runs.@." !error_count !num_runs ;
   ()
diff --git a/src/tracedAtomic.mli b/src/tracedAtomic.mli
index 3f6ee63..c4a1c2e 100644
--- a/src/tracedAtomic.mli
+++ b/src/tracedAtomic.mli
@@ -46,8 +46,13 @@ val incr : int t -> unit
 val decr : int t -> unit
 (** [decr r] atomically decrements the value of [r] by [1]. *)
 
-val trace : (unit -> unit) -> unit
-(** start the simulation trace *)
+val trace : ?count:int -> ?errors:int -> ?graphviz:string -> (unit -> unit) -> unit
+(** [trace fn] starts the simulation trace on the function [fn].
+
+    - [?count] is the max number of traces to test (defaults to infinity)
+    - [?errors] is the max number of errors to find (defaults to 1)
+    - [?graphviz] is a filename to use for outputting a dot file (defaults to no output)
+ *)
 
 val spawn : (unit -> unit) -> unit
 (** spawn [f] as a new 'thread' *)
@@ -59,4 +64,4 @@ val final : (unit -> unit) -> unit
 (** run [f] after all processes complete *)
 
 val every : (unit -> unit) -> unit
-(** run [f] between every possible interleaving *)
\ No newline at end of file
+(** run [f] between every possible interleaving *)