test.ml

let ( <.> ) f g = fun x -> f (g x)

module Buffer : sig 
  type 'a t = private bytes

  and 'a rd = < rd : unit; .. > as 'a
  and 'a wr = < wr : unit; .. > as 'a
  and ro = < rd: unit; >
  and rdwr = < rd: unit; wr: unit; >

  val ro : bytes -> ro t
  val rdwr : bytes -> rdwr t
end = struct
  type 'a t = bytes
  and 'a rd = < rd : unit; .. > as 'a
  and 'a wr = < wr : unit; .. > as 'a
  and ro = < rd: unit; >
  and rdwr = < rd: unit; wr: unit; >

  let ro : bytes -> ro t = fun x -> x
  let rdwr : bytes -> rdwr t = fun x -> x
end

let failwith fmt = Format.kasprintf failwith fmt

type encode =
  | Done | Wr of { k : int -> encode; o : int; l : int; v : Buffer.ro Buffer.t }
type 'a decode =
  | Done of 'a | Rd of { k : int -> 'a decode; o : int; l : int; v : Buffer.rdwr Buffer.t }

(* GROUND *)

let io_buffer_size = 65536

type encoder =
  { payload : bytes
  ; mutable pos : int }

let make_encoder () =
  { payload= Bytes.create io_buffer_size
  ; pos= 0 }

let flush k0 encoder =
  if encoder.pos > 0
  then
    let rec k1 n =
      if n < encoder.pos
      then Wr { k= (fun m -> k1 (n + m))
              ; o= n; l= encoder.pos - n
              ; v= Buffer.ro encoder.payload }
      else ( encoder.pos <- 0 ; k0 () ) in
    k1 0
  else k0 ()

let write encoder s =
  let max = Bytes.length encoder.payload in
  let go j l encoder =
    let rem = max - encoder.pos in
    let len = if l > rem then rem else l in
    Bytes.blit_string s j encoder.payload encoder.pos len ;
    if len < l then assert false in
  go 0 (String.length s) encoder

type decoder =
  { buffer : bytes
  ; mutable pos : int
  ; mutable max : int }

let make_decoder () =
  { buffer= Bytes.create io_buffer_size
  ; pos= 0
  ; max= 0 }

let at_least_one_line decoder =
  let pos = ref decoder.pos in
  let chr = ref '\000' in
  let has_cr = ref false in
  while !pos < decoder.max
        && ( chr := Bytes.get decoder.buffer !pos
           ; not (!chr = '\n' && !has_cr) )
  do has_cr := !chr = '\r' ; incr pos done ;
  !pos < decoder.max && !chr = '\n' && !has_cr

let safe k decoder : 'v decode =
  try k decoder with exn -> assert false

let prompt k decoder =
  if decoder.pos > 0
  then
    ( let rest = decoder.max - decoder.pos in
      Bytes.blit decoder.buffer decoder.pos decoder.buffer 0 rest ;
      decoder.max <- rest ;
      decoder.pos <- 0 ) ;
  let rec go off =
    if off = Bytes.length decoder.buffer
    then assert false
    else if not (at_least_one_line { decoder with max= off })
    then Rd { k= (fun len -> go (off + len))
            ; o= off; l= Bytes.length decoder.buffer - off
            ; v= Buffer.rdwr decoder.buffer }
    else
      ( decoder.max <- off
      ; safe k decoder ) in
  go decoder.max

let end_of_input decoder = decoder.max

let peek_while_eol decoder =
  let idx = ref decoder.pos in
  let chr = ref '\000' in
  let has_cr = ref false in

  while !idx < end_of_input decoder
        && ( chr := Bytes.get decoder.buffer !idx
           ; not (!chr = '\n' && !has_cr) )
  do has_cr := !chr = '\r' ; incr idx done ;

  if !idx < end_of_input decoder && !chr = '\n' && !has_cr
  then decoder.buffer, decoder.pos, !idx + 1 - decoder.pos
  else assert false

let line k decoder =
  let raw_crlf, off, len = peek_while_eol decoder in
  k (Bytes.sub_string raw_crlf off (len - 2))

let decode decoder k =
  if at_least_one_line decoder
  then safe (line k) decoder
  else prompt (line k) decoder

type context =
  { encoder : encoder
  ; decoder : decoder }

let make_context () =
  { encoder= make_encoder ()
  ; decoder= make_decoder () }

type value =
  [ `EHLO | `PP_250 ]

let encode encoder v : encode =
  let return () : encode = Done in
  match v with
  | `EHLO -> write encoder "EHLO\r\n" ; flush return encoder
  | `PP_250 -> write encoder "250\r\n" ; flush return encoder

let decode decoder : value decode =
  let return v : _ decode = Done v in
  decode decoder @@ function
  | "EHLO" -> return `EHLO
  | "250" -> return `PP_250
  | _ -> assert false

module Scheduler = struct
  type 'a t =
    | Rd of { k : int -> 'a t
            ; o : int; l : int
            ; v : Buffer.rdwr Buffer.t }
    | Wr of { k : int -> 'a t
            ; o : int; l : int
            ; v : Buffer.ro Buffer.t }
    | Done of 'a

  let return x = Done x

  let rec go ~f m len = match m len with
    | Done v -> f v
    | Rd { k; o; l; v; } -> Rd { k= go ~f k; o; l; v; }
    | Wr { k; o; l; v; } -> Wr { k= go ~f k; o; l; v; }

  let bind : 'a t -> f:('a -> 'b t) -> 'b t = fun m ~f -> match m with
    | Done v -> f v
    | Rd { k; o; l; v; } ->
      Rd { k= go ~f k; o; l; v; }
    | Wr { k; o; l; v; } ->
      Wr { k= go ~f k; o; l; v; }

  let encode
    : type a. context -> value -> (context -> 'r t) -> 'r t
    = fun ctx x k ->
      let rec go : encode -> 'a t = function
        | Wr { k; o; l; v; } ->
          let continue n = go (k n) in
          Wr { k= continue; o; l; v; }
        | Done -> k ctx in
      (go <.> encode ctx.encoder) x

  let decode
    : context -> (context -> value -> 'r t) -> 'r t
    = fun ctx k ->
      let rec go : value decode -> 'a t = function
        | Rd { k; o; l; v; } ->
          let continue n = go (k n) in
          Rd { k= continue; o; l; v; }
        | Done v -> k ctx v in
      go (decode ctx.decoder)
end

type 'a linear = { linear : 'a }
type 'a data = { data : 'a }

type (_, _, _, _) lens =
  | Zero : ('a, 'b, [ `cons of 'a * 'xs ], [ `cons of 'b * 'xs ]) lens
  | Succ : ('x, 'y, 'xs, 'ys) lens -> ('x, 'y, [ `cons of 'a * 'xs ], [ `cons of 'a * 'ys ]) lens
  | Other : ('xs -> 'x) * ('xs -> 'y -> 'ys) -> ('x, 'y, 'xs, 'ys) lens

let rec get
  : type x y xs ys. (x, y, xs, ys) lens -> xs -> x
  = fun l xs -> match l, xs with
    | Zero, (`cons (hd, _)) -> hd
    | Succ l, (`cons (_, tl)) -> get l tl
    | Other (get, _), xs -> get xs

let rec put
  : type x y xs ys. (x, y, xs, ys) lens -> xs -> y -> ys
  = fun l xs b -> match l, xs with
  | Zero, (`cons (_, tl)) -> `cons (b, tl)
  | Succ l, (`cons (hd, tl)) -> `cons (hd, put l tl b)
  | Other (_, put), xs -> put xs b

module type SCHEDULER = sig
  type 'a t

  val bind : 'a t -> f:('a -> 'b t) -> 'b t
  val return : 'a -> 'a t
end

module type MONAD = sig
  module Scheduler : SCHEDULER

  type 'a s = 'a Scheduler.t
  type ('p, 'q, 'a) t = { m : ('p -> ('q * 'a) s) }

  val return : 'a -> ('p, 'p, 'a data) t
  val bind : ('p, 'q, 'a data) t -> f:('a -> ('q, 'r, 'b) t) -> ('p, 'r, 'b) t
  val (>>=) : ('p, 'q, 'a data) t -> ('a -> ('q, 'r, 'b) t) -> ('p, 'r, 'b) t
  val (>>) : ('p, 'q, 'a data) t -> ('q, 'r, 'b) t -> ('p, 'r, 'b) t
  val run : f:(unit -> (unit, unit, 'a data) t) -> 'a s
end

module type LENS = sig
  module Monad : MONAD

  type all_empty = [ `cons of unit * 'xs ] as 'xs

  val ( @> ) : ('p, 'q, 'a) Monad.t -> ('p, 'q, 'pr, 'ps) lens -> ('pr, 'ps, 'a) Monad.t
  val ( <@ ) : ('p, 'q, 'pr, 'ps) lens -> ('p, 'q, 'a) Monad.t -> ('pr, 'ps, 'a) Monad.t

  val l0 : ('a, 'b, [ `cons of 'a * 'xs ], [ `cons of 'b * 'xs ]) lens
  val l1 : ('a, 'b, [ `cons of 'x1 * [ `cons of 'a * 'xs ] ], [ `cons of 'x1 * [ `cons of 'b * 'xs ] ]) lens
  val l2 : ('a, 'b, [ `cons of 'x1 * [ `cons of 'x2 * [ `cons of 'a * 'xs ] ] ],
                    [ `cons of 'x1 * [ `cons of 'x2 * [ `cons of 'b * 'xs ] ] ]) lens
  val run : ('a -> (all_empty, all_empty, 'b data) Monad.t) -> 'a -> 'b Monad.s

  val extend : ('p, [ `cons of unit * 'p ], unit data) Monad.t
  val shrink : ([ `cons of unit * 'p ], 'p, unit data) Monad.t
end

module type LINEAR = sig
  module Monad : MONAD

  type 'f bind

  val ( >>- ) : ('p, 'q, 'a linear) Monad.t -> ('a linear -> ('q, 'r, 'b) Monad.t) bind -> ('p, 'r, 'b) Monad.t

  val put_linear : (unit, 'a linear, 'q, 'r) lens -> ('p, 'q, 'a linear) Monad.t -> ('p, 'r, unit data) Monad.t
  val put_linear_val : (unit, 'a linear, 'p, 'q) lens -> 'a -> ('p, 'q, unit data) Monad.t
  val get_linear : ('a linear, unit, 'p, 'q) lens -> ('p, 'q, 'a linear) Monad.t
  val return_linear : 'a -> ('p, 'p, 'a linear) Monad.t

  val bind : 'f -> 'f bind
end

module Monad (Scheduler : SCHEDULER) : MONAD with module Scheduler = Scheduler = struct
  module Scheduler = Scheduler

  type 'a s = 'a Scheduler.t
  type ('p, 'q, 'a) t = { m : ('p -> ('q * 'a) s) }

  let return a = { m= fun p -> Scheduler.return (p, { data= a }) }
  let bind m ~f = { m= fun p -> Scheduler.bind (m.m p) (fun (q, { data= a }) -> (f a).m q) }
  let (>>=) x f = bind x ~f
  let (>>) m1 m2 = { m= fun p -> Scheduler.bind (m1.m p) (fun (q, _) -> m2.m q) }
  let run ~f = Scheduler.bind ((f ()).m ()) ~f:(fun (_, { data }) -> Scheduler.return data)
end

module Lens (Scheduler : SCHEDULER) (Monad : MONAD with module Scheduler = Scheduler) : LENS with module Monad := Monad = struct
  type all_empty = [ `cons of unit * 'xs ] as 'xs

  let ( @> )
    : 'p 'q 'pr 'ps 'a. ('p, 'q, 'a) Monad.t -> ('p, 'q, 'pr, 'ps) lens -> ('pr, 'ps, 'a) Monad.t
    = fun m l ->
      { m= fun p -> Scheduler.bind (m.m (get l p)) ~f:(fun (q, a) -> Scheduler.return (put l p q, a)) }

  let ( <@ ) l m = m @> l

  let l0 = Zero
  let l1 = Succ Zero
  let l2 = Succ (Succ Zero)

  let run = fun f x ->
    let rec all_empty = `cons ((), all_empty) in
    Scheduler.bind ((f x).Monad.m all_empty) ~f:(fun (_, { data }) -> Scheduler.return data)

  let extend : 'pr. ('pr, [ `cons of unit * 'pr ], unit data) Monad.t
    = { m= fun pr -> Scheduler.return (`cons ((), pr), { data= () }) }

  let shrink : type pr. ([ `cons of unit * pr ], pr, unit data) Monad.t
    = { m= fun (`cons ((), pr)) -> Scheduler.return (pr, { data= () }) }
end

module Linear
  (Scheduler : SCHEDULER)
  (Monad : MONAD with module Scheduler = Scheduler) : LINEAR with module Monad := Monad = struct
  type 'f bind = 'f

  let ( >>- )
    : 'p 'q 'a 'r 'b. ('p, 'q, 'a linear) Monad.t -> ('a linear -> ('q, 'r, 'b) Monad.t) bind -> ('p, 'r, 'b) Monad.t
    = fun m f -> { m= fun p -> Scheduler.bind (m.Monad.m p) (fun (q, x) -> (f x).m q) }

  let put_linear
    : 'a 'q 'r 'p. (unit, 'a linear, 'q, 'r) lens -> ('p, 'q, 'a linear) Monad.t -> ('p, 'r, unit data) Monad.t
    = fun l m -> { m= fun p -> Scheduler.bind (m.m p) ~f:(fun (q, v) -> Scheduler.return (put l q v, { data= () })) }

  let put_linear_val
    : 'a 'p 'q. (unit, 'a linear, 'p, 'q) lens -> 'a -> ('p, 'q, unit data) Monad.t
    = fun l v ->
      { m= fun p -> Scheduler.return (put l p { linear= v }, { data= () }) }

  let get_linear
    : 'a 'p 'q. ('a linear, unit, 'p, 'q) lens -> ('p, 'q, 'a linear) Monad.t
    = fun l -> { m= fun p -> Scheduler.return (put l p (), get l p) }

  let return_linear
    : 'a 'p. 'a -> ('p, 'p, 'a linear) Monad.t
    = fun v -> { m= fun p -> Scheduler.return (p, { linear= v }) }

  let bind : 'f. 'f -> 'f bind = fun f -> f
end

module Make (Scheduler : SCHEDULER) = struct
  module Monad = Monad (Scheduler)

  include Monad
  include Lens(Scheduler)(Monad)
  include Linear(Scheduler)(Monad)
end

type ehlo = EHLO
type auth = AUTH
type password = PASSWORD
type recipient = RECIPIENT

type pp_250 = PP_250
type tp_334 = TP_334

type authentication =
  | PP_235
  | PN_500
  | PN_501
  | PN_535

module Protocol = struct
  type 'a t =
    | EHLO : ehlo t
    | AUTH : auth t
    | PASSWORD : password t
    | PP_250 : pp_250 t
    | TP_334 : tp_334 t
    | Authentication_result : authentication t
    | MAIL_FROM : recipient t

    | Magic : 'a t

  let to_value : type a. a t -> a -> value = fun w v -> assert false
  let of_value : type a. a t * value -> a = fun (_w, _v) -> assert false
end

module Channel : sig
  type 'a t = 'a Protocol.t

  val encode : context -> 'a t -> 'a -> unit Scheduler.t
  val decode : context -> 'a t -> 'a Scheduler.t
end = struct
  type 'a t = 'a Protocol.t

  let ( <.> ) f g = fun x -> f (g x)

  let encode ctx w v =
    Scheduler.encode ctx (Protocol.to_value w v) (fun ctx -> Scheduler.return ())
  let decode ctx w =
    Scheduler.decode ctx (fun ctx v ->
      let v = Protocol.of_value (w, v) in
      Scheduler.return v)
end

module Ground = Make (Scheduler)

module Session = struct
  let unlinear { linear } = linear

  type ('v, 's) send = 'v Protocol.t * 's linear
  type ('v, 's) recv = 'v Protocol.t * 's linear
  type close = Close

  type 'a prot = 'a

  let (@?>)
    : 'v Protocol.t -> ('s * 'c) prot -> (('v, 's) send * ('v, 'c) recv) prot
    = fun w (s, c) ->
      ((w, { linear= s }), (w, { linear = c }))

  let (@!>)
    : 'v Protocol.t -> ('s * 'c) prot -> (('v, 's) recv * ('v, 's) send) prot
    = fun a b -> a @?> b

  let finish : (close * close) prot = Close, Close

  let send
    : context -> (('v, 's) send linear, unit, 'p, 'q) lens -> 'v -> ('p, 'q, 's linear) Ground.t
    = fun ctx l v ->
      { m= (fun p ->
            let w, cont = unlinear (get l p) in
            let encode = Channel.encode ctx w v in
            Scheduler.bind encode ~f:(fun () -> Scheduler.return (put l p (), cont))) }

  let receive
    : context -> (('v, 's) recv linear, unit, 'p, 'q) lens -> ('p, 'q, ('s linear * 'v data) linear) Ground.t
    = fun ctx l ->
      { m= (fun p ->
            let w, cont = unlinear (get l p) in
            let decode = Channel.decode ctx w in
            Scheduler.bind decode ~f:(fun v ->
                let s = put l p () in
                let v = { linear= (cont, { data= v; }) } in
                Scheduler.return (s, v))) }

  let close
    : (close linear, unit, 'p, 'q) lens -> ('p, 'q, unit data) Ground.t
    = fun l -> { m= fun p -> Scheduler.return (put l p (), { data= () }) }

  let accept
    : context -> ('s * 'c) Channel.t -> ('p, 'p, 's linear) Ground.t
    = fun ctx w -> { m= fun p ->
        let decode = Channel.decode ctx w in
        Scheduler.bind decode ~f:(fun (s, _) -> Scheduler.return (p, { linear= s })) }

  let connect
    : context -> ('s * 'c) Channel.t -> ('s * 'c) -> ('p, 'p, 'c linear) Ground.t
    = fun ctx w (s, c) ->
      { m= fun p ->
           let encode = Channel.encode ctx w (s, c) in
           Scheduler.bind encode ~f:(fun () -> Scheduler.return (p, { linear= c })) }
end

module SList = struct
  type 'a w =
    | [] : 'a w
    | (::) : 'a linear * 'a t -> 'a w
  and 'a t = 'a w linear
end

module SEither = struct
  type ('a, 'b) w =
    | L : 'a linear -> ('a, 'b) w
    | R : 'b linear -> ('a, 'b) w
  and ('a, 'b) t = ('a, 'b) w linear
end

let run_client ctx w recipients =
  let rec go = function
    | Scheduler.Rd { k; o; l; v; } ->
      let n = input stdin (v :> bytes) o l in go (k n)
    | Scheduler.Wr { k; o; l; v; } ->
      output stdout (v :> bytes) o l ; go (k l)
    | Scheduler.Done v -> v in
  
  let open Ground in

  let s = l0 in

  let m =
    let ( let* ) x f = x >>- (bind (fun { linear= w } -> put_linear_val s w >> f s)) in
    let ( let+ ) x f = x >>- (bind (fun { linear= { linear= w }, { data= v } } -> put_linear_val s w >> f (s, v))) in

    let* s = Session.connect ctx w
        Session.(Protocol.EHLO
                 @?> Protocol.PP_250
                 @!> Protocol.AUTH
                 @?> Protocol.TP_334
                 @!> Protocol.PASSWORD
                 @?> Protocol.Authentication_result
                 @!> finish) in
    let* s = Session.send ctx s EHLO in
    let+ s, PP_250 = Session.receive ctx s in
    let* s = Session.send ctx s AUTH in
    let+ s, TP_334 = Session.receive ctx s in
    let* s = Session.send ctx s PASSWORD in
    let+ s, v = Session.receive ctx s in
    match v with
    | PP_235 ->
      Session.close s
    | PN_500 ->
      Session.close s
    | PN_501 ->
      Session.close s
    | PN_535 ->
      Session.close s in
  
  let fiber = run (fun m -> return ()) m in
  go fiber