Use polymorphic variants to avoid the impossible

Author: polytypic

src/michael_scott_queue.ml

@@ -18,32 +18,30 @@
 
 (* Michael-Scott queue *)
 
-type 'a node = Nil | Next of 'a * 'a node Atomic.t
-type 'a t = { head : 'a node Atomic.t; tail : 'a node Atomic.t Atomic.t }
+type 'a node = 'a * [ `Nil | `Node of 'a node ] Atomic.t
+
+type 'a t = {
+  head : [ `Node of 'a node ] Atomic.t;
+  tail : [ `Nil | `Node of 'a node ] Atomic.t Atomic.t;
+}
 
 let create () =
-  let tail = Atomic.make Nil in
-  let head = Next (Obj.magic (), tail) in
+  let tail = Atomic.make `Nil in
+  let head = `Node (Obj.magic (), tail) in
   { head = Atomic.make head; tail = Atomic.make tail }
 
-let is_empty q =
-  match Atomic.get q.head with
-  | Nil -> failwith "MSQueue.is_empty: impossible"
-  | Next (_, x) -> ( match Atomic.get x with Nil -> true | _ -> false)
+let is_empty { head; _ } =
+  let (`Node (_, x)) = Atomic.get head in
+  Atomic.get x == `Nil
 
-let pop q =
+let pop { head; _ } =
   let b = Backoff.create () in
   let rec loop () =
-    let s = Atomic.get q.head in
-    let nhead =
-      match s with
-      | Nil -> failwith "MSQueue.pop: impossible"
-      | Next (_, x) -> Atomic.get x
-    in
-    match nhead with
-    | Nil -> None
-    | Next (v, x) when Atomic.compare_and_set q.head s (Next (Obj.magic (), x))
-      ->
+    let (`Node (_, x) as old_head) = Atomic.get head in
+    match Atomic.get x with
+    | `Nil -> None
+    | `Node (v, x)
+      when Atomic.compare_and_set head old_head (`Node (Obj.magic (), x)) ->
         Some v
     | _ ->
         Backoff.once b;
@@ -54,40 +52,35 @@ let pop q =
 let rec fix_tail tail old_tail new_tail =
   if Atomic.compare_and_set tail old_tail new_tail then
     match Atomic.get new_tail with
-    | Nil -> ()
-    | Next (_, new_new_tail) -> fix_tail tail new_tail new_new_tail
+    | `Nil -> ()
+    | `Node (_, new_new_tail) -> fix_tail tail new_tail new_new_tail
 
-let push q v =
+let push { tail; _ } v =
   let rec find_tail_and_enq curr_end node =
-    if Atomic.compare_and_set curr_end Nil node then ()
+    if Atomic.compare_and_set curr_end `Nil node then ()
     else
       match Atomic.get curr_end with
-      | Nil -> find_tail_and_enq curr_end node
-      | Next (_, n) -> find_tail_and_enq n node
+      | `Nil -> find_tail_and_enq curr_end node
+      | `Node (_, n) -> find_tail_and_enq n node
   in
-  let new_tail = Atomic.make Nil in
-  let newnode = Next (v, new_tail) in
-  let old_tail = Atomic.get q.tail in
+  let new_tail = Atomic.make `Nil in
+  let newnode = `Node (v, new_tail) in
+  let old_tail = Atomic.get tail in
   find_tail_and_enq old_tail newnode;
-  if Atomic.compare_and_set q.tail old_tail new_tail then
+  if Atomic.compare_and_set tail old_tail new_tail then
     match Atomic.get new_tail with
-    | Nil -> ()
-    | Next (_, new_new_tail) -> fix_tail q.tail new_tail new_new_tail
+    | `Nil -> ()
+    | `Node (_, new_new_tail) -> fix_tail tail new_tail new_new_tail
 
-let clean_until q f =
+let clean_until { head; _ } f =
   let b = Backoff.create () in
   let rec loop () =
-    let s = Atomic.get q.head in
-    let nhead =
-      match s with
-      | Nil -> failwith "MSQueue.pop: impossible"
-      | Next (_, x) -> Atomic.get x
-    in
-    match nhead with
-    | Nil -> ()
-    | Next (v, x) ->
+    let (`Node (_, x) as old_head) = Atomic.get head in
+    match Atomic.get x with
+    | `Nil -> ()
+    | `Node (v, x) ->
         if not (f v) then
-          if Atomic.compare_and_set q.head s (Next (Obj.magic (), x)) then (
+          if Atomic.compare_and_set head old_head (`Node (Obj.magic (), x)) then (
             Backoff.reset b;
             loop ())
           else (
@@ -97,11 +90,10 @@ let clean_until q f =
   in
   loop ()
 
-type 'a cursor = 'a node
+type 'a cursor = [ `Nil | `Node of 'a node ]
 
-let snapshot q =
-  match Atomic.get q.head with
-  | Nil -> failwith "MSQueue.snapshot: impossible"
-  | Next (_, n) -> Atomic.get n
+let snapshot { head; _ } =
+  let (`Node (_, n)) = Atomic.get head in
+  Atomic.get n
 
-let next c = match c with Nil -> None | Next (a, n) -> Some (a, Atomic.get n)
+let next = function `Nil -> None | `Node (a, n) -> Some (a, Atomic.get n)