Add new functions and improve documentations of Treiber stack.

src/treiber_stack.ml

@@ -1,31 +1,82 @@
 (** Treiber's Lock Free stack *)
 
-type 'a node = Nil | Cons of { value : 'a; tail : 'a node }
-type 'a t = 'a node Atomic.t
+type 'a t = 'a list Atomic.t
 
-let create () = Atomic.make Nil |> Multicore_magic.copy_as_padded
-let is_empty t = Atomic.get t == Nil
+let create () = Atomic.make_contended []
+let is_empty t = Atomic.get t == []
+let of_list list = Atomic.make_contended list
+let of_seq seq = Atomic.make_contended (List.of_seq seq)
 
-let rec push t value backoff =
-  let tail = Atomic.get t in
-  let cons = Cons { value; tail } in
-  if not (Atomic.compare_and_set t tail cons) then
-    push t value (Backoff.once backoff)
-
-let push t value = push t value Backoff.default
+(* *)
 
 exception Empty
 
 type ('a, _) poly = Option : ('a, 'a option) poly | Value : ('a, 'a) poly
 
-let rec pop_as : type a r. a t -> Backoff.t -> (a, r) poly -> r =
+let peek_as : type a r. a t -> (a, r) poly -> r =
+ fun t poly ->
+  match Atomic.get t with
+  | [] -> begin match poly with Option -> None | Value -> raise Empty end
+  | hd :: _ -> ( match poly with Option -> Some hd | Value -> hd)
+
+let peek_exn t = peek_as t Value
+let peek_opt t = peek_as t Option
+
+type ('a, _) poly2 =
+  | Option : ('a, 'a option) poly2
+  | Value : ('a, 'a) poly2
+  | Unit : ('a, unit) poly2
+
+let rec pop_as : type a r. a t -> Backoff.t -> (a, r) poly2 -> r =
  fun t backoff poly ->
   match Atomic.get t with
-  | Nil -> begin match poly with Option -> None | Value -> raise Empty end
-  | Cons cons_r as cons ->
-      if Atomic.compare_and_set t cons cons_r.tail then
-        match poly with Option -> Some cons_r.value | Value -> cons_r.value
+  | [] -> begin
+      match poly with Option -> None | Value | Unit -> raise Empty
+    end
+  | hd :: tail as before ->
+      if Atomic.compare_and_set t before tail then
+        match poly with Option -> Some hd | Value -> hd | Unit -> ()
       else pop_as t (Backoff.once backoff) poly
 
-let pop t = pop_as t Backoff.default Value
+let pop_exn t = pop_as t Backoff.default Value
 let pop_opt t = pop_as t Backoff.default Option
+let drop_exn t = pop_as t Backoff.default Unit
+
+let rec pop_all t backoff =
+  match Atomic.get t with
+  | [] -> []
+  | old_head ->
+      if Atomic.compare_and_set t old_head [] then old_head
+      else pop_all t (Backoff.once backoff)
+
+let pop_all t = pop_all t Backoff.default
+
+let to_seq t =
+  match Atomic.get t with [] -> Seq.empty | old_head -> List.to_seq old_head
+(* *)
+
+let rec push t value backoff =
+  let before = Atomic.get t in
+  let after = value :: before in
+  if not (Atomic.compare_and_set t before after) then
+    push t value (Backoff.once backoff)
+
+let push t value = push t value Backoff.default
+
+(**)
+
+let rec push_all_ t backoff values =
+  match Atomic.get t with
+  | [] ->
+      if Atomic.compare_and_set t [] values then ()
+      else push_all_ t (Backoff.once backoff) values
+  | _ as old_head ->
+      if Atomic.compare_and_set t old_head (values @ old_head) then ()
+      else push_all_ t (Backoff.once backoff) values
+
+let push_all t values =
+  match values with
+  | [] -> ()
+  | _ -> push_all_ t Backoff.default (List.rev values)
+
+let add_seq t seq = push_all_ t Backoff.default (List.of_seq seq |> List.rev)

src/treiber_stack.mli

@@ -1,31 +1,153 @@
 (** Classic multi-producer multi-consumer Treiber stack.
 
-    All function are lockfree. It is the recommended starting point
-    when needing LIFO structure. *)
+  All functions are lock-free. It is the recommended starting point
+  when needing a LIFO structure. *)
+
+(** {1 API} *)
 
 type 'a t
-(** Type of Treiber stack holding items of type [t]. *)
+(** Represents a lock-free Treiber stack holding elements of type ['a]. *)
 
 val create : unit -> 'a t
-(** [create ()] returns a new and empty Treiber stack. *)
+(** [create ()] creates a new empty Treiber stack. *)
+
+val of_list : 'a list -> 'a t
+(** [of_list list] creates a new Treiber stack from a list. *)
 
 val is_empty : 'a t -> bool
-(** [is_empty s] checks whether stack [s] is empty. *)
+(** [is_empty stack] returns [true] if the [stack] is empty, otherwise [false]. *)
 
-val push : 'a t -> 'a -> unit
-(** [push s v] adds the element [v] at the top of stack [s]. *)
+(** {2 Consumer functions} *)
 
 exception Empty
-(** Raised when {!pop} is applied to an empty queue. *)
+(** Raised when {!pop_exn}, {!peek_exn}, {!drop_exn}, or {!set_exn} is 
+  applied to an empty stack.
+*)
 
-val pop : 'a t -> 'a
-(** [pop s] removes and returns the topmost element in the
-    stack [s].
+val peek_exn : 'a t -> 'a
+(** [peek_exn stack] returns the top element of the [stack] without removing it.
+
+  @raises Empty if the [stack] is empty. *)
 
-    @raise Empty if [a] is empty.
-*)
+val peek_opt : 'a t -> 'a option
+(** [peek_opt stack] returns [Some] of the top element of the [stack] without
+  removing it, or [None] if the [stack] is empty. *)
+
+val pop_exn : 'a t -> 'a
+(** [pop_exn stack] removes and returns the top element of the [stack].
+
+  @raises Empty if the [stack] is empty. *)
 
 val pop_opt : 'a t -> 'a option
-(** [pop_opt s] removes and returns the topmost element in the
-    stack [s], or returns [None] if the stack is empty.
+(** [pop_opt stack] removes and returns [Some] of the top element of the [stack],
+  or [None] if the [stack] is empty. *)
+
+val drop_exn : 'a t -> unit
+(** [drop_exn stack] removes the top element of the [stack]. 
+
+  @raises Empty if the [stack] is empty. *)
+
+val pop_all : 'a t -> 'a list
+(** [pop_all stack] removes and returns all elements of the [stack] in LIFO 
+order. 
+
+  {[
+  # open Saturn_lockfree.Stack
+  # let t : int t = create ()
+  val t : int t = <abstr>
+  # push t 1
+  - : unit = ()
+  # push t 2
+  - : unit = ()
+  # push t 3
+  - : unit = ()
+  # pop_all t
+  - : int list = [3; 2; 1]
+  ]}
 *)
+
+(** {2 Producer functions} *)
+
+val push : 'a t -> 'a -> unit
+(** [push stack element] adds [element] to the top of the [stack]. *)
+
+val push_all : 'a t -> 'a list -> unit
+(** [push_all stack elements] adds all [elements] to the top of the [stack]. 
+
+  {[
+  # let t : int t = create ()
+  val t : int t = <abstr>
+  # push_all t [1; 2; 3; 4]
+  - : unit = ()
+  # pop_opt t
+  - : int option = Some 4
+  # pop_opt t 
+  - : int option = Some 3
+  # pop_all t
+  - : int list = [2; 1]
+  ]}
+  *)
+
+(** {2 With Sequences }*)
+val to_seq : 'a t -> 'a Seq.t
+(** [to_seq stack] takes a snapshot of [stack] and returns its value top to 
+bottom.
+
+  🐌 This is a linear time operation. *)
+
+val of_seq : 'a Seq.t -> 'a t
+(** [of_seq seq] creates a stack from a [seq]. It must be finite. *)
+
+val add_seq : 'a t -> 'a Seq.t -> unit
+(** [add_seq stack seq] adds all elements of [seq] to the top of the 
+[stack]. [seq] must be finite. *)
+
+(** {1 Examples}
+  An example top-level session:
+  {[
+    # open Saturn_lockfree.Stack
+    # let t : int t = create ()
+    val t : int t = <abstr>
+    # push t 42
+    - : unit = ()
+    # push_all t [1; 2; 3]
+    - : unit = ()
+    # pop_exn t
+    - : int = 3
+    # peek_opt t
+    - : int option = Some 2
+    # pop_all t
+    - : int list = [2; 1; 42]
+    # pop_exn t
+    Exception: Saturn_lockfree__Treiber_stack.Empty.]}
+
+  A multicore example: 
+  {@ocaml non-deterministic[
+    # open Saturn_lockfree.Stack
+    # let t : int t = create ()
+    val t : int t = <abstr>
+    # let barrier =  Atomic.make 2
+    val barrier : int Atomic.t = <abstr>
+    # let pusher () = 
+        Atomic.decr barrier;
+        while Atomic.get barrier != 0 do Domain.cpu_relax () done;
+        push_all t [1;2;3] |> ignore;
+        push t 42;
+        push t 12
+    val pusher : unit -> unit = <fun>
+    # let popper () = 
+        Atomic.decr barrier;
+        while Atomic.get barrier != 0 do Domain.cpu_relax () done;
+        List.init 6 (fun i -> Domain.cpu_relax (); pop_opt t)
+    val popper : unit -> int option list = <fun>
+    # let domain_pusher = Domain.spawn pusher
+    val domain_pusher : unit Domain.t = <abstr>
+    # let domain_popper = Domain.spawn popper
+    val domain_popper : int option list Domain.t = <abstr>
+    # Domain.join domain_pusher
+    - : unit = ()
+    # Domain.join domain_popper
+    - : int option list = [Some 42; Some 3; Some 2; Some 1; None; Some 12]
+    ]}
+
+  *)

test/treiber_stack/stm_treiber_stack.ml

@@ -1,56 +1,87 @@
-(** Sequential and Parallel model-based tests of treiber_stack *)
+(** Sequential and Parallel model-based tests of bounded_queue *)
 
 open QCheck
 open STM
-module Treiber_stack = Saturn.Stack
+module Stack = Saturn.Stack
 
 module Spec = struct
-  type cmd = Push of int | Pop | Is_empty
+  type cmd =
+    | Push of int
+    | Push_all of int list
+    | Pop_opt
+    (* peek_exn and drop_exn use the same function as pop_exn*)
+    | Pop_all
+    | Peek_opt (* peek_exn uses the same function as peek_exn *)
+    | To_seq
+    | Is_empty
+
+  let string_of_int_list l =
+    "[" ^ String.concat "; " (List.map string_of_int l) ^ "]"
 
   let show_cmd c =
     match c with
     | Push i -> "Push " ^ string_of_int i
-    | Pop -> "Pop"
+    | Push_all l -> "Push_all " ^ string_of_int_list l
+    (* | Push_all_exn l -> "Push_all_exn " ^ string_of_int_list l *)
+    | Pop_opt -> "Pop_opt"
+    | Pop_all -> "Pop_all"
+    | Peek_opt -> "Peek_opt"
+    | To_seq -> "To_seq"
     | Is_empty -> "Is_empty"
 
   type state = int list
-  type sut = int Treiber_stack.t
+  type sut = int Stack.t
 
   let arb_cmd _s =
     let int_gen = Gen.nat in
     QCheck.make ~print:show_cmd
       (Gen.oneof
          [
            Gen.map (fun i -> Push i) int_gen;
-           Gen.return Pop;
+           Gen.map (fun l -> Push_all l) (Gen.list int_gen);
+           Gen.return Pop_opt;
+           Gen.return Pop_all;
+           Gen.return Peek_opt;
+           Gen.return To_seq;
            Gen.return Is_empty;
          ])
 
   let init_state = []
-  let init_sut () = Treiber_stack.create ()
+  let init_sut () = Stack.create ()
   let cleanup _ = ()
 
   let next_state c s =
     match c with
     | Push i -> i :: s
-    | Pop -> ( match s with [] -> s | _ :: s' -> s')
+    | Push_all l -> List.rev l @ s
+    (* | Push_all_exn l ->
+        if List.length s + List.length l > capacity then s else List.rev l @ s *)
+    | Pop_opt -> ( match s with [] -> s | _ :: s' -> s')
+    | Pop_all -> []
+    | Peek_opt -> s
+    | To_seq -> s
     | Is_empty -> s
 
   let precond _ _ = true
 
   let run c d =
     match c with
-    | Push i -> Res (unit, Treiber_stack.push d i)
-    | Pop -> Res (result int exn, protect Treiber_stack.pop d)
-    | Is_empty -> Res (bool, Treiber_stack.is_empty d)
+    | Push i -> Res (unit, Stack.push d i)
+    | Push_all l -> Res (unit, Stack.push_all d l)
+    | Pop_opt -> Res (option int, Stack.pop_opt d)
+    | Pop_all -> Res (list int, Stack.pop_all d)
+    | Peek_opt -> Res (option int, Stack.peek_opt d)
+    | To_seq -> Res (seq int, Stack.to_seq d)
+    | Is_empty -> Res (bool, Stack.is_empty d)
 
   let postcond c (s : state) res =
     match (c, res) with
-    | Push _, Res ((Unit, _), _) -> true
-    | Pop, Res ((Result (Int, Exn), _), res) -> (
-        match s with
-        | [] -> res = Error Treiber_stack.Empty
-        | j :: _ -> res = Ok j)
+    | Push _, Res ((Unit, _), ()) -> true
+    | Push_all _, Res ((Unit, _), _) -> true
+    | (Pop_opt | Peek_opt), Res ((Option Int, _), res) -> (
+        match s with [] -> res = None | j :: _ -> res = Some j)
+    | Pop_all, Res ((List Int, _), res) -> res = s
+    | To_seq, Res ((Seq Int, _), res) -> List.of_seq res = s
     | Is_empty, Res ((Bool, _), res) -> res = (s = [])
     | _, _ -> false
 end