src/utils/cdk/bzip2.ml

   1
   2 (* Module [Bzip2]: reading and writing to/from [bzip2] compressed files *)
   3
   4 exception Error of string
   5
   6 let buffer_size = 1024
   7
   8 type in_channel =
   9   { in_chan: Pervasives.in_channel;
  10     in_buffer: string;
  11     mutable in_pos: int;
  12     mutable in_avail: int;
  13     mutable in_eof: bool;
  14     in_stream: Bzlib.stream;
  15     mutable in_size: int32; }
  16
  17 let open_in_chan ic =
  18   { in_chan = ic;
  19     in_buffer = String.create buffer_size;
  20     in_pos = 0;
  21     in_avail = 0;
  22     in_eof = false;
  23     in_stream = Bzlib.decompress_init 0 false;
  24     in_size = Int32.zero }
  25
  26 let open_in filename =
  27   let ic = Pervasives.open_in_bin filename in
  28   try
  29     open_in_chan ic
  30   with e -> Pervasives.close_in ic; raise e
  31
  32 let read_byte iz =
  33   if iz.in_avail = 0 then begin
  34     let n = Pervasives.input iz.in_chan iz.in_buffer 0
  35                              (String.length iz.in_buffer) in
  36     if n = 0 then raise End_of_file;
  37     iz.in_pos <- 0;
  38     iz.in_avail <- n
  39   end;
  40   let c = iz.in_buffer.[iz.in_pos] in
  41   iz.in_pos <- iz.in_pos + 1;
  42   iz.in_avail <- iz.in_avail - 1;
  43   Char.code c
  44
  45 let read_int32 iz =
  46   let b1 = read_byte iz in
  47   let b2 = read_byte iz in
  48   let b3 = read_byte iz in
  49   let b4 = read_byte iz in
  50   Int32.logor (Int32.of_int b1)
  51     (Int32.logor (Int32.shift_left (Int32.of_int b2) 8)
  52       (Int32.logor (Int32.shift_left (Int32.of_int b3) 16)
  53                    (Int32.shift_left (Int32.of_int b4) 24)))
  54
  55 let rec input iz buf pos len =
  56   if pos < 0 || len < 0 || pos + len > String.length buf then
  57     invalid_arg "Bzip2.input";
  58   if iz.in_eof then 0 else begin
  59     if iz.in_avail = 0 then begin
  60       let n = Pervasives.input iz.in_chan iz.in_buffer 0
  61                                (String.length iz.in_buffer) in
  62       if n = 0 then raise(Error("truncated file"));
  63       iz.in_pos <- 0;
  64       iz.in_avail <- n
  65     end;
  66     let (finished, used_in, used_out) =
  67       try
  68         Bzlib.decompress iz.in_stream iz.in_buffer iz.in_pos iz.in_avail
  69                                    buf pos len
  70       with Bzlib.Error(_, e) ->
  71         raise(Error(Bzlib.string_of_error e)) in
  72     iz.in_pos <- iz.in_pos + used_in;
  73     iz.in_avail <- iz.in_avail - used_in;
  74     iz.in_size <- Int32.add iz.in_size (Int32.of_int used_out);
  75     if finished then begin
  76       iz.in_eof <- true;
  77       used_out
  78     end else if used_out = 0 then
  79       input iz buf pos len
  80     else
  81       used_out
  82   end
  83
  84 let rec really_input iz buf pos len =
  85   if len <= 0 then () else begin
  86     let n = input iz buf pos len in
  87     if n = 0 then raise End_of_file;
  88     really_input iz buf (pos + n) (len - n)
  89   end
  90
  91 let char_buffer = String.create 1
  92
  93 let input_char iz =
  94   if input iz char_buffer 0 1 = 0 then raise End_of_file else char_buffer.[0]
  95
  96 let input_byte iz =
  97   Char.code (input_char iz)
  98
  99 let dispose iz =
 100   iz.in_eof <- true;
 101   Bzlib.decompress_end iz.in_stream
 102
 103 let close_in iz =
 104   dispose iz;
 105   Pervasives.close_in iz.in_chan
 106
 107 type out_channel =
 108   { out_chan: Pervasives.out_channel;
 109     out_buffer: string;
 110     mutable out_pos: int;
 111     mutable out_avail: int;
 112     out_stream: Bzlib.stream;
 113     mutable out_size: int32; }
 114
 115 let open_out_chan ?(level = 6) oc =
 116   if level < 1 || level > 9 then invalid_arg "Bzip2.open_out: bad level";
 117   { out_chan = oc;
 118     out_buffer = String.create buffer_size;
 119     out_pos = 0;
 120     out_avail = buffer_size;
 121     out_stream = Bzlib.compress_init level 0 0;
 122     out_size = Int32.zero }
 123
 124 let open_out ?(level = 6) filename =
 125   open_out_chan ~level (Pervasives.open_out_bin filename)
 126
 127 let rec output oz buf pos len =
 128   if pos < 0 || len < 0 || pos + len > String.length buf then
 129     invalid_arg "Bzlib2.output";
 130   (* If output buffer is full, flush it *)
 131   if oz.out_avail = 0 then begin
 132   (* Printf.printf "Flushing out_avail\n"; *)
 133     Pervasives.output oz.out_chan oz.out_buffer 0 oz.out_pos;
 134     oz.out_pos <- 0;
 135     oz.out_avail <- String.length oz.out_buffer
 136   end;
 137   let (_, used_in, used_out) =
 138     try
 139       Bzlib.compress oz.out_stream buf pos len
 140                                  oz.out_buffer oz.out_pos oz.out_avail
 141                                  Bzlib.BZ_RUN
 142     with Bzlib.Error(f, e) ->
 143       raise (Error(Bzlib.string_of_error e)) in
 144   oz.out_pos <- oz.out_pos + used_out;
 145   oz.out_avail <- oz.out_avail - used_out;
 146   oz.out_size <- Int32.add oz.out_size (Int32.of_int used_in);
 147   if used_in < len then output oz buf (pos + used_in) (len - used_in)
 148
 149 let output_char oz c =
 150   char_buffer.[0] <- c;
 151   output oz char_buffer 0 1
 152
 153 let output_byte oz b =
 154   output_char oz (Char.unsafe_chr b)
 155
 156 let flush oz =
 157   let rec do_flush () =
 158     (* If output buffer is full, flush it *)
 159     if oz.out_avail = 0 then begin
 160       Pervasives.output oz.out_chan oz.out_buffer 0 oz.out_pos;
 161       oz.out_pos <- 0;
 162       oz.out_avail <- String.length oz.out_buffer
 163     end;
 164     let (finished, _, used_out) =
 165       Bzlib.compress oz.out_stream oz.out_buffer 0 0
 166                                  oz.out_buffer oz.out_pos oz.out_avail
 167                                  Bzlib.BZ_FINISH in
 168     oz.out_pos <- oz.out_pos + used_out;
 169     oz.out_avail <- oz.out_avail - used_out;
 170     if not finished then do_flush() in
 171   do_flush();
 172   (* Final data flush *)
 173   if oz.out_pos > 0 then
 174     Pervasives.output oz.out_chan oz.out_buffer 0 oz.out_pos;
 175   Bzlib.compress_end oz.out_stream
 176
 177 let close_out oz =
 178   flush oz;
 179   Pervasives.close_out oz.out_chan