(* Copyright 2005-2007 Berke DURAK, INRIA Rocquencourt.

This file is part of Dose2.

Dose2 is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

Dose2 is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>. *)

open Fragments;;

exception Unknown_tag;;
exception Unknown_constructor;;

type io_in = Fragments.io_in
type io_out = Fragments.io_out

(*** r, w, literate, sum_io_spec, record_io_spec *)
type 'a r = io_in -> 'a
type 'a w = 'a -> io_out -> unit
type 'a literate = 'a r * 'a w
type 'a sum_io_spec = (string -> 'a r) * ('a -> string * (io_out -> unit));;
type 'a record_io_spec = (string * ('a -> io_in -> 'a) * ('a -> io_out -> unit)) list;;
type ('thing,'bag) iterator = 'bag -> ('thing -> unit) -> unit
type ('thing,'bag) builder = unit -> ('thing -> unit) * (unit -> 'bag)
type ('thing,'bag) collection_io_spec = 'thing literate * ('thing,'bag) iterator * ('thing,'bag) builder
(* ***)
(*** read, write, flush *)
let read : 'a literate -> io_in -> 'a = fun (r,_) i -> r i;;
let write : 'a literate -> 'a -> io_out -> unit = fun (_,w) x o -> w x o;;
let flush : io_out -> unit = fun o -> Fragments.flush o;;
let finish i = finish i;;
let write_and_flush io x o = write io x o; flush o;;
(* ***)
(*** convert *)
let convert ~(reader : io_in) ~(writer : io_out) () =
  let rec loop () =
    let t = read_token reader in
    (*if t <> Tag T_Stop_data then
      begin*)
        write_token writer t;
        loop ()
      (*end
    else
      flush writer*)
  in
  loop ()
;;
(* ***)
(*** dump *)
let dump ~(reader : io_in) () =
  try
    let i = ref 0 in
    while true do
      let t = read_token reader in
      Printf.printf "%d %s\n%!" !i (string_of_token t);
      incr i
    done
  with
  | End_of_file -> finish reader
;;
(* ***)
(*** io_{int,string,char,bool,float,unit,int64} *)
let io_int : int literate = (fun i -> read_int i), (fun x o -> write_int o x);;
let io_string : string literate = (fun i -> read_string i), (fun x o -> write_string o x);;
let io_char : char literate = (fun i -> read_char i), (fun x o -> write_char o x);;
let io_bool : bool literate = (fun i -> read_bool i), (fun x o -> write_bool o x);;
let io_float : float literate = (fun i -> read_float i), (fun x o -> write_float o x);;
let io_int64 : int64 literate = (fun i -> read_int64 i), (fun x o -> write_int64 o x);;
let io_unit : unit literate = (fun i -> ()), (fun x o -> ());;
(* ***)
(*** io_convert *)
let io_convert : ('a -> 'b) -> ('b -> 'a) -> 'a literate -> 'b literate =
  fun f g io ->
    (fun i -> f (read io i)),
    (fun x o -> write io (g x) o)
;;
(* ***)
(*** assoc3 *)
let rec assoc3 key = function
  | [] -> raise Not_found
  | (x,y,z)::rest when x = key -> (y,z)
  | _::rest -> assoc3 key rest
;;
(* ***)
(*** io_pair *)
let io_pair : 'a literate -> 'b literate -> ('a * 'b) literate = fun ((xr,xw) as xio) ((yr,yw) as yio) ->
  begin fun i ->
    read_this_tag i T_LParen;
    let x = read xio i in
    read_this_tag i T_Comma;
    let y = read yio i in
    read_this_tag i T_RParen;
    (x, y)
  end,
  begin fun (x,y) o ->
    write_tag o T_LParen;
    write xio x o;
    write_tag o T_Comma;
    write yio y o;
    write_tag o T_RParen
  end
;;
(* ***)
(*** io_triple *)
let io_triple : 'a literate -> 'b literate -> 'c literate -> ('a * 'b * 'c) literate =
  fun
    ((xr,xw) as xio)
    ((yr,yw) as yio)
    ((zr,zw) as zio)
  ->
  begin fun i ->
    read_this_tag i T_LParen;
    let x = read xio i in
    read_this_tag i T_Comma;
    let y = read yio i in
    read_this_tag i T_Comma;
    let z = read zio i in
    read_this_tag i T_RParen;
    (x, y, z)
  end,
  begin fun (x,y,z) o ->
    write_tag o T_LParen;
    write xio x o;
    write_tag o T_Comma;
    write yio y o;
    write_tag o T_Comma;
    write zio z o;
    write_tag o T_RParen
  end
;;
(* ***)
(*** io_quadruple *)
let io_quadruple : 'a literate -> 'b literate -> 'c literate -> 'd literate -> ('a * 'b * 'c * 'd) literate =
  fun
    ((xr,xw) as xio)
    ((yr,yw) as yio)
    ((zr,zw) as zio)
    ((tr,tw) as tio)
  ->
  begin fun i ->
    read_this_tag i T_LParen;
    let x = read xio i in
    read_this_tag i T_Comma;
    let y = read yio i in
    read_this_tag i T_Comma;
    let z = read zio i in
    read_this_tag i T_Comma;
    let t = read tio i in
    read_this_tag i T_RParen;
    (x, y, z, t)
  end,
  begin fun (x,y,z,t) o ->
    write_tag o T_LParen;
    write xio x o;
    write_tag o T_Comma;
    write yio y o;
    write_tag o T_Comma;
    write zio z o;
    write_tag o T_Comma;
    write tio t o;
    write_tag o T_RParen
  end
;;
(* ***)
(*** discard *)
let rec discard i =
  match read_token i with
  | Tag T_LBrack -> discard_list i
  | Tag T_LParen -> discard_tuple i
  | Tag T_LBrace -> discard_record i
  | Tag (T_Array|T_Hashtbl) ->
      ignore (read_int i);
      discard_list i
  | Tag T_Set -> discard_list i
  | Bool _ | Char _ | Int _ | Int64 _ | Float _ | String _ -> ()
  | t -> raise (Error(sf "Unexpected %S in discard" (string_of_token t)))
and discard_record i =
  match read_token i with
  | Tag T_RBrace -> ()
  | Field _ -> discard_record_field i
  | t -> raise (Error(sf "Unexpected %S in discard_record" (string_of_token t)))
and discard_record_field i =
  read_this_tag i T_Equal;
  discard i;
  begin
    match read_token i with
    | Tag T_RBrace -> ()
    | Tag T_Semicolon ->
        begin
          match read_token i with
          | Field _ -> discard_record_field i
          | t -> raise (Error(sf "Unexpected %S in discard_record_field after semicolon" (string_of_token t)))
        end
    | t -> raise (Error(sf "Unexpected %S in discard_record_field" (string_of_token t)))
  end
and discard_list i =
  match read_token i with
  | Tag T_RBrack -> ()
  | Tag T_Semicolon ->
      discard i;
      discard_list i
  | t -> raise (Error(sf "Unexpected %S in discard_list" (string_of_token t)))
and discard_tuple i =
  match read_token i with
  | Tag T_RParen -> ()
  | Tag T_Comma ->
      discard i;
      discard_list i
  | t -> raise (Error(sf "Unexpected %S in discard_tuple" (string_of_token t)))
(* ***)
(*** io_collection *)
let io_collection : ('thing,'bag) collection_io_spec -> 'bag literate =
  fun (xio, iterator, builder) ->
  begin fun i ->
    (*read_this_tag i T_Collection;*)
    let (add, get) = builder () in
    let rec loop () =
      match peek_token i with
      | Tag T_RBrack ->
          drop_token i;
          get ()
      | _ ->
          let x = read xio i in
          add x;
          match read_token i with
          | Tag T_RBrack -> get ()
          | Tag T_Semicolon -> loop ()
          | _ -> raise Unknown_tag
    in
    read_this_tag i T_Set;
    read_this_tag i T_LBrack;
    loop ()
  end,
  begin fun xb o ->
    write_tag o T_Set;
    write_tag o T_LBrack;
    let first = ref true in
    iterator xb
      begin fun x ->
        if !first then first := false else write_tag o T_Semicolon;
        write xio x o;
      end;
    write_tag o T_RBrack
  end
;;
(* ***)
(*** io_list *)
let io_list : 'a literate -> 'a list literate = fun ((xr,xw) as xio) ->
  begin fun i ->
    read_this_tag i T_LBrack;
    let rec loop () =
      match peek_token i with
      | Tag T_RBrack ->
          drop_token i;
          []
      | _ ->
          let x = read xio i in
          match read_token i with
          | Tag T_RBrack -> [x]
          | Tag T_Semicolon -> x :: (loop ())
          | _ -> raise Unknown_tag
    in
    loop ()
  end,
  begin fun xl o ->
    write_tag o T_LBrack;
    let rec loop semicolon = function
      | [] -> write_tag o T_RBrack
      | x :: rest ->
          if semicolon then write_tag o T_Semicolon;
          write xio x o;
          loop true rest
    in
    loop false xl
  end
;;
(* ***)
(*** io_array *)
let io_array : 'a literate -> 'a array literate = fun ((xr,xw) as xio) ->
  begin fun i ->
    read_this_tag i T_Array;
    let m = read_int i in
    read_this_tag i T_LBrack;
    let x =
      Array.init m
        begin fun j ->
          if j > 0 then read_this_tag i T_Semicolon;
          read xio i
        end
    in
    read_this_tag i T_RBrack;
    x
  end,
  begin fun x o ->
    write_tag o T_Array;
    write_int o (Array.length x);
    write_tag o T_LBrack;
    Array.iteri
      begin fun j x ->
        if j > 0 then write_tag o T_Semicolon;
        write xio x o
      end
      x;
    write_tag o T_RBrack
  end
;;
(* ***)
(*** io_hashtbl *)
let io_hashtbl : 'a literate -> 'b literate -> ('a,'b) Hashtbl.t literate = fun ((xr,xw) as xio) ((yr,yw) as yio) ->
  begin fun i ->
    read_this_tag i T_Hashtbl;
    let m = read_int i in
    let h = Hashtbl.create m in
    read_this_tag i T_LBrack;
    for j = 0 to m - 1 do
      if j > 0 then read_this_tag i T_Semicolon;
      read_this_tag i T_LParen;
      let x = read xio i in
      read_this_tag i T_Comma;
      let y = read yio i in
      read_this_tag i T_RParen;
      Hashtbl.add h x y
    done;
    read_this_tag i T_RBrack;
    h
  end,
  begin fun x o ->
    write_tag o T_Hashtbl;
    write_int o (Hashtbl.length x);
    write_tag o T_LBrack;
    let first = ref true in
    Hashtbl.iter
      begin fun x y ->
        if !first then first := false else write_tag o T_Semicolon;
        write_tag o T_LParen;
        write xio x o;
        write_tag o T_Comma;
        write yio y o;
        write_tag o T_RParen;
      end
      x;
    write_tag o T_RBrack
  end
;;
(* ***)
(*** io_record *)
let io_record : 'a record_io_spec -> 'a -> 'a literate =
  fun rl initial ->
    begin fun i ->
      read_this_tag i T_LBrace;
      let rec loop1 semicolon current =
        match read_token i with
        | Tag T_RBrace -> current
        | Tag T_Semicolon when semicolon ->
            begin
              match read_token i with
              | Field field ->
                read_this_tag i T_Equal;
                do_field current field rl
              | t -> raise (Error(sf "Expecting field in record after semicolon, not %S" (string_of_token t)))
            end
        | Field field ->
            read_this_tag i T_Equal;
            do_field current field rl
        | t -> raise (Error(sf "Expecting field in record, not %S" (string_of_token t)))
      and do_field current field = function
        | [] ->
            loss i; (* We've just seen an unknown field *)
            discard i; (* Magical discard *)
            loop1 true current
        | (field',input,_)::rest ->
            if field = field' then
              begin
                let current = input current i in
                loop1 true current
              end
            else
              do_field current field rest
      in
      loop1 false initial
    end,
    begin fun x o ->
      write_tag o T_LBrace;
      let first = ref true in
      List.iter
        begin fun (field,_,output) ->
          if !first then first := false else write_tag o T_Semicolon;
          write_field o field;
          write_tag o T_Equal;
          output x o
        end
        rl;
      write_tag o T_RBrace
    end
;;
(* ***)
(*** io_sum *)
let io_sum : 'a sum_io_spec -> 'a literate =
  fun (sr,sw) ->
  begin fun i ->
    match read_token i with
    | Constructor name ->
      begin
        try
          sr name i
        with
        | Unknown_constructor -> raise (Error(sf "Constructor %S unknown" name))
      end
    | t -> raise (Error(sf "Expecting constructor, not %S" (string_of_token t)))
  end,
  begin fun x o ->
    let (name, f) = sw x in
    write_token o (Constructor name);
    f o
  end
;;
(* ***)
(*** io_option *)
let io_option : 'a literate -> 'a option literate =
  fun ((xr,xw) as xio) ->
    io_sum
      begin
        begin function
          | "none" -> fun i -> None
          | "some" -> fun i -> Some(read xio i)
          | _      -> raise Unknown_tag
        end,
        begin function
          | None   -> "none", write io_unit ()
          | Some x -> "some", write xio x
        end
      end
;;
(* ***)
(*** io_not_implemented *)
let io_not_implemented =
  (fun _ -> assert false),
  (fun _ _ -> assert false)
;;
(* ***)
