(* 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 Io;;
open Fragments;;
open Channels;;
let sf = Printf.sprintf;;

(*** CHUNK_OUT *)
module type CHUNK_OUT =
  sig
    include OUT
    val output_int : o -> int -> unit
    val output_int64 : o -> int64 -> unit
    val output_string : o -> string -> unit
    val output_float : o -> float -> unit
  end
;;
(* ***)
(*** CHUNK_IN *)
module type CHUNK_IN =
  sig
    include IN
    val input_int : i -> int
    val input_int64 : i -> int64
    val input_string : i -> string
    val input_float : i -> float
  end
;;
(* ***)
(*** Compact_in *)
module Compact_in(I : IN) =
  struct
    include I
    (*** input_int *)
    let input_int i =
      let x = I.input_byte i in
      let y =
        if x land 0x40 = 0 then
          x land 0x3f
        else
          let rec loop () =
            let x = I.input_byte i in
            if x land 0x80 = 0 then
              x land 0x7f
            else
              (x land 0x7f) lor ((loop ()) lsl 7)
          in
          (x land 0x3f) lor ((loop ()) lsl 6)
      in
      if 0 <> x land 0x80 then
        - y - 1
      else
        y
    ;;
    (* ***)
    (*** input_string *)
    let input_string i =
      let x = I.input_byte i in
      if x = 0 then
        ""
      else
        if 1 <= x && x <= 26 then
          String.make 1 (Char.chr (x + 95))
        else
          let n =
            if 27 <= x && x <= 254 then
              x - 25
            else
              input_int i
          in
          let u = String.create n in
          I.input i u 0 n;
          u
    ;;
    (* ***)
    (*** input_int64 *)
    let input_int64 i =
      let g o x = Int64.logor (Int64.shift_left (Int64.of_int (I.input_byte i)) o) x in
      let x = g 56 0L in
      let x = g 48 x in
      let x = g 40 x in
      let x = g 32 x in
      let x = g 24 x in
      let x = g 16 x in
      let x = g  8 x in
      let x = g  0 x in
      x
    ;;
    (* ***)
    (*** input_float *)
    let input_float i = Int64.float_of_bits (input_int64 i)
    ;;
    (* ***)
  end
;;
(* ***)
(*** Compact_out *)
module Compact_out(O : OUT) =
  struct
    include O
    (*** output_int *)
    let output_int o x =
      let (sign,x) =
        (if x < 0 then 0x80 else 0x00),
        (if x < 0 then - (x + 1) else x)
      in
      let f = O.output_byte o in
      if x < 64 then
        f (sign lor x)
      else
        let rec loop q = 
          if q < 128 then
            f q
          else
            begin
              f (0x80 lor (q land 0x7f));
              loop (q lsr 7)
            end
        in
        f (0x40 lor sign lor (x land 0x3f));
        loop (x lsr 6)
    (* ***)
    (*** output_string  *)
    let output_string o x =
      let f = O.output_byte o in
      let m = String.length x in
      if m = 0 then
        f 0
      else
        if m = 1 && 'a' <= x.[0] && x.[0] <= 'z' then
          f ((Char.code x.[0]) - 95)
        else
          if 2 <= m && m < 229 then
            begin
              f (m + 25);
              O.output o x 0 m
            end
          else
            begin
              f 255;
              output_int o m;
              O.output o x 0 m
            end
    (* ***)
    (*** output_int64 *)
    let output_int64 o x = (* little endian *)
      let g s = output_byte o (Int64.to_int (Int64.logand (Int64.shift_right x s) 0xffL)) in
      g 56; g 48; g 40; g 32; g 24; g 16; g 8; g 0
    (* ***)
    (*** output_float  *)
    let output_float o x =
      output_int64 o (Int64.bits_of_float x)
    (* ***)
  end
;;
(* ***)
(*** code_{bool,char,...} *)
let code_bool             = max_tag + 1
let code_char             = max_tag + 2
let code_int              = max_tag + 3
let code_float            = max_tag + 4
let code_string           = max_tag + 5
let code_int64            = max_tag + 6
let code_field            = max_tag + 7
let code_field_ref        = max_tag + 8
let code_constructor      = max_tag + 9
let code_constructor_ref  = max_tag + 10
;;
(* ***)
(*** Outputter *)
module Outputter(O : CHUNK_OUT) =
  struct
    open O;;
    (*** create_out *)
    let create_out o =
      create_io_out
        ~write_token:
          begin fun t ->
            let c = output_byte o in
            match t with
            | Bool false ->     c code_bool;   output_byte o 0
            | Bool true ->      c code_bool;   output_byte o 1
            | Char x ->         c code_int;    output_byte o (Char.code x)
            | Int x ->          c code_int;    output_int o x
            | Int64 x ->        c code_int64;  output_int64 o x
            | Float x ->        c code_float;  output_float o x
            | String x ->       c code_string; output_string o x
            | Constructor x ->  c code_constructor;  output_string o x
            | Field x ->        c code_field;  output_string o x
            | Tag x ->          c (int_of_tag x);
            | EOF ->            ()
          end
        ~flush:
          begin fun () ->
            O.flush o
          end
        ()
    ;;
    (* ***)
  end
;;
(* ***)
(*** Inputter *)
module Inputter(I : CHUNK_IN) =
  struct
    open I;;
    (*** in_of_channel *)
    let create_in i =
      let token = ref None in
      let peek_token () =
        match !token with
        | None ->
          let t =
            match input_byte i with
            | x when x <= max_tag             -> Tag(tag_of_int x)
            | x when x = code_bool            -> Bool(input_byte i = 1)
            | x when x = code_int             -> Int(input_int i)
            | x when x = code_string          -> String(input_string i)
            | x when x = code_float           -> Float(input_float i)
            | x when x = code_field           -> Field(input_string i)
            | x when x = code_constructor     -> Constructor(input_string i)
            | x                               -> raise (Error(sf "Invalid tag byte %d" x))
          in
          token := Some t;
          t
        | Some t -> t
      in
      create_io_in
        ~peek_token
        ~read_token:
          begin fun () ->
            let t = peek_token () in
            token := None;
            t
          end
        ()
    (* ***)
  end
;;
(* ***)
(*** Compressing_outputter *)
module Compressing_outputter(O : CHUNK_OUT) =
  struct
    open O;;
    (*** create_out *)
    let create_out o =
      let token_id = ref 0 in
      let dictionary = Hashtbl.create 256 in
      let c = output_byte o in
      let do_memorized code code_ref x =
        try
          let id = Hashtbl.find dictionary x in
          c code_constructor_ref;
          output_int o id
        with
        | Not_found ->
            let id = !token_id in
            incr token_id;
            Hashtbl.add dictionary x id;
            c code_constructor;
            output_string o x
      in
      create_io_out
        ~write_token:
          begin fun t ->
            match t with
            | Bool false ->    c code_bool;   output_byte o 0
            | Bool true ->     c code_bool;   output_byte o 1
            | Char x ->        c code_int;    output_byte o (Char.code x)
            | Int x ->         c code_int;    output_int o x
            | Int64 x ->       c code_int64;  output_int64 o x
            | Float x ->       c code_float;  output_float o x
            | String x ->      c code_string; output_string o x
            | Tag x ->         c (int_of_tag x);
            | Constructor x -> do_memorized code_constructor code_constructor_ref x
            | Field x ->       do_memorized code_field code_field_ref x
            | EOF ->           ()
          end
        ~flush:
          begin fun () ->
            O.flush o
          end
        ()
    ;;
    (* ***)
  end
;;
(* ***)
(*** Compressing_inputter *)
module Compressing_inputter(I : CHUNK_IN) =
  struct
    open I;;
    (*** in_of_channel *)
    let create_in i =
      let token_id = ref 0 in
      let dictionary = Hashtbl.create 256
      and reverse_dictionary = Hashtbl.create 256
      in
      let get_and_memorize () =
        let x = input_string i in
        if not (Hashtbl.mem reverse_dictionary x) then
          begin
            let id = !token_id in
            incr token_id;
            Hashtbl.add dictionary id x;
            Hashtbl.add reverse_dictionary x id;
            x
          end
        else
          raise (Error(sf "Ident %S already in dictionary." x));
      in
      let get_and_recall () =
        let id = input_int i in
        try
          Hashtbl.find dictionary id
        with
        | Not_found -> raise (Error(sf "Invalid ident reference %d" id))
      in
      let token = ref None in
      let peek_token () =
        match !token with
        | None ->
          let t =
            match input_byte i with
            | x when x <= max_tag             -> Tag(tag_of_int x)
            | x when x = code_bool            -> Bool(input_byte i = 1)
            | x when x = code_int             -> Int(input_int i)
            | x when x = code_string          -> String(input_string i)
            | x when x = code_float           -> Float(input_float i)
            | x when x = code_field           -> Field(get_and_memorize ())
            | x when x = code_constructor     -> Constructor(get_and_memorize ())
            | x when x = code_field_ref       -> Field(get_and_recall ())
            | x when x = code_constructor_ref -> Constructor(get_and_recall ())
            | x                               -> raise (Error(sf "Invalid tag byte %d" x))
          in
          token := Some t;
          t
        | Some t -> t
      in
      create_io_in
        ~peek_token
        ~read_token:
          begin fun () ->
            let t = peek_token () in
            token := None;
            t
          end
        ()
    (* ***)
  end
;;
(* ***)

module O_CHANNEL = Outputter(Compact_out(CHANNEL_OUT));;
module O_STRING  = Outputter(Compact_out(STRING_OUT));;
module I_CHANNEL = Inputter (Compact_in(CHANNEL_IN));;
module I_STRING  = Inputter (Compact_in(STRING_IN));;

module C_O_CHANNEL = Compressing_outputter(Compact_out(CHANNEL_OUT));;
module C_O_STRING  = Compressing_outputter(Compact_out(STRING_OUT));;
module C_I_CHANNEL = Compressing_inputter (Compact_in(CHANNEL_IN));;
module C_I_STRING  = Compressing_inputter (Compact_in(STRING_IN));;

let out_of_channel ~out_channel ?(compress=false) () =
  if compress then
    C_O_CHANNEL.create_out out_channel
  else
    O_CHANNEL.create_out out_channel
;;
let in_of_channel ~in_channel ?(compress=false) () =
  if compress then
    C_I_CHANNEL.create_in in_channel
  else
    I_CHANNEL.create_in in_channel
;;

let in_of_string ~str ?offset ?length ?(compress=false) () =
  let i = STRING_IN.of_string ?offset ?length str in
  if compress then
    C_I_STRING.create_in i
  else
    I_STRING.create_in i
;;

let out_to_string ?(size=128) ?(compress=false) () =
  let o = STRING_OUT.create size in
  if compress then
    (C_O_STRING.create_out o,
     fun () -> STRING_OUT.contents o)
  else
    (O_STRING.create_out o,
     fun () -> STRING_OUT.contents o)
;;
