Skip to content

Commit

Permalink
Add new functions and improve documentations of Treiber stack. (#158)
Browse files Browse the repository at this point in the history
Add new functions and improve documentations of Treiber stack
  • Loading branch information
lyrm authored Nov 21, 2024
1 parent 37d41e2 commit 8ed5de4
Show file tree
Hide file tree
Showing 4 changed files with 353 additions and 86 deletions.
85 changes: 68 additions & 17 deletions src/treiber_stack.ml
Original file line number Diff line number Diff line change
@@ -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)
152 changes: 137 additions & 15 deletions src/treiber_stack.mli
Original file line number Diff line number Diff line change
@@ -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]
]}
*)
63 changes: 47 additions & 16 deletions test/treiber_stack/stm_treiber_stack.ml
Original file line number Diff line number Diff line change
@@ -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
Expand Down
Loading

0 comments on commit 8ed5de4

Please sign in to comment.