Source file accumulator.ml

open Kcas

let n_way_max = Domain.recommended_domain_count () |> Bits.ceil_pow_2
let n_way_default = n_way_max |> Int.min 8

type t = int Loc.t array

let make ?n_way n =
  let n_way =
    match n_way with
    | None -> n_way_default
    | Some n_way -> n_way |> Int.min n_way_max |> Bits.ceil_pow_2
  in
  Array.init n_way (fun i -> Loc.make (if i = 0 then n else 0))

let n_way_of = Array.length

let get_self a =
  let h = (Domain.self () :> int) in
  (* TODO: Consider mixing the bits of [h] to get better distribution *)
  Array.unsafe_get a (h land (Array.length a - 1))

module Xt = struct
  let add ~xt a n = if n <> 0 then Xt.fetch_and_add ~xt (get_self a) n |> ignore
  let incr ~xt a = Xt.incr ~xt (get_self a)
  let decr ~xt a = Xt.decr ~xt (get_self a)

  let rec get ~xt a s i =
    let s = s + Xt.get ~xt (Array.unsafe_get a i) in
    if i = 0 then s else get ~xt a s (i - 1)

  let get ~xt a =
    let i = Array.length a - 1 in
    let s = Xt.get ~xt (Array.unsafe_get a i) in
    if i = 0 then s else get ~xt a s (i - 1)

  let set ~xt a n = add ~xt a (n - get ~xt a)
end

module Tx = struct
  let add a n = Kcas.Xt.to_tx { tx = Xt.add a n }
  let incr a = Kcas.Xt.to_tx { tx = Xt.incr a }
  let decr a = Kcas.Xt.to_tx { tx = Xt.decr a }
  let get a = Kcas.Xt.to_tx { tx = Xt.get a }
  let set a n = Kcas.Xt.to_tx { tx = Xt.set a n }
end

let add a n = if n <> 0 then Loc.fetch_and_add (get_self a) n |> ignore
let incr a = Loc.incr (get_self a)
let decr a = Loc.decr (get_self a)
let get a = Kcas.Xt.commit { tx = Xt.get a }
let set a n = Kcas.Xt.commit { tx = Xt.set a n }