Source file lole.ml

module Option = struct
  let map_default default f = function Some v -> f v | None -> default
end

module type VALUE = sig
  type t

  val sentinel : t

  val weight : t -> int

  val merge : t -> t -> t option

  val pequal : t -> t -> bool

  val pp : t Fmt.t
end

module RBQ (V : VALUE) = struct
  module Queue = FQueue

  type t = {c: int; w: int; q: V.t Queue.t}

  and value = V.t

  let make capacity = {c= capacity; w= 0; q= Queue.empty}

  let pp ppf {c; w; q} =
    Fmt.pf ppf "{ @[<hov>c = %d;@ w = %d;@ q = %a;@] }" c w
      (Fmt.hvbox (Queue.pp V.pp))
      q

  let available t = t.c - t.w

  let push t v =
    let w = t.w + V.weight v in
    if w > t.c then Error t else Ok {t with w; q= Queue.push t.q v}

  let shift_exn t =
    match Queue.shift t.q with v, q -> (v, {t with w= t.w - V.weight v; q})

  let cons t v =
    let w = t.w + V.weight v in
    if w > t.c then Error t else Ok {t with w; q= Queue.cons t.q v}

  let cons_exn t v =
    if t.w + V.weight v > t.c then invalid_arg "cons_exn: no enough space" ;
    {t with w= t.w + V.weight v; q= Queue.cons t.q v}

  let weight t = Queue.fold (fun acc x -> acc + V.weight x) 0 t.q

  let to_list t = Queue.to_list t.q
end

type bigstring =
  (char, Bigarray.int8_unsigned_elt, Bigarray.c_layout) Bigarray.Array1.t

type 'a blitter = 'a -> int -> bigstring -> int -> int -> unit

let pp_chr =
  Fmt.using (function '\032' .. '\126' as x -> x | _ -> '.') Fmt.char

let pp_scalar : type buffer.
    get:(buffer -> int -> char) -> length:(buffer -> int) -> buffer Fmt.t =
 fun ~get ~length ppf b ->
  let l = length b in
  for i = 0 to l / 16 do
    Fmt.pf ppf "%08x: " (i * 16) ;
    let j = ref 0 in
    while !j < 16 do
      if (i * 16) + !j < l then
        Fmt.pf ppf "%02x" (Char.code @@ get b ((i * 16) + !j))
      else Fmt.pf ppf "  " ;
      if !j mod 2 <> 0 then Fmt.pf ppf " " ;
      incr j
    done ;
    Fmt.pf ppf "  " ;
    j := 0 ;
    while !j < 16 do
      if (i * 16) + !j < l then Fmt.pf ppf "%a" pp_chr (get b ((i * 16) + !j))
      else Fmt.pf ppf " " ;
      incr j
    done ;
    Fmt.pf ppf "@\n"
  done

let pp_string = pp_scalar ~get:String.get ~length:String.length

let pp_bytes = pp_scalar ~get:Bytes.get ~length:Bytes.length

let pp_bigstring =
  pp_scalar ~get:Bigarray.Array1.get ~length:Bigarray.Array1.dim

module RBA = struct
  type t = {r: int; w: int; c: int; b: bigstring}

  let is_power_of_two x = x <> 0 && x land (lnot x + 1) = x

  let create capacity =
    if not (is_power_of_two capacity) then
      invalid_arg "RBA.create: the capacity need to be a power of two" ;
    { r= 0
    ; w= 0
    ; c= capacity
    ; b= Bigarray.Array1.create Bigarray.char Bigarray.c_layout capacity }

  let[@inline] mask t v = v land (t.c - 1)

  let[@inline] empty t = t.r = t.w

  let[@inline] size t = t.w - t.r

  let[@inline] available t = t.c - (t.w - t.r)

  let[@inline] full t = size t = t.c

  let pp ppf ({r; w; c; b} as t) =
    let pp_rb r w m ppf b =
      if mask t r < mask t w then
        Fmt.pf ppf "%a" (Fmt.hvbox pp_bigstring)
          (Bigarray.Array1.sub b (mask t r) (w - r))
      else if empty t then Fmt.pf ppf "<empty>"
      else
        Fmt.pf ppf "@[<hov>%a and %a@]" (Fmt.hvbox pp_bigstring)
          (Bigarray.Array1.sub b (mask t r) (m - mask t r))
          (Fmt.hvbox pp_bigstring)
          (Bigarray.Array1.sub b 0 (mask t w))
    in
    Fmt.pf ppf "{ @[<hov>r = %d;@ w = %d;@ c = %d;@ b = %a;@] }" r w c
      (Fmt.hvbox (pp_rb r w c))
      b

  let push t v =
    let[@inline] mask t v = v land (t.c - 1) in
    t.b.{(mask [@inlined]) t t.w} <- v ;
    {t with w= t.w + 1}

  let shift t _ =
    let[@inline] mask t v = v land (t.c - 1) in
    let r = t.b.{(mask [@inlined]) t t.r} in
    (r, {t with r= t.r + 1})

  module N = struct
    let push t ~blit ~length ?(off = 0) ?len v =
      let len = match len with None -> length v | Some len -> len in
      let[@inline] mask t v = v land (t.c - 1) in
      let pre = t.c - (mask [@inlined]) t t.w in
      let extra = len - pre in
      let areas =
        if extra > 0 then (
          blit v off t.b ((mask [@inlined]) t t.w) pre ;
          blit v (off + pre) t.b 0 extra ;
          [ Bigarray.Array1.sub t.b ((mask [@inlined]) t t.w) pre
          ; Bigarray.Array1.sub t.b 0 extra ] )
        else (
          blit v off t.b ((mask [@inlined]) t t.w) len ;
          [Bigarray.Array1.sub t.b ((mask [@inlined]) t t.w) len] )
      in
      (areas, {t with w= t.w + len})

    let keep t ~blit ~length ?(off = 0) ?len v =
      let len = match len with None -> length v | Some len -> len in
      assert (size t >= len) ;
      let pre = t.c - mask t t.r in
      let extra = len - pre in
      if extra > 0 then (
        blit t.b (mask t t.r) v off pre ;
        blit t.b 0 v (off + pre) extra )
      else blit t.b (mask t t.r) v off len

    let shift t len = {t with r= t.r + len}
  end
end

module Buffer = struct
  type t = Bigstring of bigstring | String of string | Bytes of Bytes.t

  let weight = function
    | Bigstring raw -> Bigarray.Array1.dim raw
    | String raw -> String.length raw
    | Bytes raw -> Bytes.length raw

  let ppw_bigstring ppf b =
    let len = Bigarray.Array1.dim b in
    for i = 0 to len - 1 do
      Fmt.char ppf (Bigarray.Array1.unsafe_get b i)
    done

  let ppw ppf = function
    | Bigstring b -> ppw_bigstring ppf b
    | String b -> Fmt.string ppf b
    | Bytes b -> Fmt.string ppf (Bytes.unsafe_to_string b)

  let pp ppf = function
    | Bigstring b -> Fmt.pf ppf "(Bigstring %a)" (Fmt.hvbox pp_bigstring) b
    | Bytes b -> Fmt.pf ppf "(Bytes %a)" (Fmt.hvbox pp_bytes) b
    | String b -> Fmt.pf ppf "(String %a)" (Fmt.hvbox pp_string) b

  let sub buffer off len =
    match buffer with
    | Bigstring b -> Bigstring (Bigarray.Array1.sub b off len)
    | String b -> String (String.sub b off len)
    | Bytes b -> Bytes (Bytes.sub b off len)
end

module IOVec = struct
  type t = {buffer: Buffer.t; off: int; len: int}

  let weight {len; _} = len

  let sentinel =
    let deadbeef = "\222\173\190\239" in
    {buffer= Buffer.String deadbeef; off= 0; len= String.length deadbeef}

  let make buffer off len = {buffer; off; len}

  let length {len; _} = len

  let lengthv = List.fold_left (fun acc x -> length x + acc) 0

  let shift {buffer; off; len} n =
    assert (n <= len) ;
    {buffer; off= off + n; len= len - n}

  let split {buffer; off; len} n =
    assert (n <= len) ;
    ( {buffer= Buffer.sub buffer off n; off= 0; len= n}
    , {buffer= Buffer.sub buffer (off + n) (len - n); off= 0; len= len - n} )

  let pequal a b =
    match (a, b) with
    | {buffer= Buffer.Bytes a; _}, {buffer= Buffer.Bytes b; _} -> a == b
    | {buffer= Buffer.Bigstring a; _}, {buffer= Buffer.Bigstring b; _} ->
        a == b
    | _, _ -> false

  let merge a b =
    match (a, b) with
    | {buffer= Buffer.Bytes a'; _}, {buffer= Buffer.Bytes b'; _} ->
        assert (a' == b') ;
        if a.off + a.len = b.off then
          Some {buffer= Buffer.Bytes a'; off= a.off; len= a.len + b.len}
        else None
    | {buffer= Buffer.Bigstring a'; _}, {buffer= Buffer.Bigstring b'; _} ->
        assert (a' == b') ;
        if a.off + a.len = b.off then
          Some {buffer= Buffer.Bigstring a'; off= a.off; len= a.len + b.len}
        else None
    | _, _ -> None

  let ppw ppf = function
    | {buffer= Buffer.Bigstring b; off; len} ->
        Buffer.ppw_bigstring ppf (Bigarray.Array1.sub b off len)
    | {buffer= Buffer.String b; off; len} ->
        Fmt.string ppf (String.sub b off len)
    | {buffer= Buffer.Bytes b; off; len} ->
        Fmt.string ppf (Bytes.sub_string b off len)

  let pp ppf {buffer; off; len} =
    Fmt.pf ppf "{ @[<hov>buffer = %a;@ off = %d;@ len = %d:@] }"
      (Fmt.hvbox Buffer.pp) buffer off len
end

module RBS = RBQ (IOVec)

type encoder =
  { sched: RBS.t
  ; write: RBA.t
  ; flush: (int * (int -> encoder -> unit)) FQueue.t
  ; written: int
  ; received: int }

let pp ppf {sched; write; _} =
  Fmt.pf ppf "{ @[<hov>sched = %a;@ write = %a;@] }" (Fmt.hvbox RBS.pp) sched
    (Fmt.hvbox RBA.pp) write

type 'v state =
  | Flush of {continue: int -> 'v state; iovecs: IOVec.t list}
  | Continue of {continue: encoder -> 'v state; encoder: encoder}
  | End of 'v

let create len =
  { sched= RBS.make (len * 2)
  ; write= RBA.create len
  ; flush= FQueue.empty
  ; written= 0
  ; received= 0 }

let check iovec t =
  match iovec with
  | {IOVec.buffer= Buffer.Bigstring bs; _} ->
      let be =
        Bigarray.Array1.sub t.write.RBA.b (Bigarray.Array1.dim t.write.RBA.b) 0
      in
      let sub_ptr : int = Obj.magic @@ Obj.field (Obj.repr bs) 1 in
      let raw_ptr : int = Obj.magic @@ Obj.field (Obj.repr t.write.RBA.b) 1 in
      let end_ptr : int = Obj.magic @@ Obj.field (Obj.repr be) 1 in
      sub_ptr >= raw_ptr && sub_ptr <= end_ptr
  | _ -> false

let shift_buffers n t =
  let rec aux rest acc t =
    match RBS.shift_exn t.sched with
    | iovec, shifted ->
        let len = IOVec.length iovec in
        if rest > len then
          aux (rest - len) (iovec :: acc)
            { t with
              sched= shifted
            ; write=
                (if check iovec t then RBA.N.shift t.write len else t.write) }
        else if rest > 0 then
          let last, rest = IOVec.split iovec rest in
          ( List.rev (last :: acc)
          , { t with
              sched= RBS.cons_exn shifted rest
            ; write=
                ( if check iovec t then RBA.N.shift t.write (IOVec.length last)
                else t.write ) } )
        else (List.rev acc, t)
    | exception RBS.Queue.Empty -> (List.rev acc, t)
  in
  aux n [] t

let shift_flushes n t =
  let rec aux t =
    try
      let (threshold, f), flush = FQueue.shift t.flush in
      if
        compare (t.written + n - min_int) (threshold - min_int) >= 0
        (* unsigned int *)
      then
        let () = f n {t with flush} in
        aux {t with flush}
      else t
    with FQueue.Empty -> t
  in
  aux t

let shift n t =
  let lst, t = shift_buffers n t in
  ( lst
  , shift_flushes (IOVec.lengthv lst) t
    |> fun t -> {t with written= t.written + n} )

let has t = RBS.weight t.sched

let drain drain t =
  let rec go rest t =
    match RBS.shift_exn t.sched with
    | iovec, shifted ->
        let len = IOVec.length iovec in
        if rest > len then
          go (rest - len)
            { t with
              sched= shifted
            ; write=
                (if check iovec t then RBA.N.shift t.write len else t.write) }
        else
          { t with
            sched= RBS.cons_exn shifted (IOVec.shift iovec rest)
          ; write= (if check iovec t then RBA.N.shift t.write rest else t.write)
          }
    | exception RBS.Queue.Empty -> t
  in
  go drain t |> fun t -> {t with written= t.written + drain}

let flush k t =
  let t = shift_flushes (has t) t in
  let continue n =
    let t = drain n t in
    k {t with written= t.written + n}
  in
  Flush {continue; iovecs= RBS.to_list t.sched}

let continue continue encoder = Continue {continue; encoder}

let rec schedule k ~length ~buffer ?(off = 0) ?len v t =
  let len = match len with Some len -> len | None -> length v - off in
  match RBS.push t.sched (IOVec.make (buffer v) off len) with
  | Ok sched -> continue k {t with sched; received= t.received + len}
  | Error _ ->
      let max = RBS.available t.sched in
      let k t =
        (schedule [@tailcall]) k ~length ~buffer ~off:(off + max)
          ~len:(len - max) v t
      in
      schedule (flush k) ~length ~buffer ~off ~len:max v t

let schedule_string =
  let length = String.length in
  let buffer x = Buffer.String x in
  fun k t ?(off = 0) ?len v -> schedule k ~length ~buffer ~off ?len v t

let schedule_bytes =
  let length = Bytes.length in
  let buffer x = Buffer.Bytes x in
  fun k t ?(off = 0) ?len v -> schedule k ~length ~buffer ~off ?len v t

let schedule_bigstring =
  let length = Bigarray.Array1.dim in
  let buffer x = Buffer.Bigstring x in
  fun k t ?(off = 0) ?len v -> schedule k ~length ~buffer ~off ?len v t

let schedule_flush f t = {t with flush= FQueue.push t.flush (t.received, f)}

external identity : 'a -> 'a = "%identity"

let schedulev k l t =
  let rec aux t = function
    | [] -> continue k t
    | (length, off, len, buffer) :: r ->
        schedule
          (fun t -> (aux [@tailcall]) t r)
          ~length ?off ?len ~buffer:identity buffer t
  in
  aux t l

let schedulev_bigstring k l t =
  let rec aux t = function
    | [] -> continue k t
    | buffer :: r ->
        schedule_bigstring (fun t -> (aux [@tailcall]) t r) t buffer
  in
  aux t l

let rec write k ~blit ~length ?(off = 0) ?len buffer t =
  let len = match len with Some len -> len | None -> length buffer - off in
  let available = RBA.available t.write in
  (* XXX(dinosaure): we can factorize the first and the second branch. *)
  if available >= len then
    let areas, write = RBA.N.push t.write ~blit ~length ~off ~len buffer in
    schedulev_bigstring k areas {t with write}
  else if available > 0 then
    let k t =
      (write [@tailcall]) k ~blit ~length ~off:(off + available)
        ~len:(len - available) buffer t
    in
    let areas, write =
      RBA.N.push t.write ~blit ~length ~off ~len:available buffer
    in
    schedulev_bigstring (flush k) areas {t with write}
  else
    let k t = (write [@tailcall]) k ~blit ~length ~off ~len buffer t in
    flush k t

let writev k l t =
  let rec aux t = function
    | [] -> continue k t
    | (blit, length, off, len, buffer) :: r ->
        write (fun t -> (aux [@tailcall]) t r) ~blit ~length ?off ?len buffer t
  in
  aux t l

let bigarray_blit_from_string src src_off dst dst_off len =
  for i = 0 to len - 1 do
    Bigarray.Array1.unsafe_set dst (dst_off + i)
      (String.unsafe_get src (src_off + i))
  done

let bigarray_blit_from_bytes src src_off dst dst_off len =
  for i = 0 to len - 1 do
    Bigarray.Array1.unsafe_set dst (dst_off + i)
      (Bytes.unsafe_get src (src_off + i))
  done

let bigarray_blit src src_off dst dst_off len =
  Bigarray.Array1.(blit (sub src src_off len) (sub dst dst_off len))

let bigarray_blit_to_bytes src src_off dst dst_off len =
  for i = 0 to len - 1 do
    Bytes.set dst (dst_off + i) (Bigarray.Array1.unsafe_get src (src_off + i))
  done

let write_string =
  let length = String.length in
  let blit = bigarray_blit_from_string in
  fun ?(off = 0) ?len a k t -> write k ~blit ~length ~off ?len a t

let write_bytes =
  let length = Bytes.length in
  let blit = bigarray_blit_from_bytes in
  fun ?(off = 0) ?len a k t -> write k ~blit ~length ~off ?len a t

let write_bigstring =
  let length = Bigarray.Array1.dim in
  let blit = bigarray_blit in
  fun ?(off = 0) ?len a k t -> write k ~blit ~length ~off ?len a t

let write_char =
  let length _ = assert false in
  let blit src src_off dst dst_off len =
    assert (src_off = 0) ;
    assert (len = 1) ;
    EndianBigstring.BigEndian_unsafe.set_char dst dst_off src
  in
  fun a k t -> write k ~length ~blit ~off:0 ~len:1 a t

let write_uint8 =
  let length _ = assert false in
  let blit src src_off dst dst_off len =
    assert (src_off = 0) ;
    assert (len = 1) ;
    EndianBigstring.BigEndian_unsafe.set_int8 dst dst_off src
  in
  fun a k t -> write k ~length ~blit ~off:0 ~len:1 a t

module type EndianBigstringSig = EndianBigstring.EndianBigstringSig

module type EndianBytesSig = EndianBytes.EndianBytesSig

module type SE = sig
  val write_uint16 : int -> (encoder -> 'r state) -> encoder -> 'r state

  val write_uint32 : int32 -> (encoder -> 'r state) -> encoder -> 'r state

  val write_uint64 : int64 -> (encoder -> 'r state) -> encoder -> 'r state
end

module MakeE (EBigstring : EndianBigstringSig) : SE = struct
  let _length _ = assert false

  let write_uint16 =
    let length = _length in
    let blit src src_off dst dst_off len =
      assert (src_off = 0) ;
      assert (len = 2) ;
      EBigstring.set_int16 dst dst_off src
    in
    fun a k t -> write k ~length ~blit ~off:0 ~len:2 a t

  let write_uint32 =
    let length = _length in
    let blit src src_off dst dst_off len =
      assert (src_off = 0) ;
      assert (len = 4) ;
      EBigstring.set_int32 dst dst_off src
    in
    fun a k t -> write k ~length ~blit ~off:0 ~len:4 a t

  let write_uint64 =
    let length = _length in
    let blit src src_off dst dst_off len =
      assert (src_off = 0) ;
      assert (len = 8) ;
      EBigstring.set_int64 dst dst_off src
    in
    fun a k t -> write k ~length ~blit ~off:0 ~len:8 a t
end

module LE = MakeE (EndianBigstring.LittleEndian_unsafe)
module BE = MakeE (EndianBigstring.BigEndian_unsafe)