Chris@4: /* fitblk.c: example of fitting compressed output to a specified size
Chris@4:    Not copyrighted -- provided to the public domain
Chris@4:    Version 1.1  25 November 2004  Mark Adler */
Chris@4: 
Chris@4: /* Version history:
Chris@4:    1.0  24 Nov 2004  First version
Chris@4:    1.1  25 Nov 2004  Change deflateInit2() to deflateInit()
Chris@4:                      Use fixed-size, stack-allocated raw buffers
Chris@4:                      Simplify code moving compression to subroutines
Chris@4:                      Use assert() for internal errors
Chris@4:                      Add detailed description of approach
Chris@4:  */
Chris@4: 
Chris@4: /* Approach to just fitting a requested compressed size:
Chris@4: 
Chris@4:    fitblk performs three compression passes on a portion of the input
Chris@4:    data in order to determine how much of that input will compress to
Chris@4:    nearly the requested output block size.  The first pass generates
Chris@4:    enough deflate blocks to produce output to fill the requested
Chris@4:    output size plus a specfied excess amount (see the EXCESS define
Chris@4:    below).  The last deflate block may go quite a bit past that, but
Chris@4:    is discarded.  The second pass decompresses and recompresses just
Chris@4:    the compressed data that fit in the requested plus excess sized
Chris@4:    buffer.  The deflate process is terminated after that amount of
Chris@4:    input, which is less than the amount consumed on the first pass.
Chris@4:    The last deflate block of the result will be of a comparable size
Chris@4:    to the final product, so that the header for that deflate block and
Chris@4:    the compression ratio for that block will be about the same as in
Chris@4:    the final product.  The third compression pass decompresses the
Chris@4:    result of the second step, but only the compressed data up to the
Chris@4:    requested size minus an amount to allow the compressed stream to
Chris@4:    complete (see the MARGIN define below).  That will result in a
Chris@4:    final compressed stream whose length is less than or equal to the
Chris@4:    requested size.  Assuming sufficient input and a requested size
Chris@4:    greater than a few hundred bytes, the shortfall will typically be
Chris@4:    less than ten bytes.
Chris@4: 
Chris@4:    If the input is short enough that the first compression completes
Chris@4:    before filling the requested output size, then that compressed
Chris@4:    stream is return with no recompression.
Chris@4: 
Chris@4:    EXCESS is chosen to be just greater than the shortfall seen in a
Chris@4:    two pass approach similar to the above.  That shortfall is due to
Chris@4:    the last deflate block compressing more efficiently with a smaller
Chris@4:    header on the second pass.  EXCESS is set to be large enough so
Chris@4:    that there is enough uncompressed data for the second pass to fill
Chris@4:    out the requested size, and small enough so that the final deflate
Chris@4:    block of the second pass will be close in size to the final deflate
Chris@4:    block of the third and final pass.  MARGIN is chosen to be just
Chris@4:    large enough to assure that the final compression has enough room
Chris@4:    to complete in all cases.
Chris@4:  */
Chris@4: 
Chris@4: #include <stdio.h>
Chris@4: #include <stdlib.h>
Chris@4: #include <assert.h>
Chris@4: #include "zlib.h"
Chris@4: 
Chris@4: #define local static
Chris@4: 
Chris@4: /* print nastygram and leave */
Chris@4: local void quit(char *why)
Chris@4: {
Chris@4:     fprintf(stderr, "fitblk abort: %s\n", why);
Chris@4:     exit(1);
Chris@4: }
Chris@4: 
Chris@4: #define RAWLEN 4096    /* intermediate uncompressed buffer size */
Chris@4: 
Chris@4: /* compress from file to def until provided buffer is full or end of
Chris@4:    input reached; return last deflate() return value, or Z_ERRNO if
Chris@4:    there was read error on the file */
Chris@4: local int partcompress(FILE *in, z_streamp def)
Chris@4: {
Chris@4:     int ret, flush;
Chris@4:     unsigned char raw[RAWLEN];
Chris@4: 
Chris@4:     flush = Z_NO_FLUSH;
Chris@4:     do {
Chris@4:         def->avail_in = fread(raw, 1, RAWLEN, in);
Chris@4:         if (ferror(in))
Chris@4:             return Z_ERRNO;
Chris@4:         def->next_in = raw;
Chris@4:         if (feof(in))
Chris@4:             flush = Z_FINISH;
Chris@4:         ret = deflate(def, flush);
Chris@4:         assert(ret != Z_STREAM_ERROR);
Chris@4:     } while (def->avail_out != 0 && flush == Z_NO_FLUSH);
Chris@4:     return ret;
Chris@4: }
Chris@4: 
Chris@4: /* recompress from inf's input to def's output; the input for inf and
Chris@4:    the output for def are set in those structures before calling;
Chris@4:    return last deflate() return value, or Z_MEM_ERROR if inflate()
Chris@4:    was not able to allocate enough memory when it needed to */
Chris@4: local int recompress(z_streamp inf, z_streamp def)
Chris@4: {
Chris@4:     int ret, flush;
Chris@4:     unsigned char raw[RAWLEN];
Chris@4: 
Chris@4:     flush = Z_NO_FLUSH;
Chris@4:     do {
Chris@4:         /* decompress */
Chris@4:         inf->avail_out = RAWLEN;
Chris@4:         inf->next_out = raw;
Chris@4:         ret = inflate(inf, Z_NO_FLUSH);
Chris@4:         assert(ret != Z_STREAM_ERROR && ret != Z_DATA_ERROR &&
Chris@4:                ret != Z_NEED_DICT);
Chris@4:         if (ret == Z_MEM_ERROR)
Chris@4:             return ret;
Chris@4: 
Chris@4:         /* compress what was decompresed until done or no room */
Chris@4:         def->avail_in = RAWLEN - inf->avail_out;
Chris@4:         def->next_in = raw;
Chris@4:         if (inf->avail_out != 0)
Chris@4:             flush = Z_FINISH;
Chris@4:         ret = deflate(def, flush);
Chris@4:         assert(ret != Z_STREAM_ERROR);
Chris@4:     } while (ret != Z_STREAM_END && def->avail_out != 0);
Chris@4:     return ret;
Chris@4: }
Chris@4: 
Chris@4: #define EXCESS 256      /* empirically determined stream overage */
Chris@4: #define MARGIN 8        /* amount to back off for completion */
Chris@4: 
Chris@4: /* compress from stdin to fixed-size block on stdout */
Chris@4: int main(int argc, char **argv)
Chris@4: {
Chris@4:     int ret;                /* return code */
Chris@4:     unsigned size;          /* requested fixed output block size */
Chris@4:     unsigned have;          /* bytes written by deflate() call */
Chris@4:     unsigned char *blk;     /* intermediate and final stream */
Chris@4:     unsigned char *tmp;     /* close to desired size stream */
Chris@4:     z_stream def, inf;      /* zlib deflate and inflate states */
Chris@4: 
Chris@4:     /* get requested output size */
Chris@4:     if (argc != 2)
Chris@4:         quit("need one argument: size of output block");
Chris@4:     ret = strtol(argv[1], argv + 1, 10);
Chris@4:     if (argv[1][0] != 0)
Chris@4:         quit("argument must be a number");
Chris@4:     if (ret < 8)            /* 8 is minimum zlib stream size */
Chris@4:         quit("need positive size of 8 or greater");
Chris@4:     size = (unsigned)ret;
Chris@4: 
Chris@4:     /* allocate memory for buffers and compression engine */
Chris@4:     blk = malloc(size + EXCESS);
Chris@4:     def.zalloc = Z_NULL;
Chris@4:     def.zfree = Z_NULL;
Chris@4:     def.opaque = Z_NULL;
Chris@4:     ret = deflateInit(&def, Z_DEFAULT_COMPRESSION);
Chris@4:     if (ret != Z_OK || blk == NULL)
Chris@4:         quit("out of memory");
Chris@4: 
Chris@4:     /* compress from stdin until output full, or no more input */
Chris@4:     def.avail_out = size + EXCESS;
Chris@4:     def.next_out = blk;
Chris@4:     ret = partcompress(stdin, &def);
Chris@4:     if (ret == Z_ERRNO)
Chris@4:         quit("error reading input");
Chris@4: 
Chris@4:     /* if it all fit, then size was undersubscribed -- done! */
Chris@4:     if (ret == Z_STREAM_END && def.avail_out >= EXCESS) {
Chris@4:         /* write block to stdout */
Chris@4:         have = size + EXCESS - def.avail_out;
Chris@4:         if (fwrite(blk, 1, have, stdout) != have || ferror(stdout))
Chris@4:             quit("error writing output");
Chris@4: 
Chris@4:         /* clean up and print results to stderr */
Chris@4:         ret = deflateEnd(&def);
Chris@4:         assert(ret != Z_STREAM_ERROR);
Chris@4:         free(blk);
Chris@4:         fprintf(stderr,
Chris@4:                 "%u bytes unused out of %u requested (all input)\n",
Chris@4:                 size - have, size);
Chris@4:         return 0;
Chris@4:     }
Chris@4: 
Chris@4:     /* it didn't all fit -- set up for recompression */
Chris@4:     inf.zalloc = Z_NULL;
Chris@4:     inf.zfree = Z_NULL;
Chris@4:     inf.opaque = Z_NULL;
Chris@4:     inf.avail_in = 0;
Chris@4:     inf.next_in = Z_NULL;
Chris@4:     ret = inflateInit(&inf);
Chris@4:     tmp = malloc(size + EXCESS);
Chris@4:     if (ret != Z_OK || tmp == NULL)
Chris@4:         quit("out of memory");
Chris@4:     ret = deflateReset(&def);
Chris@4:     assert(ret != Z_STREAM_ERROR);
Chris@4: 
Chris@4:     /* do first recompression close to the right amount */
Chris@4:     inf.avail_in = size + EXCESS;
Chris@4:     inf.next_in = blk;
Chris@4:     def.avail_out = size + EXCESS;
Chris@4:     def.next_out = tmp;
Chris@4:     ret = recompress(&inf, &def);
Chris@4:     if (ret == Z_MEM_ERROR)
Chris@4:         quit("out of memory");
Chris@4: 
Chris@4:     /* set up for next reocmpression */
Chris@4:     ret = inflateReset(&inf);
Chris@4:     assert(ret != Z_STREAM_ERROR);
Chris@4:     ret = deflateReset(&def);
Chris@4:     assert(ret != Z_STREAM_ERROR);
Chris@4: 
Chris@4:     /* do second and final recompression (third compression) */
Chris@4:     inf.avail_in = size - MARGIN;   /* assure stream will complete */
Chris@4:     inf.next_in = tmp;
Chris@4:     def.avail_out = size;
Chris@4:     def.next_out = blk;
Chris@4:     ret = recompress(&inf, &def);
Chris@4:     if (ret == Z_MEM_ERROR)
Chris@4:         quit("out of memory");
Chris@4:     assert(ret == Z_STREAM_END);    /* otherwise MARGIN too small */
Chris@4: 
Chris@4:     /* done -- write block to stdout */
Chris@4:     have = size - def.avail_out;
Chris@4:     if (fwrite(blk, 1, have, stdout) != have || ferror(stdout))
Chris@4:         quit("error writing output");
Chris@4: 
Chris@4:     /* clean up and print results to stderr */
Chris@4:     free(tmp);
Chris@4:     ret = inflateEnd(&inf);
Chris@4:     assert(ret != Z_STREAM_ERROR);
Chris@4:     ret = deflateEnd(&def);
Chris@4:     assert(ret != Z_STREAM_ERROR);
Chris@4:     free(blk);
Chris@4:     fprintf(stderr,
Chris@4:             "%u bytes unused out of %u requested (%lu input)\n",
Chris@4:             size - have, size, def.total_in);
Chris@4:     return 0;
Chris@4: }