(* RLE *)

type t = string;;
exception Invalid_compressed_data of string;;

let decompress b u =
  Buffer.clear b;
  let m = String.length u in
  let rec loop1 x s i =
    if i = m then raise (Invalid_compressed_data "EOF in integer");
    let y = Char.code u.[i] in
    if y < 128 then
      ((y lsl s) lor x, i + 1)
    else
      loop1 (((y land 127) lsl s) lor x) (s + 7) (i + 1)
  and loop2 i =
    if i = m then
      begin
        let u = Buffer.contents b in
        Buffer.clear b;
        u
      end
    else
      let (x,i) = loop1 0 0 i in
      if x land 1 = 0 then
        (* walk *)
        begin
          let x = x lsr 1 in
          if i + x > m then raise (Invalid_compressed_data "Walk too long");
          Buffer.add_substring b u i x;
          loop2 (i + x)
        end
      else
        (* run *)
        begin
          let x = x lsr 1 in
          if i >= m then raise (Invalid_compressed_data "Run too long");
          Buffer.add_string b (String.make x u.[i]);
          loop2 (i + 1)
        end
  in
  loop2 0
;;

let emit_int b x =
  let add_byte y = Buffer.add_char b (Char.chr y) in
  if x < 128 then
    add_byte x
  else if x < 16384 then
    begin
      add_byte (128 lor (x land 127));
      add_byte (x lsr 7);
    end
  else if x < 2097152 then
    begin
      add_byte (128 lor (x land 127));
      add_byte (128 lor ((x lsr 7) land 127));
      add_byte ((x lsr 14) land 127)
    end
  else (* should be enough *)
    begin
      add_byte (128 lor (x land 127));
      add_byte (128 lor ((x lsr 7) land 127));
      add_byte (128 lor ((x lsr 14) land 127));
      add_byte ((x lsr 21) land 127)
    end
;;

let compress b u =
  Buffer.clear b;
  let m = String.length u in
  let count_run i =
    let rec loop q c j =
      if j = m or u.[j] <> q then
        c
      else
        loop q (c + 1) (j + 1)
    in
    loop u.[i] 1 (i + 1)
  in
  let emit_run c q =
    if c > 0 then
      begin
        emit_int b (1 lor (c lsl 1));
        Buffer.add_char b q
      end
  in
  let emit_walk i j =
    let n = j - i in
    if n > 0 then
      begin
        emit_int b (n lsl 1);
        Buffer.add_substring b u i n
      end
  in
  let finish i =
    if i < m then emit_walk i m;
    let u = Buffer.contents b in
    Buffer.clear b;
    u
  in
  (* bytes starting from i have not yet been emitted *)
  (* we are examining byte j *)
  let rec walk i j =
    if j = m then
      finish i
    else
      let c = count_run j in
      if c < 8 then
        walk i (j + c)
      else
        begin
          emit_walk i j;
          emit_run c u.[j];
          walk (j + c) (j + c)
        end
  in
  walk 0 0
;;

let test fn =
  let ic = open_in fn in
  let b = Buffer.create 16 in
  try
    let i = ref 0 in
    while true do
      let u = input_line ic in
      incr i;
      try
        let u' = compress b u in
        let u'' = decompress b u' in
        if u <> u'' then
          begin
            Printf.printf "FAILURE line %d: %S <> %S\n" !i u u''
          end
      with
      | x -> Printf.printf "EXCEPTION line %d: %S %s\n" !i u (Printexc.to_string x)
    done;
    assert false
  with
  | End_of_file -> close_in ic
;;
