Huffman Tree

encode-decode.ml

(* This contains the meat of the algorithm *)

(* Returns a list of all encodings stored in hm *)
let get_codes (hm : huffmantree) : encoding list =
  let rec fold_tree tree enc acc =
    match tree with
      | _, Leaf(c) -> (c, List.rev(enc)) :: acc
      | _, Node(l, r) ->
        let a = (fold_tree r (1::enc) acc) in
        (fold_tree l (0::enc) a) in
  fold_tree hm [] []

(* Returns the longest encoding in encl*)
let longest_code (encl : encoding list) : int list =
  let sorted_enc = List.sort (fun (_,a) (_,b) -> (List.length b) - (List.length a)) encl in
  snd(List.hd(sorted_enc))

(* Returns the bit vector in encl that encodes chr *)
let lookup_code (chr: char) (encl : encoding list) : int list =
  snd (List.find (fun ((c:char), _) -> c = chr) encl)

(* Returns the huffman tree created from cl and the encoded bit stream*)
let encode (cl : char list) : huffmantree * int list=
  (* Helper function for generating a range of numbers *)
  let rec range (a : int) (b : int) : int list = 
    if a > b then [] else a :: (range (a+1) b) in
  
  (* Our frequencies in an array associating char ascii with index *)
  let freq_assoc : int array = countchars cl in
  
  (* We place our nodes of the tree in a fake priority queue, aka a list, starting with the leaves *)
  let nodes_pq : (int * huffmantree) list = List.fold_left 
    (fun (acc : (int * huffmantree) list) (i : int) -> 
      (* Frequency, huffmantree Leaf with index = character ascii number*)
      (freq_assoc.(i), (i,Leaf(char_of_int i)))::acc) [] (range 0 255) in
  
  (* A recursive function that takes in an initial list of frequency*tree *)
  (* and builds a complete tree. It also includes a sanity check.*)
  let rec build_tree (pq : (int * huffmantree) list) (i : int) : (int * huffmantree) =
    (* Sort our list so that it's in ascending order (lower frequencies first) *)
    let pq = List.sort (fun ((a : int), _) ((b : int), _) -> a - b) pq in
    match pq with
      (* PQ should never be empty since we take out 2 and put in 1 node *)
      | [] -> failwith "The impossible happened"
      (* PQ contains only a single element, which must be the root (index = 510) *)
      | [t] -> if i = 511 then t else (failwith "Invariant violated.")
      (* We take the two smallest leaves/nodes and combine them, *)
      (* we then put them back into pq and call buildtree with the new node in pq *)
      | (a,l) :: (b,r) :: tl -> 
        build_tree ((a+b, (i,Node(l,r))) :: tl) (i+1) in
  
  (* Our root huffmantree *)
  let htree : huffmantree = snd (build_tree nodes_pq 256) in
  
  (* Our encoding list holds the association list for lookup *)
  let enc : encoding list = get_codes htree in
  
  (* A helper function for fold_left that generates the encoding bitstream *)
  let fold (acc : int list) (c : char) : int list = 
    let rec helper t =
      match t with
        | [] -> acc
        | hd :: tl -> hd :: (helper tl) in
    helper (List.rev (lookup_code c enc)) in
  
  htree, List.rev(List.fold_left fold [] cl)

(* Returns the char list decoded from bitstream, based on the bit vectors
    in hm *)
let decode (hm : huffmantree) (bitstream: int list) : char list =
  (* Helper function for tree traversal of hm *)
  let rec traverse (tree : huffmantree) (stream : int list) : char * int list =
    match tree with 
      | _, Leaf(c) -> c, stream
      | _, Node(l, r) ->
        match stream with
          | hd :: tl -> 
            (* If we encounter a 0, then branch left, else right *)
            if hd = 0 then traverse l tl else traverse r tl
          | [] -> char_of_int 0, [] in
  
  (* Helper function to keep the stream going *)
  let rec fold (stream : int list) (acc : char list) : char list =
    match stream with 
      | [] -> acc
      | _ -> 
        let c, tl = traverse hm stream in
        fold tl (c::acc) in
        
  List.rev (fold bitstream [])

(* Returns bitstream with hm compressed and stored on the beginning of the
   stream in preorder *)
let prepend_tree (hm : huffmantree) (bitstream: int list) : int list =
  (* Traverses the tree and appends 510*9 bits representing the tree onto bitstream *)
  (* New function because I wasn't sure if we're allowed to change function to rec *)
  let rec fold_tree (tree : huffmantree) (acc : int list) : int list =
    match tree with
      | i, Leaf(_) -> to_bits 9 acc i
      | i, Node(l, r) ->
        let a = fold_tree r acc in
        to_bits 9 (fold_tree l a) i in
        
  fold_tree hm bitstream

(* Returns bitstream with the encoding for a huffman tree removed from the
   beginning and decoded.
   Requires: a valid encoding for a huffman tree at the beginning of
   bitstream *)
let regrow_tree (bitstream : int list) : huffmantree * int list =
  (* Iterates over the bitstream and given the assumption that tree is*)
  (* stored in preorder and Leaf terminated, reconstructs our huffman tree *)
  let rec maketree (stream : int list) : huffmantree * int list =
    let index, tl = read_n_bits_to_int 9 stream in
    if index < 256 then
      (index, Leaf(char_of_int index)), tl
    else
      (* First grow out the left side *)
      let l, t = maketree tl in
      (* The grow out the right side *)
      let r, t2 = maketree t in
      (index, Node(l, r)), t2 in
  
  maketree bitstream

huffman.ml

(* Full program to encode and decode *)

(* huffman tree consists of an index and a node of either a Leaf or a binary Node of 2 subtrees *)
type huffmantree = int*node and node = Leaf of char | Node of huffmantree*huffmantree
type encoding = char * (int list)

(* Returns a char list of the contents of the file at fname *)
let load_chars (fname : string) : char list =
  let explode (s : string) : char list =
    let rec addtolist n acc =
      if n < 0 then acc else
        addtolist (n - 1) (s.[n] :: acc) in
    addtolist (String.length s - 1) ['\n'] in
  let file = open_in fname in
  let stringlist =
    try
      let rec read_lines acc =
        try read_lines (input_line file :: acc) with End_of_file -> acc in
      read_lines []
    with exc -> close_in file; raise exc in
  List.concat (List.rev_map explode stringlist)

let to_bits (n : int) (bitstream : int list) (num : int) : int list =
  let rec each_bit accum i =
    if i < n then
      each_bit (((num lsr i) land 1) :: accum) (i + 1)
    else accum in
  each_bit bitstream 0

(* Returns bitstream with the first n bits removed and converted back
 * to an integer *)
let read_n_bits_to_int (n : int) (bitstream : int list) : int * int list =
  let rec pop_n i word lst =
    if i = n then (word, lst)
    else match lst with
        h :: t -> pop_n (i + 1) ((word lsl 1) lor h) t
      | [] -> failwith "not enough bits remaining"
  in pop_n 0 0 bitstream

(* Write a stream of bits to a byte-oriented out channel.*)
(* Pad last byte to length 8 with the supplied padding. *)
let write_bits (f : out_channel) (bits : int list) (padding : int list) : unit =
  assert (List.length padding >= 8); (* sanity check *)
  let next_byte bits =
    let rec nb n byte bits =
      if n = 8 then (byte, bits)
      else match bits with
        | [] -> (fst (nb n byte padding), [])
        | h :: t -> nb (n + 1) (byte lsl 1 lor h) t
    in nb 0 0 bits in
  let rec write_char bits =
    if bits = [] then () else
    let (byte, bits) = next_byte bits in
    output_byte f byte; write_char bits in
  write_char bits;
  close_out f

let write_chars (f : out_channel) (chrs : char list) : unit =
  List.iter (output_char f) chrs

(* Returns the contents of f as a bitstream *)
let read_bits (f : in_channel) : int list =
  (* reversed list of bytes *)
  let rec read_bytes bytes =
    let res = try Some (input_byte f)
      with End_of_file -> None in
    match res with
      None -> bytes
    | Some(b) -> read_bytes (b :: bytes) in
  List.fold_left (to_bits 8) [] (read_bytes [])


(* Returns an array with the frequency of appearance of each character
   in cl stored in its ASCII-based index *)
let countchars (cl : char list) : int array =
  let ar = Array.make 256 0 in
  List.iter (fun elt -> ar.(int_of_char elt) <- (ar.(int_of_char elt) + 1)) cl;
  ar

(* Returns a list of all encodings stored in hm *)
let get_codes (hm : huffmantree) : encoding list =
  let rec fold_tree tree enc acc =
    match tree with
      | _, Leaf(c) -> (c, List.rev(enc)) :: acc
      | _, Node(l, r) ->
        let a = (fold_tree r (1::enc) acc) in
        (fold_tree l (0::enc) a) in
  fold_tree hm [] []

(* Returns the longest encoding in encl*)
let longest_code (encl : encoding list) : int list =
  let sorted_enc = List.sort (fun (_,a) (_,b) -> (List.length b) - (List.length a)) encl in
  snd(List.hd(sorted_enc))

(* Returns the bit vector in encl that encodes chr *)
let lookup_code (chr: char) (encl : encoding list) : int list =
  snd (List.find (fun ((c:char), _) -> c = chr) encl)

(* Returns the huffman tree created from cl and the encoded bit stream*)
let encode (cl : char list) : huffmantree * int list=
  (* Helper function for generating a range of numbers *)
  let rec range (a : int) (b : int) : int list = 
    if a > b then [] else a :: (range (a+1) b) in
  
  (* Our frequencies in an array associating char ascii with index *)
  let freq_assoc : int array = countchars cl in
  
  (* We place our nodes of the tree in a fake priority queue, aka a list, starting with the leaves *)
  let nodes_pq : (int * huffmantree) list = List.fold_left 
    (fun (acc : (int * huffmantree) list) (i : int) -> 
      (* Frequency, huffmantree Leaf with index = character ascii number*)
      (freq_assoc.(i), (i,Leaf(char_of_int i)))::acc) [] (range 0 255) in
  
  (* A recursive function that takes in an initial list of frequency*tree *)
  (* and builds a complete tree. It also includes a sanity check.*)
  let rec build_tree (pq : (int * huffmantree) list) (i : int) : (int * huffmantree) =
    (* Sort our list so that it's in ascending order (lower frequencies first) *)
    let pq = List.sort (fun ((a : int), _) ((b : int), _) -> a - b) pq in
    match pq with
      (* PQ should never be empty since we take out 2 and put in 1 node *)
      | [] -> failwith "The impossible happened"
      (* PQ contains only a single element, which must be the root (index = 510) *)
      | [t] -> if i = 511 then t else (failwith "Invariant violated.")
      (* We take the two smallest leaves/nodes and combine them, *)
      (* we then put them back into pq and call buildtree with the new node in pq *)
      | (a,l) :: (b,r) :: tl -> 
        build_tree ((a+b, (i,Node(l,r))) :: tl) (i+1) in
  
  (* Our root huffmantree *)
  let htree : huffmantree = snd (build_tree nodes_pq 256) in
  
  (* Our encoding list holds the association list for lookup *)
  let enc : encoding list = get_codes htree in
  
  (* A helper function for fold_left that generates the encoding bitstream *)
  let fold (acc : int list) (c : char) : int list = 
    let rec helper t =
      match t with
        | [] -> acc
        | hd :: tl -> hd :: (helper tl) in
    helper (List.rev (lookup_code c enc)) in
  
  htree, List.rev(List.fold_left fold [] cl)

(* Returns the char list decoded from bitstream, based on the bit vectors
    in hm *)
let decode (hm : huffmantree) (bitstream: int list) : char list =
  (* Helper function for tree traversal of hm *)
  let rec traverse (tree : huffmantree) (stream : int list) : char * int list =
    match tree with 
      | _, Leaf(c) -> c, stream
      | _, Node(l, r) ->
        match stream with
          | hd :: tl -> 
            (* If we encounter a 0, then branch left, else right *)
            if hd = 0 then traverse l tl else traverse r tl
          | [] -> char_of_int 0, [] in
  
  (* Helper function to keep the stream going *)
  let rec fold (stream : int list) (acc : char list) : char list =
    match stream with 
      | [] -> acc
      | _ -> 
        let c, tl = traverse hm stream in
        fold tl (c::acc) in
        
  List.rev (fold bitstream [])

(* Returns bitstream with hm compressed and stored on the beginning of the
   stream in preorder *)
let prepend_tree (hm : huffmantree) (bitstream: int list) : int list =
  (* Traverses the tree and appends 510*9 bits representing the tree onto bitstream *)
  (* New function because I wasn't sure if we're allowed to change function to rec *)
  let rec fold_tree (tree : huffmantree) (acc : int list) : int list =
    match tree with
      | i, Leaf(_) -> to_bits 9 acc i
      | i, Node(l, r) ->
        let a = fold_tree r acc in
        to_bits 9 (fold_tree l a) i in
        
  fold_tree hm bitstream

(* Returns bitstream with the encoding for a huffman tree removed from the
   beginning and decoded.
   Requires: a valid encoding for a huffman tree at the beginning of
   bitstream *)
let regrow_tree (bitstream : int list) : huffmantree * int list =
  (* Iterates over the bitstream and given the assumption that tree is*)
  (* stored in preorder and Leaf terminated, reconstructs our huffman tree *)
  let rec maketree (stream : int list) : huffmantree * int list =
    let index, tl = read_n_bits_to_int 9 stream in
    if index < 256 then
      (index, Leaf(char_of_int index)), tl
    else
      (* First grow out the left side *)
      let l, t = maketree tl in
      (* The grow out the right side *)
      let r, t2 = maketree t in
      (index, Node(l, r)), t2 in
  
  maketree bitstream
  
let encode_file (fname : string) : unit =
  let chars = load_chars fname in
  let (hm, codelist) = encode (chars) in
  let prepended = prepend_tree hm codelist in
  let compressedf = open_out_bin (fname ^ "-compressed.txt") in
  write_bits compressedf prepended (longest_code (get_codes hm));
  close_out compressedf

let decode_file (fname : string) : unit =
  let readcompf = open_in_bin (fname ^ "-compressed.txt") in
  let decodedf = open_out (fname ^ "-decoded.txt") in
  let readin = read_bits readcompf in
  let (regrown, bitstream) = regrow_tree readin in
  let chrs = decode regrown bitstream in
  write_chars decodedf chrs;
  close_in readcompf;
  close_out decodedf

let main (fname : string) : unit =
  encode_file fname; decode_file fname