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.Array1.dim b in

  try for i = 0 to l - 1 do if not (predicate (Bigarray.Array1.get 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.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 (Bigarray.Array1.get s (off + i)) done)
        l;
      Lole.IOVec.lengthv l in
    eval writer (Lole.create 0x100) t v;
    Buffer.contents buf