Source file ke.ml

type ('a, 'b) t =
  { mutable r: int
  ; mutable w: int
  ; c: int
  ; k: ('a, 'b) Bigarray.kind
  ; v: ('a, 'b, Bigarray.c_layout) Bigarray.Array1.t }

exception Empty
exception Full

let[@inline always] to_power_of_two v =
  let res = ref (pred v) in
  res := !res lor (!res lsr 1) ;
  res := !res lor (!res lsr 2) ;
  res := !res lor (!res lsr 4) ;
  res := !res lor (!res lsr 8) ;
  res := !res lor (!res lsr 16) ;
  succ !res

let[@inline always] mask t v = v land (t.c - 1)
let[@inline always] empty t = t.r = t.w
let[@inline always] size t = t.w - t.r
let[@inline always] full t = size t = t.c
let[@inline always] available t = t.c - (t.w - t.r)
let is_empty t = (empty [@inlined]) t
let length q = size q

let create ?capacity kind =
  let capacity =
    match capacity with
    | None | Some 0 -> 1
    | Some n ->
      if n < 0 then Fmt.invalid_arg "Ke.create"
      else to_power_of_two n
  in
  ( { r= 0
    ; w= 0
    ; c= capacity
    ; k= kind
    ; v= Bigarray.Array1.create kind Bigarray.c_layout capacity }
  , capacity )

type ('a, 'b) bigarray = ('a, 'b, Bigarray.c_layout) Bigarray.Array1.t
type ('a, 'b) blit = 'a -> int -> 'b -> int -> int -> unit
type 'a length = 'a -> int

let push_exn t ~blit ~length ?(off = 0) ?len v =
  let len = match len with None -> length v - off | Some len -> len in
  if (available [@inlined]) t < len then raise Full ;
  let msk = (mask [@inlined]) t t.w in
  let pre = t.c - msk in
  let rst = len - pre in
  let ret =
    if rst > 0 then (
      blit v off t.v msk pre ;
      blit v (off + pre) t.v 0 rst ;
      [ Bigarray.Array1.sub t.v ((mask [@inlined]) t t.w) pre
      ; Bigarray.Array1.sub t.v 0 rst ] )
    else (
      blit v off t.v msk len ;
      [Bigarray.Array1.sub t.v ((mask [@inlined]) t t.w) len] )
  in
  t.w <- t.w + len ;
  ret

let push t ~blit ~length ?off ?len v =
  try Some (push_exn t ~blit ~length ?off ?len v) with Full -> None

let keep_exn t ~blit ~length ?(off = 0) ?len v =
  let len = match len with None -> length v - off | Some len -> len in
  if (size [@inlined]) t < len then raise Empty ;
  let msk = (mask [@inlined]) t t.r in
  let pre = t.c - msk in
  let rst = len - pre in
  if rst > 0 then (
    blit t.v msk v off pre ;
    blit t.v 0 v (off + pre) rst )
  else blit t.v msk v off len

let keep t ~blit ~length ?off ?len v =
  try Some (keep_exn t ~blit ~length ?off ?len v) with Empty -> None

let peek t =
  let len = (size [@inlined]) t in
  if len == 0 then []
  else
    let msk = (mask [@inlined]) t t.r in
    let pre = t.c - msk in
    let rst = len - pre in
    if rst > 0 then
      [ Bigarray.Array1.sub t.v msk pre
      ; Bigarray.Array1.sub t.v 0 rst ]
    else [ Bigarray.Array1.sub t.v msk len ]

let unsafe_shift t len = t.r <- t.r + len

let shift_exn t len =
  if (size [@inlined]) t < len then raise Empty ;
  unsafe_shift t len