1*212397c6Schristos /* fitblk.c: example of fitting compressed output to a specified size
2*212397c6Schristos Not copyrighted -- provided to the public domain
3*212397c6Schristos Version 1.1 25 November 2004 Mark Adler */
4*212397c6Schristos
5*212397c6Schristos /* Version history:
6*212397c6Schristos 1.0 24 Nov 2004 First version
7*212397c6Schristos 1.1 25 Nov 2004 Change deflateInit2() to deflateInit()
8*212397c6Schristos Use fixed-size, stack-allocated raw buffers
9*212397c6Schristos Simplify code moving compression to subroutines
10*212397c6Schristos Use assert() for internal errors
11*212397c6Schristos Add detailed description of approach
12*212397c6Schristos */
13*212397c6Schristos
14*212397c6Schristos /* Approach to just fitting a requested compressed size:
15*212397c6Schristos
16*212397c6Schristos fitblk performs three compression passes on a portion of the input
17*212397c6Schristos data in order to determine how much of that input will compress to
18*212397c6Schristos nearly the requested output block size. The first pass generates
19*212397c6Schristos enough deflate blocks to produce output to fill the requested
20*212397c6Schristos output size plus a specfied excess amount (see the EXCESS define
21*212397c6Schristos below). The last deflate block may go quite a bit past that, but
22*212397c6Schristos is discarded. The second pass decompresses and recompresses just
23*212397c6Schristos the compressed data that fit in the requested plus excess sized
24*212397c6Schristos buffer. The deflate process is terminated after that amount of
25*212397c6Schristos input, which is less than the amount consumed on the first pass.
26*212397c6Schristos The last deflate block of the result will be of a comparable size
27*212397c6Schristos to the final product, so that the header for that deflate block and
28*212397c6Schristos the compression ratio for that block will be about the same as in
29*212397c6Schristos the final product. The third compression pass decompresses the
30*212397c6Schristos result of the second step, but only the compressed data up to the
31*212397c6Schristos requested size minus an amount to allow the compressed stream to
32*212397c6Schristos complete (see the MARGIN define below). That will result in a
33*212397c6Schristos final compressed stream whose length is less than or equal to the
34*212397c6Schristos requested size. Assuming sufficient input and a requested size
35*212397c6Schristos greater than a few hundred bytes, the shortfall will typically be
36*212397c6Schristos less than ten bytes.
37*212397c6Schristos
38*212397c6Schristos If the input is short enough that the first compression completes
39*212397c6Schristos before filling the requested output size, then that compressed
40*212397c6Schristos stream is return with no recompression.
41*212397c6Schristos
42*212397c6Schristos EXCESS is chosen to be just greater than the shortfall seen in a
43*212397c6Schristos two pass approach similar to the above. That shortfall is due to
44*212397c6Schristos the last deflate block compressing more efficiently with a smaller
45*212397c6Schristos header on the second pass. EXCESS is set to be large enough so
46*212397c6Schristos that there is enough uncompressed data for the second pass to fill
47*212397c6Schristos out the requested size, and small enough so that the final deflate
48*212397c6Schristos block of the second pass will be close in size to the final deflate
49*212397c6Schristos block of the third and final pass. MARGIN is chosen to be just
50*212397c6Schristos large enough to assure that the final compression has enough room
51*212397c6Schristos to complete in all cases.
52*212397c6Schristos */
53*212397c6Schristos
54*212397c6Schristos #include <stdio.h>
55*212397c6Schristos #include <stdlib.h>
56*212397c6Schristos #include <assert.h>
57*212397c6Schristos #include "zlib.h"
58*212397c6Schristos
59*212397c6Schristos #define local static
60*212397c6Schristos
61*212397c6Schristos /* print nastygram and leave */
quit(char * why)62*212397c6Schristos local void quit(char *why)
63*212397c6Schristos {
64*212397c6Schristos fprintf(stderr, "fitblk abort: %s\n", why);
65*212397c6Schristos exit(1);
66*212397c6Schristos }
67*212397c6Schristos
68*212397c6Schristos #define RAWLEN 4096 /* intermediate uncompressed buffer size */
69*212397c6Schristos
70*212397c6Schristos /* compress from file to def until provided buffer is full or end of
71*212397c6Schristos input reached; return last deflate() return value, or Z_ERRNO if
72*212397c6Schristos there was read error on the file */
partcompress(FILE * in,z_streamp def)73*212397c6Schristos local int partcompress(FILE *in, z_streamp def)
74*212397c6Schristos {
75*212397c6Schristos int ret, flush;
76*212397c6Schristos unsigned char raw[RAWLEN];
77*212397c6Schristos
78*212397c6Schristos flush = Z_NO_FLUSH;
79*212397c6Schristos do {
80*212397c6Schristos def->avail_in = fread(raw, 1, RAWLEN, in);
81*212397c6Schristos if (ferror(in))
82*212397c6Schristos return Z_ERRNO;
83*212397c6Schristos def->next_in = raw;
84*212397c6Schristos if (feof(in))
85*212397c6Schristos flush = Z_FINISH;
86*212397c6Schristos ret = deflate(def, flush);
87*212397c6Schristos assert(ret != Z_STREAM_ERROR);
88*212397c6Schristos } while (def->avail_out != 0 && flush == Z_NO_FLUSH);
89*212397c6Schristos return ret;
90*212397c6Schristos }
91*212397c6Schristos
92*212397c6Schristos /* recompress from inf's input to def's output; the input for inf and
93*212397c6Schristos the output for def are set in those structures before calling;
94*212397c6Schristos return last deflate() return value, or Z_MEM_ERROR if inflate()
95*212397c6Schristos was not able to allocate enough memory when it needed to */
recompress(z_streamp inf,z_streamp def)96*212397c6Schristos local int recompress(z_streamp inf, z_streamp def)
97*212397c6Schristos {
98*212397c6Schristos int ret, flush;
99*212397c6Schristos unsigned char raw[RAWLEN];
100*212397c6Schristos
101*212397c6Schristos flush = Z_NO_FLUSH;
102*212397c6Schristos do {
103*212397c6Schristos /* decompress */
104*212397c6Schristos inf->avail_out = RAWLEN;
105*212397c6Schristos inf->next_out = raw;
106*212397c6Schristos ret = inflate(inf, Z_NO_FLUSH);
107*212397c6Schristos assert(ret != Z_STREAM_ERROR && ret != Z_DATA_ERROR &&
108*212397c6Schristos ret != Z_NEED_DICT);
109*212397c6Schristos if (ret == Z_MEM_ERROR)
110*212397c6Schristos return ret;
111*212397c6Schristos
112*212397c6Schristos /* compress what was decompresed until done or no room */
113*212397c6Schristos def->avail_in = RAWLEN - inf->avail_out;
114*212397c6Schristos def->next_in = raw;
115*212397c6Schristos if (inf->avail_out != 0)
116*212397c6Schristos flush = Z_FINISH;
117*212397c6Schristos ret = deflate(def, flush);
118*212397c6Schristos assert(ret != Z_STREAM_ERROR);
119*212397c6Schristos } while (ret != Z_STREAM_END && def->avail_out != 0);
120*212397c6Schristos return ret;
121*212397c6Schristos }
122*212397c6Schristos
123*212397c6Schristos #define EXCESS 256 /* empirically determined stream overage */
124*212397c6Schristos #define MARGIN 8 /* amount to back off for completion */
125*212397c6Schristos
126*212397c6Schristos /* compress from stdin to fixed-size block on stdout */
main(int argc,char ** argv)127*212397c6Schristos int main(int argc, char **argv)
128*212397c6Schristos {
129*212397c6Schristos int ret; /* return code */
130*212397c6Schristos unsigned size; /* requested fixed output block size */
131*212397c6Schristos unsigned have; /* bytes written by deflate() call */
132*212397c6Schristos unsigned char *blk; /* intermediate and final stream */
133*212397c6Schristos unsigned char *tmp; /* close to desired size stream */
134*212397c6Schristos z_stream def, inf; /* zlib deflate and inflate states */
135*212397c6Schristos
136*212397c6Schristos /* get requested output size */
137*212397c6Schristos if (argc != 2)
138*212397c6Schristos quit("need one argument: size of output block");
139*212397c6Schristos ret = strtol(argv[1], argv + 1, 10);
140*212397c6Schristos if (argv[1][0] != 0)
141*212397c6Schristos quit("argument must be a number");
142*212397c6Schristos if (ret < 8) /* 8 is minimum zlib stream size */
143*212397c6Schristos quit("need positive size of 8 or greater");
144*212397c6Schristos size = (unsigned)ret;
145*212397c6Schristos
146*212397c6Schristos /* allocate memory for buffers and compression engine */
147*212397c6Schristos blk = malloc(size + EXCESS);
148*212397c6Schristos def.zalloc = Z_NULL;
149*212397c6Schristos def.zfree = Z_NULL;
150*212397c6Schristos def.opaque = Z_NULL;
151*212397c6Schristos ret = deflateInit(&def, Z_DEFAULT_COMPRESSION);
152*212397c6Schristos if (ret != Z_OK || blk == NULL)
153*212397c6Schristos quit("out of memory");
154*212397c6Schristos
155*212397c6Schristos /* compress from stdin until output full, or no more input */
156*212397c6Schristos def.avail_out = size + EXCESS;
157*212397c6Schristos def.next_out = blk;
158*212397c6Schristos ret = partcompress(stdin, &def);
159*212397c6Schristos if (ret == Z_ERRNO)
160*212397c6Schristos quit("error reading input");
161*212397c6Schristos
162*212397c6Schristos /* if it all fit, then size was undersubscribed -- done! */
163*212397c6Schristos if (ret == Z_STREAM_END && def.avail_out >= EXCESS) {
164*212397c6Schristos /* write block to stdout */
165*212397c6Schristos have = size + EXCESS - def.avail_out;
166*212397c6Schristos if (fwrite(blk, 1, have, stdout) != have || ferror(stdout))
167*212397c6Schristos quit("error writing output");
168*212397c6Schristos
169*212397c6Schristos /* clean up and print results to stderr */
170*212397c6Schristos ret = deflateEnd(&def);
171*212397c6Schristos assert(ret != Z_STREAM_ERROR);
172*212397c6Schristos free(blk);
173*212397c6Schristos fprintf(stderr,
174*212397c6Schristos "%u bytes unused out of %u requested (all input)\n",
175*212397c6Schristos size - have, size);
176*212397c6Schristos return 0;
177*212397c6Schristos }
178*212397c6Schristos
179*212397c6Schristos /* it didn't all fit -- set up for recompression */
180*212397c6Schristos inf.zalloc = Z_NULL;
181*212397c6Schristos inf.zfree = Z_NULL;
182*212397c6Schristos inf.opaque = Z_NULL;
183*212397c6Schristos inf.avail_in = 0;
184*212397c6Schristos inf.next_in = Z_NULL;
185*212397c6Schristos ret = inflateInit(&inf);
186*212397c6Schristos tmp = malloc(size + EXCESS);
187*212397c6Schristos if (ret != Z_OK || tmp == NULL)
188*212397c6Schristos quit("out of memory");
189*212397c6Schristos ret = deflateReset(&def);
190*212397c6Schristos assert(ret != Z_STREAM_ERROR);
191*212397c6Schristos
192*212397c6Schristos /* do first recompression close to the right amount */
193*212397c6Schristos inf.avail_in = size + EXCESS;
194*212397c6Schristos inf.next_in = blk;
195*212397c6Schristos def.avail_out = size + EXCESS;
196*212397c6Schristos def.next_out = tmp;
197*212397c6Schristos ret = recompress(&inf, &def);
198*212397c6Schristos if (ret == Z_MEM_ERROR)
199*212397c6Schristos quit("out of memory");
200*212397c6Schristos
201*212397c6Schristos /* set up for next reocmpression */
202*212397c6Schristos ret = inflateReset(&inf);
203*212397c6Schristos assert(ret != Z_STREAM_ERROR);
204*212397c6Schristos ret = deflateReset(&def);
205*212397c6Schristos assert(ret != Z_STREAM_ERROR);
206*212397c6Schristos
207*212397c6Schristos /* do second and final recompression (third compression) */
208*212397c6Schristos inf.avail_in = size - MARGIN; /* assure stream will complete */
209*212397c6Schristos inf.next_in = tmp;
210*212397c6Schristos def.avail_out = size;
211*212397c6Schristos def.next_out = blk;
212*212397c6Schristos ret = recompress(&inf, &def);
213*212397c6Schristos if (ret == Z_MEM_ERROR)
214*212397c6Schristos quit("out of memory");
215*212397c6Schristos assert(ret == Z_STREAM_END); /* otherwise MARGIN too small */
216*212397c6Schristos
217*212397c6Schristos /* done -- write block to stdout */
218*212397c6Schristos have = size - def.avail_out;
219*212397c6Schristos if (fwrite(blk, 1, have, stdout) != have || ferror(stdout))
220*212397c6Schristos quit("error writing output");
221*212397c6Schristos
222*212397c6Schristos /* clean up and print results to stderr */
223*212397c6Schristos free(tmp);
224*212397c6Schristos ret = inflateEnd(&inf);
225*212397c6Schristos assert(ret != Z_STREAM_ERROR);
226*212397c6Schristos ret = deflateEnd(&def);
227*212397c6Schristos assert(ret != Z_STREAM_ERROR);
228*212397c6Schristos free(blk);
229*212397c6Schristos fprintf(stderr,
230*212397c6Schristos "%u bytes unused out of %u requested (%lu input)\n",
231*212397c6Schristos size - have, size, def.total_in);
232*212397c6Schristos return 0;
233*212397c6Schristos }
234