Source file encoder.ml

type vec = { off : int option; len : int option }

type 'a state = 'a Lole.state

type encoder = Lole.encoder

type bigstring = Lole.bigstring

type iovecs = Lole.IOVec.t list

type -'a t = { run : 'r. (encoder -> 'r state) -> encoder -> 'a -> 'r state }

type -'a s = {
  sub :
    'r. (encoder -> 'r state) -> encoder -> ?off:int -> ?len:int -> 'a ->
    'r state;
}

let peek : 'a t -> 'b t -> ('a, 'b) Either.t t =
 fun a b ->
  { run = (fun k e -> function L x -> a.run k e x | R y -> b.run k e y) }

let char : char t = { run = (fun k e v -> Lole.write_char v k e) }

let int8 : int t = { run = (fun k e v -> Lole.write_uint8 v k e) }

let beint16 : int t = { run = (fun k e v -> Lole.BE.write_uint16 v k e) }

let beint32 : int32 t = { run = (fun k e v -> Lole.BE.write_uint32 v k e) }

let beint64 : int64 t = { run = (fun k e v -> Lole.BE.write_uint64 v k e) }

let leint16 : int t = { run = (fun k e v -> Lole.LE.write_uint16 v k e) }

let leint32 : int32 t = { run = (fun k e v -> Lole.LE.write_uint32 v k e) }

let leint64 : int64 t = { run = (fun k e v -> Lole.LE.write_uint64 v k e) }

let bool : bool t =
  {
    run =
      (fun k e -> function
        | true -> char.run k e '1'
        | false -> char.run k e '0');
  }

let substring : string s =
  { sub = (fun k e ?off ?len v -> Lole.write_string ?off ?len v k e) }

let subbytes : bytes s =
  { sub = (fun k e ?off ?len v -> Lole.write_bytes ?off ?len v k e) }

let subbigstring : bigstring s =
  { sub = (fun k e ?off ?len v -> Lole.write_bigstring ?off ?len v k e) }

let blitter length blit : _ s =
  { sub = (fun k e ?off ?len v -> Lole.write k ~blit ~length ?off ?len v e) }

let whole (a : 'v s) : 'v t =
  { run = (fun k e v -> a.sub ?off:None ?len:None k e v) }

let sub (a : 'v s) : (vec * 'v) t =
  { run = (fun k e ({ off; len }, v) -> a.sub ?off ?len k e v) }

let string : string t = whole substring

let bytes : bytes t = whole subbytes

let bigstring : bigstring t = whole subbigstring

let list ?sep a : 'a list t =
  let sep k e = match sep with None -> k e | Some a -> a.run k e () in
  let rec run k e : _ list -> _ state = function
    | [] -> k e
    | [ x ] -> a.run k e x
    | x :: r -> a.run (sep (fun e -> run k e r)) e x in
  { run }

let nop = { run = (fun k e _ -> k e) }

let option f : 'a option t =
  { run = (fun k e -> function Some v -> f.run k e v | None -> k e) }

exception Fail of string

let pure ~compare v =
  {
    run =
      (fun k e v' ->
        if compare v v' = 0
        then k e
        else raise (Fail "fail at the pure operator"));
  }

let fail s = { run = (fun _k _e _v -> raise (Fail s)) }

let const s =
  {
    run =
      (fun k e s' ->
        if String.equal s s'
        then Lole.write_string s' k e
        else raise (Fail (Fmt.strf "const: %s <> %s" s s')));
  }

let ( <|> ) pu pv =
  {
    run =
      (fun k e v ->
        try pu.run k e v with Fail _ | Bijection.Exn.Bijection -> pv.run k e v);
  }

let ( <$> ) f p = { run = (fun k e v -> p.run k e (f v)) }

let ( <*> ) a b = { run = (fun k e (x, y) -> a.run (fun e -> b.run k e y) e x) }

let prefix p r = { run = (fun k e v -> p.run (fun e -> r.run k e v) e ()) }

let suffix s r = { run = (fun k e v -> r.run (fun e -> s.run k e ()) e v) }

exception Break

let for_all predicate s =
  let l = String.length s in
  try
    for i = 0 to l - 1 do
      if not (predicate (String.unsafe_get s i)) then raise Break
    done ;
    true
  with Break -> false

let while0 predicate =
  {
    run =
      (fun k e v ->
        if for_all predicate v
        then Lole.write_string v k e
        else raise (Fail "while0"));
  }

let while1 predicate =
  {
    run =
      (fun k e v ->
        if String.length v > 0 && for_all predicate v
        then Lole.write_string v k e
        else raise (Fail "while1"));
  }

let for_all predicate b =
  let l = Bigarray_compat.Array1.dim b in
  try
    for i = 0 to l - 1 do
      if not (predicate b.{i}) then raise Break
    done ;
    true
  with Break -> false

let bigstring_while0 predicate =
  {
    run =
      (fun k e v ->
        if for_all predicate v
        then Lole.write_bigstring v k e
        else raise (Fail "bigstring_while0"));
  }

let bigstring_while1 predicate =
  {
    run =
      (fun k e v ->
        if Bigarray_compat.Array1.dim v > 0 && for_all predicate v
        then Lole.write_bigstring v k e
        else raise (Fail "bigstring_while1"));
  }

let take n =
  (* XXX(dinosaure): Angstrom.take consumes input, so [take] should produce output. *)
  {
    run =
      (fun k e s ->
        if String.length s = n then string.run k e s else raise (Fail "take"));
  }

let buffer = string

let bigstring_buffer = bigstring

let ( <* ) r s = suffix s r

let ( *> ) p r = prefix p r

let fix f =
  let rec p = lazy (f r)
  and r = { run = (fun k e v -> Lazy.(force p).run k e v) } in
  r

let commit = { run = (fun k e () -> Lole.flush k e) }

let keval :
      'v 'r. (encoder -> 'r state) -> (iovecs -> int) -> encoder -> 'v t ->
      'v -> 'r =
 fun k w e t v ->
  let rec go = function
    | Lole.End v -> v
    | Lole.Continue { continue; encoder } -> continue encoder |> go
    | Lole.Flush { continue; iovecs } ->
        let len = w iovecs in
        continue len |> go in
  t.run k e v |> go

let eval w e t v = keval (fun _e -> Lole.End ()) w e (t <* commit) v

let run t = t.run

module Make (S : sig
  type a

  val run : (encoder -> 'r state) -> encoder -> a -> 'r state
end) =
struct
  let x = { run = S.run }
end

let to_string : type a. a t -> a -> string =
 fun t v ->
  let buf = Buffer.create 16 in
  let writer l =
    List.iter
      (function
        | { Lole.IOVec.buffer = Lole.Buffer.String s; off; len } ->
            Buffer.add_substring buf s off len
        | { Lole.IOVec.buffer = Lole.Buffer.Bytes s; off; len } ->
            Buffer.add_subbytes buf s off len
        | { Lole.IOVec.buffer = Lole.Buffer.Bigstring s; off; len } ->
            for i = 0 to len - 1 do
              Buffer.add_char buf s.{off + i}
            done)
      l ;
    Lole.IOVec.lengthv l in
  eval writer (Lole.create 0x100) t v ;
  Buffer.contents buf