Source file fke.ml

[@@@warning "-37"]

module Peano = struct
  type zero = Zero
  type 'a succ = Succ
  type one = zero succ
  type two = zero succ succ
  type three = zero succ succ
end

type ('a, 'l) digit =
  | Zero : ('a, Peano.zero) digit
  | One : 'a -> ('a, Peano.one) digit
  | Two : 'a * 'a -> ('a, Peano.two) digit
  | Three : 'a * 'a * 'a -> ('a, Peano.three) digit

type 'a t =
  | Shallow : ('a, 'l) digit -> 'a t
  | Deep : {
      s : int;
      f : ('a, 'f Peano.succ) digit;
      m : ('a * 'a) t Lazy.t;
      r : ('a, 'r Peano.succ) digit;
    }
      -> 'a t

let empty = Shallow Zero

exception Empty

let _one x = Shallow (One x)
let _two x y = Shallow (Two (x, y))
let _three x y z = Shallow (Three (x, y, z))
let _deep s f m r = Deep { s; f; m; r }

let is_empty : type a. a t -> bool = function
  | Shallow Zero -> true
  | Shallow _ | Deep _ -> false

let rec push : type a. a t -> a -> a t =
 fun q x ->
  match q with
  | Shallow Zero -> _one x
  | Shallow (One y) -> _two y x
  | Shallow (Two (y, z)) -> _three y z x
  | Shallow (Three (a, b, c)) ->
      _deep 4 (Two (a, b)) (Lazy.from_val empty) (Two (c, x))
  | Deep { s; f; m; r = One y } -> _deep (s + 1) f m (Two (y, x))
  | Deep { s; f; m; r = Two (y, z) } -> _deep (s + 1) f m (Three (y, z, x))
  | Deep { s; f; m = (lazy q'); r = Three (y, z, z') } ->
      _deep (s + 1) f (lazy (push q' (y, z))) (Two (z', x))

let rec pop_exn : type a. a t -> a * a t =
 fun q ->
  match q with
  | Shallow Zero -> raise Empty
  | Shallow (One x) -> (x, empty)
  | Shallow (Two (x, y)) -> (x, _one y)
  | Shallow (Three (x, y, z)) -> (x, _two y z)
  | Deep { s; f = One x; m = (lazy q'); r } ->
      if is_empty q' then (x, Shallow r)
      else
        let (y, z), q' = pop_exn q' in
        (x, _deep (s - 1) (Two (y, z)) (Lazy.from_val q') r)
  | Deep { s; f = Two (x, y); m; r } -> (x, _deep (s - 1) (One y) m r)
  | Deep { s; f = Three (x, y, z); m; r } -> (x, _deep (s - 1) (Two (y, z)) m r)

let rec tail_exn : type a. a t -> a t * a =
 fun q ->
  match q with
  | Shallow Zero -> raise Empty
  | Shallow (One x) -> (empty, x)
  | Shallow (Two (x, y)) -> (_one x, y)
  | Shallow (Three (x, y, z)) -> (_two x y, z)
  | Deep { s; f; m = (lazy q'); r = One x } ->
      if is_empty q' then (Shallow f, x)
      else
        let q'', (y, z) = tail_exn q' in
        (_deep (s - 1) f (Lazy.from_val q'') (Two (y, z)), x)
  | Deep { s; f; m; r = Two (x, y) } -> (_deep (s - 1) f m (One x), y)
  | Deep { s; f; m; r = Three (x, y, z) } -> (_deep (s - 1) f m (Two (x, y)), z)

let peek_exn : type a. a t -> a =
 fun q ->
  match q with
  | Shallow Zero -> raise Empty
  | Shallow (One x) -> x
  | Shallow (Two (x, _)) -> x
  | Shallow (Three (x, _, _)) -> x
  | Deep { f = One x; _ } -> x
  | Deep { f = Two (x, _); _ } -> x
  | Deep { f = Three (x, _, _); _ } -> x

let pop q = try Some (pop_exn q) with Empty -> None
let tail q = try Some (tail_exn q) with Empty -> None
let peek q = try Some (peek_exn q) with Empty -> None

let rec cons : type a. a t -> a -> a t =
 fun q x ->
  match q with
  | Shallow Zero -> _one x
  | Shallow (One y) -> _two x y
  | Shallow (Two (y, z)) -> _three x y z
  | Shallow (Three (y, z, z')) ->
      _deep 4 (Two (x, y)) (Lazy.from_val empty) (Two (z, z'))
  | Deep { s; f = One y; m; r } -> _deep (s + 1) (Two (x, y)) m r
  | Deep { s; f = Two (y, z); m; r } -> _deep (s + 1) (Three (x, y, z)) m r
  | Deep { s; f = Three (y, z, z'); m = (lazy q'); r } ->
      _deep (s + 1) (Three (x, y, z)) (lazy (cons q' (z, z'))) r

let iter : type a. (a -> unit) -> a t -> unit =
 fun f q ->
  let rec go : type a. (a -> unit) -> a t -> unit =
   fun f -> function
    | Shallow Zero -> ()
    | Shallow (One x) -> f x
    | Shallow (Two (x, y)) ->
        f x;
        f y
    | Shallow (Three (x, y, z)) ->
        f x;
        f y;
        f z
    | Deep { f = hd; m = (lazy q); r = tl; _ } ->
        go f (Shallow hd);
        go
          (fun (x, y) ->
            f x;
            f y)
          q;
        go f (Shallow tl)
  in
  go f q

let rev_iter : type a. (a -> unit) -> a t -> unit =
 fun f q ->
  let rec go : type a. (a -> unit) -> a t -> unit =
   fun f -> function
    | Shallow Zero -> ()
    | Shallow (One x) -> f x
    | Shallow (Two (y, x)) ->
        f x;
        f y
    | Shallow (Three (z, y, x)) ->
        f x;
        f y;
        f z
    | Deep { f = hd; m = (lazy q); r = tl; _ } ->
        go f (Shallow tl);
        go
          (fun (y, x) ->
            f x;
            f y)
          q;
        go f (Shallow hd)
  in
  go f q

let fold : type acc x. (acc -> x -> acc) -> acc -> x t -> acc =
 fun f a q ->
  let rec go : type acc x. (acc -> x -> acc) -> acc -> x t -> acc =
   fun f a -> function
    | Shallow Zero -> a
    | Shallow (One x) -> f a x
    | Shallow (Two (x, y)) -> f (f a x) y
    | Shallow (Three (x, y, z)) -> f (f (f a x) y) z
    | Deep { f = hd; m = (lazy q); r = tl; _ } ->
        let a = go f a (Shallow hd) in
        let a = go (fun a (x, y) -> f (f a x) y) a q in
        go f a (Shallow tl)
  in
  go f a q

let length = function
  | Deep { s; _ } -> s
  | Shallow Zero -> 0
  | Shallow (One _) -> 1
  | Shallow (Two _) -> 2
  | Shallow (Three _) -> 3

let pp ?sep pp_elt = Fmt.iter ?sep iter pp_elt
let dump pp_elt = Fmt.Dump.iter iter (Fmt.any "fke") pp_elt

module Weighted = struct
  type ('a, 'b) t = {
    r : int;
    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] 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[@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] is_power_of_two v = v <> 0 && v land (lnot v + 1) = v

  let create ?capacity kind =
    let capacity =
      match capacity with
      | None | Some 0 -> 1
      | Some n ->
          if n < 0 then Fmt.invalid_arg "Rke.Weighted.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 )

  let copy t =
    let v = Bigarray.Array1.create t.k Bigarray.c_layout t.c in
    Bigarray.Array1.blit t.v v;
    { r = t.r; w = t.w; c = t.c; v; k = t.k }

  let from v =
    if not (is_power_of_two (Bigarray.Array1.dim v)) then
      Fmt.invalid_arg "RBA.from";
    let c = Bigarray.Array1.dim v in
    let k = Bigarray.Array1.kind v in
    { r = 0; w = 0; c; k; v }

  let push_exn t v =
    if (full [@inlined]) t then raise Full;
    Bigarray.Array1.unsafe_set t.v ((mask [@inlined]) t t.w) v;
    { t with w = t.w + 1 }

  let push t v = try Some (push_exn t v) with Full -> None

  let cons_exn t v =
    if (full [@inlined]) t then raise Full;
    let i = t.r - 1 in
    Bigarray.Array1.unsafe_set t.v ((mask [@inlined]) t i) v;
    { t with r = i }

  let cons t v = try Some (cons_exn t v) with Full -> None

  let pop_exn t =
    if (empty [@inlined]) t then raise Empty;
    let r = Bigarray.Array1.unsafe_get t.v ((mask [@inlined]) t t.r) in
    (r, { t with r = t.r + 1 })

  let pop t = try Some (pop_exn t) with Empty -> None

  let peek_exn t =
    if (empty [@inlined]) t then raise Empty;
    Bigarray.Array1.unsafe_get t.v ((mask [@inlined]) t t.r)

  let peek t = try Some (peek_exn t) with Empty -> None

  module N = struct
    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
      (ret, { t with w = t.w + len })

    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 | 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 unsafe_shift t len = { t with r = t.r + len }

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

    let shift t len = try Some (shift_exn t len) with Empty -> None
  end

  let iter f t =
    let idx = ref t.r in
    let max = t.w in
    while !idx <> max do
      f (Bigarray.Array1.unsafe_get t.v ((mask [@inlined]) t !idx));
      incr idx
    done

  let rev_iter f t =
    if t.r == t.w then ()
    else
      let idx = ref (pred t.w) in
      let min = t.r in
      while
        f (Bigarray.Array1.unsafe_get t.v ((mask [@inlined]) t !idx));
        !idx <> min
      do
        decr idx
      done

  let fold f a t =
    let a = ref a in
    iter (fun x -> a := f !a x) t;
    !a

  let clear t = { t with r = 0; w = 0 }
  let unsafe_bigarray { v; _ } = v
  let pp ?sep pp_elt = Fmt.iter ?sep iter pp_elt
  let dump pp_elt = Fmt.Dump.iter iter (Fmt.any "fke:weighted") pp_elt
end