sidef-scripts/Compression/huffman_file_compression.sf

#!/usr/bin/ruby

# Compress/decompress small files using Huffman coding.

define(
    CHUNK_SIZE = 1024**2,        # 1 MB
    SIGNATURE = ('HFM' + 3.chr),
)

func walk(Hash n, String s, Hash h) {
    if (n.has(:a)) {
        h{n{:a}} = s
        return nil
    }
    walk(n{:0}, s+'0', h) if n.has(:0)
    walk(n{:1}, s+'1', h) if n.has(:1)
}

func make_tree(Array bytes) {

    var nodes = bytes.freq.sort.map_2d { |b,f|
        Hash(a => b, freq => f)
    }

    loop {
        nodes.sort_by!{|h| h{:freq} }
        var(x, y) = (nodes.shift, nodes.shift)
        if (defined(x)) {
            if (defined(y)) {
                nodes << Hash(:0 => x, :1 => y, :freq => x{:freq}+y{:freq})
            }
            else {
                nodes << Hash(:0 => x, :freq => x{:freq})
            }
        }
        nodes.len > 1 || break
    }

    var n = nodes.first
    walk(n, '', n{:tree} = Hash())
    return n
}

func huffman_encode(Array bytes, Hash t) {
    join('', t{bytes...})
}

func huffman_decode (String bits, Hash tree) {
    bits.gsub(Regex('(' + tree.keys.sort_by { .len }.join('|') + ')'), {|s| tree{s} })
}

func create_huffman_entry (Array bytes, FileHandle out_fh) {

    var h   = make_tree(bytes){:tree}
    var enc = huffman_encode(bytes, h)

    var dict  = ''
    var codes = ''

    for i in (0..255) {
        var c = (h{i} \\ '')
        codes += c
        dict  += c.len.chr
    }

    out_fh.print(dict)
    out_fh.print(pack("B*", codes))
    out_fh.print(pack("N",  enc.len))
    out_fh.print(pack("B*", enc))
}

# Compress file
func huffman_compress_file(File input, File output) {

    var in_fh  =  input.open('<:raw') || die "Can't open file <<#{input}>> for reading"
    var out_fh = output.open('>:raw') || die "Can't open file <<#{output}>> for writing"

    var header = SIGNATURE

    # Print the header
    out_fh.print(header)

    while (in_fh.read(\var chunk, CHUNK_SIZE)) {
        create_huffman_entry([unpack('C*', chunk)], out_fh)
    }

    # Close the files
    in_fh.close
    out_fh.close
}

func read_bits (FileHandle fh, Number bits_len) {

    var str = ''
    fh.read(\str, bits_len>>3)
    str = unpack('B*', str)

    while (str.len < bits_len) {
        str += unpack('B*', fh.getc \\ return nil)
    }

    if (str.len > bits_len) {
        str.substr!(0, bits_len)
    }

    return str
}

func decode_huffman_entry (FileHandle fh, FileHandle out_fh) {

    var codes = []
    var codes_len = 0

    fh.read(\var buffer, 256)

    [unpack('C*', buffer)].map{ Num(_) }.each_kv {|c,l|
        if (l > 0) {
            codes_len += l
            codes << [c, l]
        }
    }

    var codes_bin = read_bits(fh, codes_len)

    var rev_dict = Hash()
    for c,l in (codes) {
        var code = substr(codes_bin, 0, l)
        codes_bin.substr!(l)
        rev_dict{code} = c.chr
    }

    var enc_len = Num(unpack('N', 4.of{ fh.getc }.join))

    if (enc_len > 0) {
        var enc_data = read_bits(fh, enc_len)
        out_fh.print(huffman_decode(enc_data, rev_dict))
        return true
    }

    return false
}

# Decompress file
func huffman_decompress_file(File input, File output) {

    var in_fh = input.open('<:raw') || die "Can't open file <<#{input}>> for reading"

    if (SIGNATURE.len.of { in_fh.getc }.join != SIGNATURE) {
        die "Not a HFM archive!\n"
    }

    var out_fh = output.open('>:raw') || die "Can't open file <<#{output}>> for writing"

    while (!in_fh.eof) {
        decode_huffman_entry(in_fh, out_fh) || break
    }

    in_fh.close
    out_fh.close
}

ARGV.getopt!('d', \var decode)

var file = File(ARGV.shift) || do {
    say "usage: #{File(__MAIN__).basename} [-d] [input file]"
    Sys.exit(2)
}

if (decode || file.match(/\.hfm\.enc\z/)) {
    huffman_decompress_file(file, File("output.hfm.dec"))
}
else {
    huffman_compress_file(file, File("output.hfm.enc"))
}