OCaml binary parser monad

Parser.ml

module Opt = struct

  module Core = struct
    type 'a t = 'a option

    let pure a = Some a
    let map f = function
      | None -> None
      | Some a -> Some (f a)
    let ap ft xt = match ft, xt with
      | Some f, Some x -> Some (f x)
      | _ -> None
    let map2 f xt yt = match xt, yt with
      | Some x, Some y -> Some (f x y)
      | _ -> None
    let bind k m = match m with
      | None -> None
      | Some a -> k a

    let fold some none = function
      | None -> none
      | Some a -> some a

    let guard p k a = if p a then k a else None

    let zero = None

    let choice x y = match x, y with
      | Some _, _ -> x
      | _ -> y
  end

  module Infix = struct
    include Core
    let ( >>= ) m k = bind k m
    let ( + ) = choice
  end

  include Core
end

module type Source = sig
  type t

  val of_bytes  : bytes -> t
  val of_string : string -> t

  val uncons : t -> (char * t) option
  val takeExactly : int -> t -> (bytes * t) option
end

module BSource : Source = struct
  type t = { backing : bytes; ix : int }

  let of_bytes  b = { backing = b; ix = 0 }
  let of_string b = of_bytes (Bytes.of_string b)

  let remains b = Bytes.length b.backing - b.ix

  let uncons b =
    if b.ix <= Bytes.length b.backing
    then Some (b.backing.[b.ix], { b with ix = b.ix + 1 })
    else None

  let takeExactly n b =
    if 0 <= n && n <= remains b
    then Some (Bytes.sub b.backing 0 n, {b with ix = b.ix + n})
    else None

end

type 'a t = BSource.t -> ('a * BSource.t) option

let run t b = Opt.map fst (t @@ BSource.of_bytes b)

let run_file t ~filepath:f =
  let ic = open_in f in
  let n = in_channel_length ic in
  let b = Bytes.create n in
  really_input ic b 0 n;
  close_in ic;
  run t b

let keep f s =
  Opt.Infix.(BSource.uncons s >>= fun (c, s) -> f c >>= fun r -> pure (r, s))

let satisfy p = keep (fun x -> if p x then Some x else None)

let byte (c : char) : char t = satisfy ((=) c)

let bytes c s =
  let open Opt.Infix in
  BSource.takeExactly (Bytes.length c) s >>= fun (sub, s) ->
  if Bytes.compare c sub = 0
  then pure (sub, s)
  else zero

let string = bytes

let notFollowedBy p s =
  match p s with
  | None   -> Some ((), s)
  | Some _ -> None

let eof = notFollowedBy (satisfy (fun _ -> true))

let take n = BSource.takeExactly n

let pure a s = Some (a, s)

let map f t s = Opt.map (fun (a, b) -> (f a, b)) (t s)

let ap ft xt s =
  let open Opt.Infix in
  ft s >>= fun (f, s) ->
  xt s >>= fun (x, s) ->
  pure (f x, s)

let map2 f xt yt s =
  let open Opt.Infix in
  xt s >>= fun (x, s) ->
  yt s >>= fun (y, s) ->
  pure (f x y, s)

let bind k m s =
  match m s with
  | None -> None
  | Some (a, s) -> k a s

let zero s = None

let (+) x y s =
  match x s with
  | Some a -> Some a
  | None   -> y s

module Infix = struct
  let ( >>= ) m k = bind k m
  let ( <$> ) = map
  let ( <*> ) = ap
  let ( <*  ) x y = map2 (fun x y -> x) x y
  let (  *> ) x y = map2 (fun x y -> y) x y
end

open Infix

module Traverse = struct
  let rec list (f : 'a -> 'b t) (ls : 'a list) : 'b list t =
    match ls with
    | [] -> pure []
    | x :: xs -> Infix.(map (fun a b -> a :: b) (f x) <*> list f xs)

  let option (f : 'a -> 'b t) (x : 'a option) : 'b option t =
    match x with
    | None -> pure None
    | Some a -> map (fun x -> Some x) (f a)
end

let oneof ls = List.fold_left (+) zero ls

let optional f = map Opt.pure f + pure None

(* ----------------------------- *)

(* These are nowhere near intelligent enough. *)

let word8 = keep (fun x -> Some (Char.code x))

let word16be =
  take 2 >>= fun bs ->
  let b n = Char.code bs.[n] in
  pure @@ (b 0 lsl 8) lor b 1

let word16le =
  take 2 >>= fun bs ->
  let b n = Char.code bs.[n] in
  pure @@ (b 1 lsl 8) lor b 0

let word32be =
  take 4 >>= fun bs ->
  let b n = Char.code bs.[n] in
  pure @@     (b 0 lsl 24)
          lor (b 1 lsl 16)
          lor (b 2 lsl 8 )
          lor (b 3)

let word32le =
  take 4 >>= fun bs ->
  let b n = Char.code bs.[n] in
  pure @@     (b 3 lsl 24)
          lor (b 2 lsl 16)
          lor (b 1 lsl 8 )
          lor (b 0)

let word64be =
  take 8 >>= fun bs ->
  let b n = Char.code bs.[n] in
  pure @@     (b 0 lsl 56)
          lor (b 1 lsl 48)
          lor (b 2 lsl 40)
          lor (b 3 lsl 32)
          lor (b 4 lsl 24)
          lor (b 5 lsl 16)
          lor (b 6 lsl 8 )
          lor (b 7)

let word64le =
  take 8 >>= fun bs ->
  let b n = Char.code bs.[n] in
  pure @@     (b 7 lsl 56)
          lor (b 6 lsl 48)
          lor (b 5 lsl 40)
          lor (b 4 lsl 32)
          lor (b 3 lsl 24)
          lor (b 2 lsl 16)
          lor (b 1 lsl 8 )
          lor (b 0)

(* val word32be_ : Int32.t t *)
(* val word32be_ : Int32.t t *)
(* val word64le_ : Int64.t t *)
(* val word64le_ : Int64.t t *)

Parser.mli

type 'a t
val run : 'a t -> bytes -> 'a option
val run_file : 'a t -> filepath:string -> 'a option

val keep : (char -> 'a option) -> 'a t
val satisfy : (char -> bool) -> char t
val byte : char -> char t
val bytes : bytes -> bytes t
val string : string -> string t
val notFollowedBy : 'a t -> unit t
val eof : unit t

val take : int -> bytes t

val word8     : int t
val word16be  : int t
val word16le  : int t
val word32be  : int t
val word32le  : int t
val word64be  : int t
val word64le  : int t
(* val word32be_ : Int32.t t *)
(* val word32be_ : Int32.t t *)
(* val word64le_ : Int64.t t *)
(* val word64le_ : Int64.t t *)

val map : ('a -> 'b) -> ('a t -> 'b t)
val pure : 'a -> 'a t
val ap : ('a -> 'b) t -> ('a t -> 'b t)
val map2 : ('a -> 'b -> 'c) -> ('a t -> 'b t -> 'c t)
val bind : ('a -> 'b t) -> ('a t -> 'b t)
val zero  : 'a t
val ( + ) : 'a t -> 'a t -> 'a t
val optional : 'a t -> 'a option t

module Infix : sig
  val ( >>= ) : 'a t -> ('a -> 'b t) -> 'b t
  val ( <$> ) : ('a -> 'b) -> ('a t -> 'b t)
  val ( <*> ) : ('a -> 'b) t -> ('a t -> 'b t)
  val ( <*  ) : 'a t -> 'b t -> 'a t
  val (  *> ) : 'a t -> 'b t -> 'b t
end

module Traverse : sig
  val list : ('a -> 'b t) -> 'a list -> 'b list t
  val option : ('a -> 'b t) -> 'a option -> 'b option t
end