cleaning and bug fix

[compress.git] / compress.c

/* depend: */
/* cflags: */
/* linker: atoi.o code.o debug.o fprintf.o */

#include <fcntl.h>
#include <unistd.h>
#include <stddef.h>
#include "atoi.h"
#include "code.h"
#include "debug.h"
#include "fprintf.h"

/* constants */

#define BUFFER_SIZE 4096

#define COMPRESS 1
#define DECOMPRESS 2

/* macros */

/* gobal variables */

char *progname = NULL;

/* help function */

int usage (int ret)
{
    int fd = ret ? _fderr : _fdout;
    fdprintf (fd, "usage: %s\n", progname);
    fdprintf (fd, " -h : help message\n");
    fdprintf (fd, " -i <file>: input file\n");
    fdprintf (fd, " -o <file>: output file\n");
    fdprintf (fd, " -v : verbose level (%d)\n", verbose);

    return ret;
}

/* create occurence table */

int *create_table (char *filename)
{
    byte_t buffer[BUFFER_SIZE] = {0};
    int nbread;
    static int table[NB_BYTES] = {0};
    int fid = 0;

    VERBOSE (DEBUG, PRINTOUT ("start creating occurence table\n"));

    /* open file */
    fid = open (filename, O_RDONLY|O_RAW);
    if (fid == -1) {
        VERBOSE (ERROR, PRINTERR ("can't open file '%s'\n", filename));
        return NULL;
    }
    VERBOSE (INFO, PRINTOUT ("file '%s' opened\n", filename));

    /* read file */
    while ((nbread = read (fid, buffer, BUFFER_SIZE)) > 0) {
        VERBOSE (DEBUG, PRINTOUT ("nbread: %d\n", nbread));
        while (nbread--) {
            table[(int)buffer[nbread]]++;
        }
    }

    /* close file */
    close (fid);

    VERBOSE (DEBUG, PRINTOUT ("end creating occurence table\n"));

    return table;
}

/* print occurence table */

void print_occ_table (int *table)
{
    int i;

    PRINTOUT ("Occurence table\n");
    for (i = 0; i < NB_BYTES; i++) {
        if (table[i]) {
            PRINTOUT ("0x%02x '%c': %d\n", i, ((i < 32) || (i > 127)) ? '.' : i, table[i]);
        }
    }
}

/* initialize forest */

leaf_t **init_forest (int *table)
{
    static leaf_t *leafs[NB_BYTES] = {0};
    int nb_leafs = 0;
    int i, l;

    VERBOSE (DEBUG, PRINTOUT ("start initiliazing forest\n"));

    /* count number of leafs */
    for (i = 0; i < NB_BYTES; i++) {
        if (table[i] > 0) {
            nb_leafs++;
        }
    }

    /* initialize leafs */
    for (i = 0, l = 0; i < NB_BYTES; i++) {
        if (table[i] > 0) {
            leafs[l] = getleaf (1);
            if (leafs[l] == NULL) {
                VERBOSE (ERROR, PRINTERR ("can't allocate memory\n"));
                return NULL;
            }
            leafs[l]->occ = table[i];
            leafs[l]->c = i;
            l++;
        }
    }

    VERBOSE (DEBUG, PRINTOUT ("end initiliazing forest\n"));

    return leafs;
}

/* create tree */

leaf_t *create_tree (leaf_t **leafs)
{
    leaf_t *branch = NULL;
    int nb_leafs = 0;
    int last = -1;
    int ante;
    int i, j;

    VERBOSE (DEBUG, PRINTOUT ("start creating tree\n"));

    /* count number of leafs */
    while (leafs[nb_leafs] != NULL) {
        nb_leafs++;
    }

    /* create tree */
    for (j = 0; j < nb_leafs - 1; j++) {

        /* look for leatest occurence */
        last = -1;
        for (i = 0; i < nb_leafs; i++) {
            if (leafs[i] == NULL) {
                continue;
            }
            if ((last == -1) || (leafs[i]->occ < leafs[last]->occ)) {
                last = i;
            }
        }

        /* look for ante leatest occurence */
        ante = -1;
        for (i = 0; i < nb_leafs; i++) {
            if ((i == last) || (leafs[i] == NULL)) {
                continue;
            }
            if ((ante == -1) || (leafs[i]->occ < leafs[ante]->occ)) {
                ante = i;
            }
        }

        /* create branch */
        if ((last == -1) || (ante == -1)) {
            VERBOSE (ERROR, PRINTERR ("error during tree building\n"));
            return NULL;
        }
        branch = getleaf (1);
        if (branch == NULL) {
            VERBOSE (ERROR, PRINTERR ("can't allocate memory\n"));
            return NULL;
        }
        branch->left = leafs[last];
        branch->right = leafs[ante];
        branch->occ = branch->left->occ + branch->right->occ;
        leafs[last] = branch;
        leafs[ante] = NULL;
    }

    VERBOSE (DEBUG, PRINTOUT ("end creating tree\n"));

    return (last != -1) ? leafs[last] : NULL;
}

/* explore tree */

void explore_tree (code_t *table, leaf_t *root, char *code, int index)
{

    VERBOSE (DEBUG, PRINTOUT ("start exploring code tree\n"));

    if ((root->left == NULL) && (root->right == NULL)) {
        codcpy ((char *)(table + (int)(root->c)), sizeof (code_t), code);
    }
    else {
        codcpy (code + index, sizeof (code_t), "1");
        explore_tree (table, root->left, code, index + 1);
        codcpy (code + index, sizeof (code_t), "0");
        explore_tree (table, root->right, code, index + 1);
    }

    VERBOSE (DEBUG, PRINTOUT ("end exploring code tree\n"));
}

/* create code table */
code_t *create_code (leaf_t *root)
{
    static code_t table[NB_BYTES] = {0};
    code_t code = {0};

    VERBOSE (DEBUG, PRINTOUT ("start creating code table\n"));

    explore_tree (table, root, (char *)&code, 0);

    VERBOSE (DEBUG, PRINTOUT ("end creating code table\n"));

    return table;
}

/* print code table */

void print_code_table (code_t *codes)
{
    char *code;
    int i;

    PRINTOUT ("Code table\n");
    for (i = 0; i < NB_BYTES; i++) {
        code = (char *)(codes + i);
        if (codlen (code) == 0) {
            continue;
        }
        PRINTOUT ("0x%02x '%c': %s\n", i, ((i < 32) || (i > 127)) ? '.' : i, code);
    }
}

/* encode header and code table */

byte_t *encode_header_table (code_t *codes, int *occ)
{
    static byte_t buffer[NB_BYTES * (NB_BYTES - 1) / 2 / 8 + NB_BYTES + 6] = {0};
    char bits[(NB_BYTES - 1) + 8 + 1] = {0};
    char *code;
    byte_t *header = buffer;
    int i, j, length, mode;
    int nb = 0;
    int size = 0;

    VERBOSE (DEBUG, PRINTOUT ("start encoding header and code table\n"));

    /* mode 1 or 2 */
    for (i = 0; i < NB_BYTES; i++) {
        code = (char *)(codes + i);
        if (codlen (code) > 0) {
            nb++;
            size += codlen (code) * occ[i];
        }
    }
    mode = (NB_BYTES < 2 * nb + 1) ? 1 : 2;
    VERBOSE (DEBUG, PRINTOUT ("nb chars: %d\n", nb));
    VERBOSE (DEBUG, PRINTOUT ("mode: %d\n", mode));
    VERBOSE (DEBUG, PRINTOUT ("size: %d\n", size));
    VERBOSE (DEBUG, PRINTOUT ("rem: %d\n", size % 256));

    /* header */
    codcpy ((char *)header, sizeof (buffer), (mode == 1) ? "MZ1   " : "MZ2   ");
    header += 6;

    /* size */
    switch (mode) {
    case 1:
        for (i = 0; i < NB_BYTES; i++) {
            code = (char *)(codes + i);
            *(header++) = (byte_t) codlen (code);
        }
        break;
    case 2:
        *(header++) = (byte_t)(nb - 1);
        for (i = 0; i < NB_BYTES; i++) {
            code = (char *)(codes + i);
            if (codlen (code) > 0) {
                *(header++) = (byte_t) i;
                *(header++) = (byte_t) codlen (code);
            }
        }
        break;
    }

    /* bits */
    for (i = 0; i < NB_BYTES; i++) {
        code = (char *)(codes + i);
        if (codlen (code) > 0) {
            codcat (bits, sizeof (code_t), code);
            while (codlen (bits) > (8 - 1)) {
                for (j = 0; j < 8; j++) {
                    *header <<= 1;
                    if (bits[j] == '1') {
                        (*header)++;
                    }
                }
                codcpy (bits, sizeof (code_t), bits + 8);
                header++;
            }
        }
    }
    if (codlen (bits) > 0) {
        for (j = 0; j < (int)codlen (bits); j++) {
            *header <<= 1;
            if (bits[j] == '1') {
                (*header)++;
            }
        }
        for (j = (int)codlen (bits); j < 8; j++) {
            *header <<= 1;
        }
        header++;
    }

    /* length */
    length = (int)(header - buffer - 6);
    VERBOSE (DEBUG, PRINTOUT ("lengh: %d %02x %02x\n", length, length >> 8, length & 0xff));
    buffer[3] = (byte_t)(length >> 8);
    buffer[4] = (byte_t)(length & 0xff);
    buffer[5] = (byte_t)(size % 256);
    header = buffer;

    VERBOSE (DEBUG, PRINTOUT ("end encoding header and code table\n"));

    return header;
}

/* print header */

void print_header (byte_t *header)
{
    int length, i;

    length = (header[3] << 8) + header[4];
    VERBOSE (DEBUG, PRINTOUT ("lengh: %d\n", length));
    for (i = 0; i < length + 6; i++) {
        PRINTOUT ("%02x", header[i]);
    }
    PRINTOUT ("\n");
}

/* write crompressed file */

int write_compress (char *output, char *input, code_t *codes, byte_t *header)
{
    byte_t bufin[BUFFER_SIZE] = {0};
    byte_t bufout[BUFFER_SIZE] = {0};
    char bits[(NB_BYTES - 1) + 8 + 1] = {0};
    int fin, fout;
    int length = 0;
    int i, j, nbread;
    byte_t *pt;

    VERBOSE (DEBUG, PRINTOUT ("start writting compressed file\n"));

    /* open input file */
    fin = open (input, O_RDONLY|O_RAW);
    if (fin == -1) {
        VERBOSE (ERROR, PRINTERR ("can't open file '%s' for reading\n", input));
        return 1;
    }
    VERBOSE (INFO, PRINTOUT ("file '%s' opened\n", input));

    /* open output file */
    fout = open (output, O_WRONLY|O_CREAT|O_RAW, 0700);
    if (fout == -1) {
        VERBOSE (ERROR, PRINTERR ("can't open file '%s' for writing\n", output));
	    close (fin);
        return 1;
    }
    VERBOSE (INFO, PRINTOUT ("file '%s' opened\n", output));

    /* write header */
    length = (header[3] << 8) + header[4];
    VERBOSE (DEBUG, PRINTOUT ("lengh: %d\n", length));
    write (fout, header, length + 6);

    /* write file */
    pt = bufout;
    while ((nbread = read (fin, bufin, BUFFER_SIZE)) > 0) {
        VERBOSE (DEBUG, PRINTOUT ("nbread: %d\n", nbread));
        for (i = 0; i < nbread; i++) {
            codcat (bits, sizeof (code_t), (char *)(codes + bufin[i]));
            while (codlen (bits) > (8 - 1)) {
                for (j = 0; j < 8; j++) {
                    *pt <<= 1;
                    if (bits[j] == '1') {
                        (*pt)++;
                    }
                }
                codcpy (bits, sizeof (code_t), bits + 8);
                if (pt - bufout == BUFFER_SIZE - 1) {
                    write (fout, bufout, BUFFER_SIZE);
                    pt = bufout;
                } else {
                    pt++;
                }
            }
        }
    }
    VERBOSE (DEBUG, PRINTOUT ("lastest bits : %d\n", codlen (bits)));
    if (codlen (bits) > 0) {
        for (j = 0; j < (int)codlen (bits); j++) {
            *pt <<= 1;
            if (bits[j] == '1') {
                (*pt)++;
            }
        }
        for (j = (int)codlen (bits); j < 8; j++) {
            *pt <<= 1;
        }
        pt++;
    }
    if (pt != bufout) {
        VERBOSE (DEBUG, PRINTOUT ("last partial buffer written: %u\n", pt - bufout));
        write (fout, bufout, pt - bufout);
    }
    
    /* closing */
    close (fin);
    close (fout);

    VERBOSE (DEBUG, PRINTOUT ("end writting compressed file\n"));

    return 0;
}

/* read header */

code_t *read_header (char *filename) {
    static code_t table[NB_BYTES] = {0};
    byte_t buffer[NB_BYTES * (NB_BYTES - 1) / 2 / 8 + NB_BYTES + 6] = {0};
    byte_t *codes = NULL;
    byte_t cur;
    int lengths[NB_BYTES] = {0};
    int fid;
    int mode = 0;
    int i, j, l, nb, size;

    VERBOSE (DEBUG, PRINTOUT ("start reading header\n"));

    /* open file */
    fid = open (filename, O_RDONLY|O_RAW);
    if (fid == -1) {
        VERBOSE (ERROR, PRINTERR ("can't open file '%s'\n", filename));
        return NULL;
    }
    VERBOSE (INFO, PRINTOUT ("file '%s' opened\n", filename));

    /* read magic number */
    nb = read (fid, buffer, 6);
    VERBOSE (DEBUG, PRINTOUT ("nb, buffer: %d 0x%02x 0x%02x\n", nb, buffer[0], buffer[1]));
    if ((nb == 6) && (buffer[0] == 'M') && (buffer[1] == 'Z')) {
        mode = (buffer[2] == '1') ? 1 : (buffer[2] == '2') ? 2 : 0;
        size = (buffer[3] << 8) + buffer[4];
        VERBOSE (DEBUG, PRINTOUT ("mode, size: %d %d\n", mode, size));
        if (size > NB_BYTES * (NB_BYTES - 1) / 2 / 8 + NB_BYTES) {
            mode = 0;
        } else {
            nb = read (fid, buffer, size);
            VERBOSE (DEBUG, PRINTOUT ("nb read: %d/%d\n", nb, size));
            if (nb != size) {
                mode = 0;
            }
        }
    }
    close (fid);
    if (mode == 0) {
        VERBOSE (ERROR, PRINTERR ("incorrect file\n"));
        return NULL;
    }

    /* analyse header */
    codes = buffer;
    switch (mode) {
    case 1:
        for (i = 0; i < NB_BYTES; i++) {
            lengths[i] = *(codes++);
        }
        break;
    case 2:
        nb = *(codes++) + 1;
        VERBOSE (DEBUG, PRINTOUT ("nb codes: %d\n", nb));
        for (i = 0; i < nb; i++) {
            j = *(codes++);
            lengths[j] = *(codes++);
        }
        break;
    }
    VERBOSE (DEBUG, for (i = 0; i < NB_BYTES; i++) if (lengths[i]) PRINTOUT ("%d: %d\n", i, lengths[i]));

    /* check lengths */
    for (i = 0, l = 0; i < NB_BYTES; i++) {
        l += lengths[i];
    }
    if (((mode == 1) && (size - 256 != (l + 7) / 8)) ||
        ((mode == 2) && (size - 2 * nb - 1 != (l + 7) / 8))) {
        VERBOSE (ERROR, PRINTERR ("incorrect code table length: %d %d %d\n", size, nb, l));
        return NULL;
    }

    /* decode code */
    cur = *(codes++);
    l = 8;
    for (i = 0; i < NB_BYTES; i++) {
        if (lengths[i] == 0) {
            continue;
        }
        while (lengths[i]--) {
            codcat ((char *)(table + i), sizeof (code_t), ((cur & 0x80) == 0) ? "0" : "1");
            l--;
            cur <<= 1;
            if (l == 0) {
                cur = *(codes++);
                l = 8;
            }
        }
    }

    VERBOSE (DEBUG, PRINTOUT ("end reading header\n"));

    return table;
}

/* write decompressed file */

int write_decompress (char *output, char *input, code_t *codes)
{
    byte_t bufin[BUFFER_SIZE] = {0};
    byte_t bufout[BUFFER_SIZE] = {0};
    byte_t bufhea[MAX(NB_BYTES * (NB_BYTES - 1) / 2 / 8 + NB_BYTES + 6, BUFFER_SIZE)] = {0};
    char bits[(NB_BYTES - 1) + 1] = {0};
    int fin, fout;
    int i, j, k, nb, size, rem;
    int is_found;
    int l = 0;
    byte_t *pt;

    VERBOSE (DEBUG, PRINTOUT ("start writing decompressed file\n"));

    /* open file for reading */
    fin = open (input, O_RDONLY|O_RAW);
    if (fin == -1) {
        VERBOSE (ERROR, PRINTERR ("can't open file '%s' for reading\n", input));
        return 1;
    }
    VERBOSE (INFO, PRINTOUT ("file '%s' opened\n", input));

    /* read magic number */
    nb = read (fin, bufhea, 6);
    if (nb != 6) {
        VERBOSE (ERROR, PRINTERR ("can't read file\n"));
        close (fin);
        return 1;
    }
    size = (bufhea[3] << 8) + bufhea[4];
    VERBOSE (DEBUG, PRINTOUT ("table size: %d\n", size));
    rem = bufhea[5];
    VERBOSE (DEBUG, PRINTOUT ("remainder: %d\n", rem));
    nb = read (fin, bufhea, size);
    if (nb != size) {
        VERBOSE (ERROR, PRINTERR ("can't read file\n"));
        close (fin);
        return 1;
    }

    /* open file for writing */
    fout = open (output, O_WRONLY|O_CREAT|O_RAW, 0700);
    if (fout == -1) {
        VERBOSE (ERROR, PRINTERR ("can't open file '%s' for writing\n", output));
        close (fin);
	return 2;
    }
    VERBOSE (INFO, PRINTOUT ("file '%s' opened\n", output));

    /* write file */
    pt = bufout;
    while ((nb = read (fin, bufin, BUFFER_SIZE)) > 0) {
        VERBOSE (DEBUG, PRINTOUT ("nbread: %d\n", nb));
        for (i = 0; i < nb; i++) {
            for (j = 0; j < 8; j++) {
                codcat (bits, sizeof (bits), ((bufin[i] & 0x80) == 0) ? "0" : "1");
                bufin[i] <<= 1;
                l++;
                VERBOSE (DEBUG, PRINTOUT ("bits: %d - %s\n", codlen (bits), bits));

                /* look for correct code */
                is_found = 0;
                for (k = 0; (k < NB_BYTES) && (!is_found); k++) {
                    if (codcmp ((char *)(codes + k), bits) == 0) {
                        is_found = 1;
                        VERBOSE (DEBUG, PRINTOUT ("found: %d\n", k));
                        *pt= k;
                        bits[0] = 0;
                        if (pt - bufout == BUFFER_SIZE - 1) {
                            VERBOSE (DEBUG, PRINTOUT ("nb buffer out: %u\n", (pt - bufout)));
                            write (fout, bufout, BUFFER_SIZE);
                            pt = bufout;
                        } else {
                            pt++;
                        }
                    }
                }
                if ((i == nb - 1) && (l % 256 == rem) && (nb != BUFFER_SIZE)) {
                    VERBOSE (DEBUG, PRINTOUT ("break\n"));
                    break;
                }
            }
        }
    }
    if (pt != bufout) {
        VERBOSE (DEBUG, PRINTOUT ("nb buffer out: %u\n", (pt - bufout)));
        write (fout, bufout, pt - bufout);
    }

    /* close files */
    close (fin);
    close (fout);

    VERBOSE (DEBUG, PRINTOUT ("end writing decompressed file\n"));

    return 0;
}

/* main function */

int main (int argc, char *argv[])
{
    char *input = NULL;
    char *output = NULL;
    int *table = NULL;
    leaf_t **leafs = NULL;
    leaf_t *root = NULL;
    code_t *codes = NULL;
    byte_t *header = NULL;
    int mode = COMPRESS;
    int rc = 1;

    progname = argv[0];

    int c;
    char * arg;
    VERBOSE (DEBUG, PRINTOUT ("start processing arguments\n"));
    while (argc-- > 1) {
        arg = *(++argv);
        if (arg[0] != '-') {
            PRINTERR ("%s: invalid option -- %s\n", progname, arg);
            return usage (1);
        }
        c = arg[1];
        VERBOSE (DEBUG, PRINTOUT ("option: %c\n", c));
        switch (c) {
        case 'c':
            mode = COMPRESS;
            break;
        case 'd':
            mode = DECOMPRESS;
            break;
        case 'i':
            input = (arg[2]) ? arg + 2 : (--argc > 0 ) ? *(++argv) : NULL;
            VERBOSE (DEBUG, PRINTOUT ("input: %s\n", input));
            break;
        case 'o':
            output = (arg[2]) ? arg + 2 : (--argc > 0 ) ? *(++argv) : NULL;
            VERBOSE (DEBUG, PRINTOUT ("output: %s\n", output));
            break;
        case 'v':
            arg = (arg[2]) ? arg + 2 : (--argc > 0 ) ? *(++argv) : NULL;
            if (arg == NULL) {
                PRINTERR ("%s: missing verbose level\n", progname);
                return usage (1);
            }
            verbose = myatoi (arg);
            VERBOSE (INFO, PRINTOUT ("verbose: %d\n", verbose));
            break;
        case 'h':
        default:
            return usage (c != 'h');
        }
    }
    if ((input == NULL) || (output == NULL)) {
        PRINTERR ("%s: missing file\n", progname);
        return usage (1);
    }
    VERBOSE (DEBUG, PRINTOUT ("end processing arguments\n"));

    switch (mode) {
    case COMPRESS:
        table = create_table (input);
        if (table == NULL) break;
        VERBOSE (INFO, print_occ_table (table));

        leafs = init_forest (table);
        if (leafs == NULL) break;
        root = create_tree (leafs);
        if (root == NULL) break;
        codes = create_code (root);
        if (codes == NULL) break;
        VERBOSE (INFO, print_code_table (codes));
        header = encode_header_table (codes, table);
        if (header == NULL) break;
        VERBOSE (INFO, print_header (header));
        rc = write_compress (output, input, codes, header);
        break;
    case DECOMPRESS:
        codes = read_header (input);
        if (codes == NULL) break;
        VERBOSE (INFO, print_code_table (codes));
        rc = write_decompress (output, input, codes);
        break;
    }

    return rc;
}

// test: compress.exe -h
// test: compress.exe -h | awk '/usage:/ { rc=1 } END { exit (1-rc) }'
// test: compress.exe -_ 2> /dev/null | awk 'END { if (NR == 0) { exit(0) } else exit (1) }'
// test: compress.exe -_ 2>&1 | awk '/usage:/ { rc=1 } END { exit (1-rc) }'
// test: compress.exe -c -i compress.c -o compress.mz
// test: ls -sS1 compress.c compress.mz | tail -1 | grep compress.mz
// test: compress.exe -d -i compress.mz -o tmp.c
// test: cmp compress.c tmp.c
// test: rm compress.mz tmp.c

/* vim: set ts=4 sw=4 et: */