1b430912de7521d5b2abb622f6b1a694f7ab6281
[wimlib] / src / xpress_compress.c
1 /*
2  * xpress_compress.c
3  *
4  * A compressor for the XPRESS compression format (Huffman variant).
5  */
6
7 /*
8  * Copyright (C) 2012, 2013, 2014 Eric Biggers
9  *
10  * This file is free software; you can redistribute it and/or modify it under
11  * the terms of the GNU Lesser General Public License as published by the Free
12  * Software Foundation; either version 3 of the License, or (at your option) any
13  * later version.
14  *
15  * This file is distributed in the hope that it will be useful, but WITHOUT
16  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
17  * FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more
18  * details.
19  *
20  * You should have received a copy of the GNU Lesser General Public License
21  * along with this file; if not, see http://www.gnu.org/licenses/.
22  */
23
24 #ifdef HAVE_CONFIG_H
25 #  include "config.h"
26 #endif
27
28 /*
29  * The maximum buffer size, in bytes, that can be compressed.  An XPRESS
30  * compressor instance must be created with a 'max_bufsize' less than or equal
31  * to this value.
32  */
33 #define XPRESS_MAX_BUFSIZE              65536
34
35 /*
36  * Define to 1 to enable the near-optimal parsing algorithm at high compression
37  * levels.  The near-optimal parsing algorithm produces a compression ratio
38  * significantly better than the greedy and lazy algorithms.  However, it is
39  * much slower.
40  */
41 #define SUPPORT_NEAR_OPTIMAL_PARSING    1
42
43 /*
44  * The lowest compression level at which near-optimal parsing is enabled.
45  */
46 #define MIN_LEVEL_FOR_NEAR_OPTIMAL      60
47
48 /*
49  * Matchfinder definitions.  For XPRESS, only a 16-bit matchfinder is needed.
50  */
51 #define mf_pos_t        u16
52 #define MF_SUFFIX
53
54 /*
55  * Note: although XPRESS can potentially use a sliding window, it isn't well
56  * suited for large buffers of data because there is no way to reset the Huffman
57  * code.  Therefore, we only allow buffers in which there is no restriction on
58  * match offsets (no sliding window).  This simplifies the code and allows some
59  * optimizations.
60  */
61
62 #include "wimlib/bitops.h"
63 #include "wimlib/compress_common.h"
64 #include "wimlib/compressor_ops.h"
65 #include "wimlib/endianness.h"
66 #include "wimlib/error.h"
67 #include "wimlib/hc_matchfinder.h"
68 #include "wimlib/unaligned.h"
69 #include "wimlib/util.h"
70 #include "wimlib/xpress_constants.h"
71
72 #if SUPPORT_NEAR_OPTIMAL_PARSING
73
74 /*
75  * CACHE_RESERVE_PER_POS is the number of lz_match structures to reserve in the
76  * match cache for each byte position.  This value should be high enough so that
77  * virtually the time, all matches found in the input buffer can fit in the
78  * match cache.  However, fallback behavior on cache overflow is still required.
79  */
80 #define CACHE_RESERVE_PER_POS   8
81
82 /*
83  * We use a binary-tree based matchfinder for optimal parsing because it can
84  * find more matches in the same number of steps compared to hash-chain based
85  * matchfinders.  In addition, since we need to find matches at almost every
86  * position, there isn't much penalty for keeping the sequences sorted in the
87  * binary trees.
88  */
89 #include "wimlib/bt_matchfinder.h"
90
91 struct xpress_optimum_node;
92
93 #endif /* SUPPORT_NEAR_OPTIMAL_PARSING */
94
95 struct xpress_item;
96
97 /* The main XPRESS compressor structure  */
98 struct xpress_compressor {
99
100         /* Pointer to the compress() implementation chosen at allocation time */
101         size_t (*impl)(struct xpress_compressor *,
102                        const void *, size_t, void *, size_t);
103
104         /* Symbol frequency counters for the Huffman code  */
105         u32 freqs[XPRESS_NUM_SYMBOLS];
106
107         /* The Huffman codewords and their lengths  */
108         u32 codewords[XPRESS_NUM_SYMBOLS];
109         u8 lens[XPRESS_NUM_SYMBOLS];
110
111         /* The "nice" match length: if a match of this length is found, then
112          * choose it immediately without further consideration.  */
113         unsigned nice_match_length;
114
115         /* The maximum search depth: consider at most this many potential
116          * matches at each position.  */
117         unsigned max_search_depth;
118
119         union {
120                 /* Data for greedy or lazy parsing  */
121                 struct {
122                         struct xpress_item *chosen_items;
123                         struct hc_matchfinder hc_mf;
124                         /* hc_mf must be last!  */
125                 };
126
127         #if SUPPORT_NEAR_OPTIMAL_PARSING
128                 /* Data for near-optimal parsing  */
129                 struct {
130                         struct xpress_optimum_node *optimum_nodes;
131                         struct lz_match *match_cache;
132                         struct lz_match *cache_overflow_mark;
133                         unsigned num_optim_passes;
134                         u32 costs[XPRESS_NUM_SYMBOLS];
135                         struct bt_matchfinder bt_mf;
136                         /* bt_mf must be last!  */
137                 };
138         #endif
139         };
140 };
141
142 #if SUPPORT_NEAR_OPTIMAL_PARSING
143
144 /*
145  * This structure represents a byte position in the input buffer and a node in
146  * the graph of possible match/literal choices.
147  *
148  * Logically, each incoming edge to this node is labeled with a literal or a
149  * match that can be taken to reach this position from an earlier position; and
150  * each outgoing edge from this node is labeled with a literal or a match that
151  * can be taken to advance from this position to a later position.
152  *
153  * But these "edges" are actually stored elsewhere (in 'match_cache').  Here we
154  * associate with each node just two pieces of information:
155  *
156  *      'cost_to_end' is the minimum cost to reach the end of the buffer from
157  *      this position.
158  *
159  *      'item' represents the literal or match that must be chosen from here to
160  *      reach the end of the buffer with the minimum cost.  Equivalently, this
161  *      can be interpreted as the label of the outgoing edge on the minimum cost
162  *      path to the "end of buffer" node from this node.
163  */
164 struct xpress_optimum_node {
165
166         u32 cost_to_end;
167
168         /*
169          * Notes on the match/literal representation used here:
170          *
171          *      The low bits of 'item' are the length: 1 if the item is a
172          *      literal, or the match length if the item is a match.
173          *
174          *      The high bits of 'item' are the actual literal byte if the item
175          *      is a literal, or the match offset if the item is a match.
176          */
177 #define OPTIMUM_OFFSET_SHIFT    16
178 #define OPTIMUM_LEN_MASK        (((u32)1 << OPTIMUM_OFFSET_SHIFT) - 1)
179         u32 item;
180 };
181
182 #endif /* SUPPORT_NEAR_OPTIMAL_PARSING */
183
184 /* An intermediate representation of an XPRESS match or literal  */
185 struct xpress_item {
186         /*
187          * Bits 0  -  8: Symbol
188          * Bits 9  - 24: Length - XPRESS_MIN_MATCH_LEN
189          * Bits 25 - 28: Number of extra offset bits
190          * Bits 29+    : Extra offset bits
191          *
192          * Unfortunately, gcc generates worse code if we use real bitfields here.
193          */
194         u64 data;
195 };
196
197 /*
198  * Structure to keep track of the current state of sending compressed data to
199  * the output buffer.
200  *
201  * The XPRESS bitstream is encoded as a sequence of little endian 16-bit coding
202  * units interwoven with literal bytes.
203  */
204 struct xpress_output_bitstream {
205
206         /* Bits that haven't yet been written to the output buffer.  */
207         u32 bitbuf;
208
209         /* Number of bits currently held in @bitbuf.  */
210         u32 bitcount;
211
212         /* Pointer to the start of the output buffer.  */
213         u8 *start;
214
215         /* Pointer to the location in the output buffer at which to write the
216          * next 16 bits.  */
217         u8 *next_bits;
218
219         /* Pointer to the location in the output buffer at which to write the
220          * next 16 bits, after @next_bits.  */
221         u8 *next_bits2;
222
223         /* Pointer to the location in the output buffer at which to write the
224          * next literal byte.  */
225         u8 *next_byte;
226
227         /* Pointer to the end of the output buffer.  */
228         u8 *end;
229 };
230
231 /* Reset the symbol frequencies for the XPRESS Huffman code.  */
232 static void
233 xpress_reset_symbol_frequencies(struct xpress_compressor *c)
234 {
235         memset(c->freqs, 0, sizeof(c->freqs));
236 }
237
238 /*
239  * Make the Huffman code for XPRESS.
240  *
241  * Input: c->freqs
242  * Output: c->lens and c->codewords
243  */
244 static void
245 xpress_make_huffman_code(struct xpress_compressor *c)
246 {
247         make_canonical_huffman_code(XPRESS_NUM_SYMBOLS, XPRESS_MAX_CODEWORD_LEN,
248                                     c->freqs, c->lens, c->codewords);
249 }
250
251 /*
252  * Initialize the output bitstream.
253  *
254  * @os
255  *      The output bitstream structure to initialize.
256  * @buffer
257  *      The output buffer.
258  * @size
259  *      Size of @buffer, in bytes.  Must be at least 4.
260  */
261 static void
262 xpress_init_output(struct xpress_output_bitstream *os, void *buffer, size_t size)
263 {
264         os->bitbuf = 0;
265         os->bitcount = 0;
266         os->start = buffer;
267         os->next_bits = os->start;
268         os->next_bits2 = os->start + 2;
269         os->next_byte = os->start + 4;
270         os->end = os->start + size;
271 }
272
273 /*
274  * Write some bits to the output bitstream.
275  *
276  * The bits are given by the low-order @num_bits bits of @bits.  Higher-order
277  * bits in @bits cannot be set.  At most 16 bits can be written at once.
278  *
279  * If the output buffer space is exhausted, then the bits will be ignored, and
280  * xpress_flush_output() will return 0 when it gets called.
281  */
282 static forceinline void
283 xpress_write_bits(struct xpress_output_bitstream *os,
284                   const u32 bits, const unsigned num_bits)
285 {
286         /* This code is optimized for XPRESS, which never needs to write more
287          * than 16 bits at once.  */
288
289         os->bitcount += num_bits;
290         os->bitbuf = (os->bitbuf << num_bits) | bits;
291
292         if (os->bitcount > 16) {
293                 os->bitcount -= 16;
294                 if (os->end - os->next_byte >= 2) {
295                         put_unaligned_le16(os->bitbuf >> os->bitcount, os->next_bits);
296                         os->next_bits = os->next_bits2;
297                         os->next_bits2 = os->next_byte;
298                         os->next_byte += 2;
299                 }
300         }
301 }
302
303 /*
304  * Interweave a literal byte into the output bitstream.
305  */
306 static forceinline void
307 xpress_write_byte(struct xpress_output_bitstream *os, u8 byte)
308 {
309         if (os->next_byte < os->end)
310                 *os->next_byte++ = byte;
311 }
312
313 /*
314  * Interweave two literal bytes into the output bitstream.
315  */
316 static forceinline void
317 xpress_write_u16(struct xpress_output_bitstream *os, u16 v)
318 {
319         if (os->end - os->next_byte >= 2) {
320                 put_unaligned_le16(v, os->next_byte);
321                 os->next_byte += 2;
322         }
323 }
324
325 /*
326  * Flush the last coding unit to the output buffer if needed.  Return the total
327  * number of bytes written to the output buffer, or 0 if an overflow occurred.
328  */
329 static size_t
330 xpress_flush_output(struct xpress_output_bitstream *os)
331 {
332         if (os->end - os->next_byte < 2)
333                 return 0;
334
335         put_unaligned_le16(os->bitbuf << (16 - os->bitcount), os->next_bits);
336         put_unaligned_le16(0, os->next_bits2);
337
338         return os->next_byte - os->start;
339 }
340
341 static forceinline void
342 xpress_write_extra_length_bytes(struct xpress_output_bitstream *os,
343                                 unsigned adjusted_len)
344 {
345         /* If length >= 18, output one extra length byte.
346          * If length >= 273, output three (total) extra length bytes.  */
347         if (adjusted_len >= 0xF) {
348                 u8 byte1 = min(adjusted_len - 0xF, 0xFF);
349                 xpress_write_byte(os, byte1);
350                 if (byte1 == 0xFF)
351                         xpress_write_u16(os, adjusted_len);
352         }
353 }
354
355 /* Output a match or literal.  */
356 static forceinline void
357 xpress_write_item(struct xpress_item item, struct xpress_output_bitstream *os,
358                   const u32 codewords[], const u8 lens[])
359 {
360         u64 data = item.data;
361         unsigned symbol = data & 0x1FF;
362
363         xpress_write_bits(os, codewords[symbol], lens[symbol]);
364
365         if (symbol >= XPRESS_NUM_CHARS) {
366                 /* Match, not a literal  */
367                 xpress_write_extra_length_bytes(os, (data >> 9) & 0xFFFF);
368                 xpress_write_bits(os, data >> 29, (data >> 25) & 0xF);
369         }
370 }
371
372 /* Output a sequence of XPRESS matches and literals.  */
373 static void
374 xpress_write_items(struct xpress_output_bitstream *os,
375                    const struct xpress_item items[], size_t num_items,
376                    const u32 codewords[], const u8 lens[])
377 {
378         for (size_t i = 0; i < num_items; i++)
379                 xpress_write_item(items[i], os, codewords, lens);
380 }
381
382 #if SUPPORT_NEAR_OPTIMAL_PARSING
383
384 /*
385  * Follow the minimum cost path in the graph of possible match/literal choices
386  * and write out the matches/literals using the specified Huffman code.
387  *
388  * Note: this is slightly duplicated with xpress_write_items().  However, we
389  * don't want to waste time translating between intermediate match/literal
390  * representations.
391  */
392 static void
393 xpress_write_item_list(struct xpress_output_bitstream *os,
394                        struct xpress_optimum_node *optimum_nodes,
395                        size_t count, const u32 codewords[], const u8 lens[])
396 {
397         struct xpress_optimum_node *cur_node = optimum_nodes;
398         struct xpress_optimum_node *end_node = optimum_nodes + count;
399         do {
400                 unsigned length = cur_node->item & OPTIMUM_LEN_MASK;
401                 unsigned offset = cur_node->item >> OPTIMUM_OFFSET_SHIFT;
402
403                 if (length == 1) {
404                         /* Literal  */
405                         unsigned literal = offset;
406
407                         xpress_write_bits(os, codewords[literal], lens[literal]);
408                 } else {
409                         /* Match  */
410                         unsigned adjusted_len;
411                         unsigned log2_offset;
412                         unsigned len_hdr;
413                         unsigned sym;
414
415                         adjusted_len = length - XPRESS_MIN_MATCH_LEN;
416                         log2_offset = bsr32(offset);
417                         len_hdr = min(0xF, adjusted_len);
418                         sym = XPRESS_NUM_CHARS + ((log2_offset << 4) | len_hdr);
419
420                         xpress_write_bits(os, codewords[sym], lens[sym]);
421                         xpress_write_extra_length_bytes(os, adjusted_len);
422                         xpress_write_bits(os, offset - (1U << log2_offset),
423                                           log2_offset);
424                 }
425                 cur_node += length;
426         } while (cur_node != end_node);
427 }
428 #endif /* SUPPORT_NEAR_OPTIMAL_PARSING */
429
430 /*
431  * Output the XPRESS-compressed data, given the sequence of match/literal
432  * "items" that was chosen to represent the input data.
433  *
434  * If @near_optimal is %false, then the items are taken from the array
435  * c->chosen_items[0...count].
436  *
437  * If @near_optimal is %true, then the items are taken from the minimum cost
438  * path stored in c->optimum_nodes[0...count].
439  */
440 static size_t
441 xpress_write(struct xpress_compressor *c, void *out, size_t out_nbytes_avail,
442              size_t count, bool near_optimal)
443 {
444         u8 *cptr;
445         struct xpress_output_bitstream os;
446         size_t out_size;
447
448         /* Account for the end-of-data symbol and make the Huffman code.  */
449         c->freqs[XPRESS_END_OF_DATA]++;
450         xpress_make_huffman_code(c);
451
452         /* Output the Huffman code as a series of 512 4-bit lengths.  */
453         cptr = out;
454         for (unsigned i = 0; i < XPRESS_NUM_SYMBOLS; i += 2)
455                 *cptr++ = (c->lens[i + 1] << 4) | c->lens[i];
456
457         xpress_init_output(&os, cptr, out_nbytes_avail - XPRESS_NUM_SYMBOLS / 2);
458
459         /* Output the Huffman-encoded items.  */
460 #if SUPPORT_NEAR_OPTIMAL_PARSING
461         if (near_optimal) {
462                 xpress_write_item_list(&os, c->optimum_nodes, count,
463                                        c->codewords, c->lens);
464
465         } else
466 #endif
467         {
468                 xpress_write_items(&os, c->chosen_items, count,
469                                    c->codewords, c->lens);
470         }
471
472         /* Write the end-of-data symbol (needed for MS compatibility)  */
473         xpress_write_bits(&os, c->codewords[XPRESS_END_OF_DATA],
474                           c->lens[XPRESS_END_OF_DATA]);
475
476         /* Flush any pending data.  Then return the compressed size if the
477          * compressed data fit in the output buffer, or 0 if it did not.  */
478         out_size = xpress_flush_output(&os);
479         if (out_size == 0)
480                 return 0;
481
482         return out_size + XPRESS_NUM_SYMBOLS / 2;
483 }
484
485 /* Tally the Huffman symbol for a literal and return the intermediate
486  * representation of that literal.  */
487 static forceinline struct xpress_item
488 xpress_record_literal(struct xpress_compressor *c, unsigned literal)
489 {
490         c->freqs[literal]++;
491
492         return (struct xpress_item) {
493                 .data = literal,
494         };
495 }
496
497 /* Tally the Huffman symbol for a match and return the intermediate
498  * representation of that match.  */
499 static forceinline struct xpress_item
500 xpress_record_match(struct xpress_compressor *c, unsigned length, unsigned offset)
501 {
502         unsigned adjusted_len = length - XPRESS_MIN_MATCH_LEN;
503         unsigned len_hdr = min(adjusted_len, 0xF);
504         unsigned log2_offset = bsr32(offset);
505         unsigned sym = XPRESS_NUM_CHARS + ((log2_offset << 4) | len_hdr);
506
507         c->freqs[sym]++;
508
509         return (struct xpress_item) {
510                 .data = (u64)sym |
511                         ((u64)adjusted_len << 9) |
512                         ((u64)log2_offset << 25) |
513                         ((u64)(offset ^ (1U << log2_offset)) << 29),
514         };
515 }
516
517 /*
518  * This is the "greedy" XPRESS compressor. It always chooses the longest match.
519  * (Exception: as a heuristic, we pass up length 3 matches that have large
520  * offsets.)
521  */
522 static size_t
523 xpress_compress_greedy(struct xpress_compressor * restrict c,
524                        const void * restrict in, size_t in_nbytes,
525                        void * restrict out, size_t out_nbytes_avail)
526 {
527         const u8 * const in_begin = in;
528         const u8 *       in_next = in_begin;
529         const u8 * const in_end = in_begin + in_nbytes;
530         struct xpress_item *next_chosen_item = c->chosen_items;
531         unsigned len_3_too_far;
532         u32 next_hashes[2] = {};
533
534         if (in_nbytes <= 8192)
535                 len_3_too_far = 2048;
536         else
537                 len_3_too_far = 4096;
538
539         hc_matchfinder_init(&c->hc_mf);
540
541         do {
542                 unsigned length;
543                 unsigned offset;
544
545                 length = hc_matchfinder_longest_match(&c->hc_mf,
546                                                       in_begin,
547                                                       in_next - in_begin,
548                                                       XPRESS_MIN_MATCH_LEN - 1,
549                                                       in_end - in_next,
550                                                       min(in_end - in_next, c->nice_match_length),
551                                                       c->max_search_depth,
552                                                       next_hashes,
553                                                       &offset);
554                 if (length >= XPRESS_MIN_MATCH_LEN &&
555                     !(length == XPRESS_MIN_MATCH_LEN && offset >= len_3_too_far))
556                 {
557                         /* Match found  */
558                         *next_chosen_item++ =
559                                 xpress_record_match(c, length, offset);
560                         in_next += 1;
561                         hc_matchfinder_skip_positions(&c->hc_mf,
562                                                       in_begin,
563                                                       in_next - in_begin,
564                                                       in_end - in_begin,
565                                                       length - 1,
566                                                       next_hashes);
567                         in_next += length - 1;
568                 } else {
569                         /* No match found  */
570                         *next_chosen_item++ =
571                                 xpress_record_literal(c, *in_next);
572                         in_next += 1;
573                 }
574         } while (in_next != in_end);
575
576         return xpress_write(c, out, out_nbytes_avail,
577                             next_chosen_item - c->chosen_items, false);
578 }
579
580 /*
581  * This is the "lazy" XPRESS compressor.  Before choosing a match, it checks to
582  * see if there's a longer match at the next position.  If yes, it outputs a
583  * literal and continues to the next position.  If no, it outputs the match.
584  */
585 static size_t
586 xpress_compress_lazy(struct xpress_compressor * restrict c,
587                      const void * restrict in, size_t in_nbytes,
588                      void * restrict out, size_t out_nbytes_avail)
589 {
590         const u8 * const in_begin = in;
591         const u8 *       in_next = in_begin;
592         const u8 * const in_end = in_begin + in_nbytes;
593         struct xpress_item *next_chosen_item = c->chosen_items;
594         unsigned len_3_too_far;
595         u32 next_hashes[2] = {};
596
597         if (in_nbytes <= 8192)
598                 len_3_too_far = 2048;
599         else
600                 len_3_too_far = 4096;
601
602         hc_matchfinder_init(&c->hc_mf);
603
604         do {
605                 unsigned cur_len;
606                 unsigned cur_offset;
607                 unsigned next_len;
608                 unsigned next_offset;
609
610                 /* Find the longest match at the current position.  */
611                 cur_len = hc_matchfinder_longest_match(&c->hc_mf,
612                                                        in_begin,
613                                                        in_next - in_begin,
614                                                        XPRESS_MIN_MATCH_LEN - 1,
615                                                        in_end - in_next,
616                                                        min(in_end - in_next, c->nice_match_length),
617                                                        c->max_search_depth,
618                                                        next_hashes,
619                                                        &cur_offset);
620                 in_next += 1;
621
622                 if (cur_len < XPRESS_MIN_MATCH_LEN ||
623                     (cur_len == XPRESS_MIN_MATCH_LEN &&
624                      cur_offset >= len_3_too_far))
625                 {
626                         /* No match found.  Choose a literal.  */
627                         *next_chosen_item++ =
628                                 xpress_record_literal(c, *(in_next - 1));
629                         continue;
630                 }
631
632         have_cur_match:
633                 /* We have a match at the current position.  */
634
635                 /* If the current match is very long, choose it immediately.  */
636                 if (cur_len >= c->nice_match_length) {
637
638                         *next_chosen_item++ =
639                                 xpress_record_match(c, cur_len, cur_offset);
640
641                         hc_matchfinder_skip_positions(&c->hc_mf,
642                                                       in_begin,
643                                                       in_next - in_begin,
644                                                       in_end - in_begin,
645                                                       cur_len - 1,
646                                                       next_hashes);
647                         in_next += cur_len - 1;
648                         continue;
649                 }
650
651                 /*
652                  * Try to find a match at the next position.
653                  *
654                  * Note: since we already have a match at the *current*
655                  * position, we use only half the 'max_search_depth' when
656                  * checking the *next* position.  This is a useful trade-off
657                  * because it's more worthwhile to use a greater search depth on
658                  * the initial match than on the next match (since a lot of the
659                  * time, that next match won't even be used).
660                  *
661                  * Note: it's possible to structure the code such that there's
662                  * only one call to longest_match(), which handles both the
663                  * "find the initial match" and "try to find a longer match"
664                  * cases.  However, it is faster to have two call sites, with
665                  * longest_match() inlined at each.
666                  */
667                 next_len = hc_matchfinder_longest_match(&c->hc_mf,
668                                                         in_begin,
669                                                         in_next - in_begin,
670                                                         cur_len,
671                                                         in_end - in_next,
672                                                         min(in_end - in_next, c->nice_match_length),
673                                                         c->max_search_depth / 2,
674                                                         next_hashes,
675                                                         &next_offset);
676                 in_next += 1;
677
678                 if (next_len > cur_len) {
679                         /* Found a longer match at the next position, so output
680                          * a literal.  */
681                         *next_chosen_item++ =
682                                 xpress_record_literal(c, *(in_next - 2));
683                         cur_len = next_len;
684                         cur_offset = next_offset;
685                         goto have_cur_match;
686                 } else {
687                         /* Didn't find a longer match at the next position, so
688                          * output the current match.  */
689                         *next_chosen_item++ =
690                                 xpress_record_match(c, cur_len, cur_offset);
691                         hc_matchfinder_skip_positions(&c->hc_mf,
692                                                       in_begin,
693                                                       in_next - in_begin,
694                                                       in_end - in_begin,
695                                                       cur_len - 2,
696                                                       next_hashes);
697                         in_next += cur_len - 2;
698                         continue;
699                 }
700         } while (in_next != in_end);
701
702         return xpress_write(c, out, out_nbytes_avail,
703                             next_chosen_item - c->chosen_items, false);
704 }
705
706 #if SUPPORT_NEAR_OPTIMAL_PARSING
707
708 /*
709  * Set Huffman symbol costs for the first optimization pass.
710  *
711  * It works well to assume that each Huffman symbol is equally probable.  This
712  * results in each symbol being assigned a cost of -log2(1.0/num_syms) where
713  * 'num_syms' is the number of symbols in the alphabet.
714  */
715 static void
716 xpress_set_default_costs(struct xpress_compressor *c)
717 {
718         for (unsigned i = 0; i < XPRESS_NUM_SYMBOLS; i++)
719                 c->costs[i] = 9;
720 }
721
722 /* Update the cost model based on the codeword lengths @c->lens.  */
723 static void
724 xpress_update_costs(struct xpress_compressor *c)
725 {
726         for (unsigned i = 0; i < XPRESS_NUM_SYMBOLS; i++)
727                 c->costs[i] = c->lens[i] ? c->lens[i] : XPRESS_MAX_CODEWORD_LEN;
728 }
729
730 /*
731  * Follow the minimum cost path in the graph of possible match/literal choices
732  * and compute the frequencies of the Huffman symbols that are needed to output
733  * those matches and literals.
734  */
735 static void
736 xpress_tally_item_list(struct xpress_compressor *c,
737                        struct xpress_optimum_node *end_node)
738 {
739         struct xpress_optimum_node *cur_node = c->optimum_nodes;
740
741         do {
742                 unsigned length = cur_node->item & OPTIMUM_LEN_MASK;
743                 unsigned offset = cur_node->item >> OPTIMUM_OFFSET_SHIFT;
744
745                 if (length == 1) {
746                         /* Literal  */
747                         unsigned literal = offset;
748
749                         c->freqs[literal]++;
750                 } else {
751                         /* Match  */
752                         unsigned adjusted_len;
753                         unsigned log2_offset;
754                         unsigned len_hdr;
755                         unsigned sym;
756
757                         adjusted_len = length - XPRESS_MIN_MATCH_LEN;
758                         log2_offset = bsr32(offset);
759                         len_hdr = min(0xF, adjusted_len);
760                         sym = XPRESS_NUM_CHARS + ((log2_offset << 4) | len_hdr);
761
762                         c->freqs[sym]++;
763                 }
764                 cur_node += length;
765         } while (cur_node != end_node);
766 }
767
768 /*
769  * Find a new minimum cost path through the graph of possible match/literal
770  * choices.  We find the minimum cost path from 'c->optimum_nodes[0]', which
771  * represents the node at the beginning of the input buffer, to
772  * 'c->optimum_nodes[in_nbytes]', which represents the node at the end of the
773  * input buffer.  Edge costs are evaluated using the cost model 'c->costs'.
774  *
775  * The algorithm works backward, starting at 'c->optimum_nodes[in_nbytes]' and
776  * proceeding backwards one position at a time.  At each position, the minimum
777  * cost to reach 'c->optimum_nodes[in_nbytes]' from that position is computed
778  * and the match/literal choice is saved.
779  */
780 static void
781 xpress_find_min_cost_path(struct xpress_compressor *c, size_t in_nbytes,
782                           struct lz_match *end_cache_ptr)
783 {
784         struct xpress_optimum_node *cur_node = c->optimum_nodes + in_nbytes;
785         struct lz_match *cache_ptr = end_cache_ptr;
786
787         cur_node->cost_to_end = 0;
788         do {
789                 unsigned literal;
790                 u32 best_item;
791                 u32 best_cost_to_end;
792                 unsigned num_matches;
793                 struct lz_match *match;
794                 unsigned len;
795
796                 cur_node--;
797                 cache_ptr--;
798
799                 literal = cache_ptr->offset;
800
801                 /* Consider coding a literal.  */
802                 best_item = ((u32)literal << OPTIMUM_OFFSET_SHIFT) | 1;
803                 best_cost_to_end = c->costs[literal] +
804                                    (cur_node + 1)->cost_to_end;
805
806                 num_matches = cache_ptr->length;
807
808                 if (num_matches == 0) {
809                         /* No matches; the only choice is the literal.  */
810                         cur_node->cost_to_end = best_cost_to_end;
811                         cur_node->item = best_item;
812                         continue;
813                 }
814
815                 /*
816                  * Consider each match length from the minimum
817                  * (XPRESS_MIN_MATCH_LEN) to the length of the longest match
818                  * found at this position.  For each length, consider only the
819                  * smallest offset for which that length is available.  Although
820                  * this is not guaranteed to be optimal due to the possibility
821                  * of a larger offset costing less than a smaller offset to
822                  * code, this is a very useful heuristic.
823                  */
824                 match = cache_ptr - num_matches;
825                 len = XPRESS_MIN_MATCH_LEN;
826                 if (cache_ptr[-1].length < 0xF + XPRESS_MIN_MATCH_LEN) {
827                         /* All lengths are small.  Optimize accordingly.  */
828                         do {
829                                 unsigned offset;
830                                 unsigned log2_offset;
831                                 u32 offset_cost;
832
833                                 offset = match->offset;
834                                 log2_offset = bsr32(offset);
835                                 offset_cost = log2_offset;
836                                 do {
837                                         unsigned len_hdr;
838                                         unsigned sym;
839                                         u32 cost_to_end;
840
841                                         len_hdr = len - XPRESS_MIN_MATCH_LEN;
842                                         sym = XPRESS_NUM_CHARS +
843                                               ((log2_offset << 4) | len_hdr);
844                                         cost_to_end =
845                                                 offset_cost + c->costs[sym] +
846                                                 (cur_node + len)->cost_to_end;
847                                         if (cost_to_end < best_cost_to_end) {
848                                                 best_cost_to_end = cost_to_end;
849                                                 best_item =
850                                                         ((u32)offset <<
851                                                          OPTIMUM_OFFSET_SHIFT) | len;
852                                         }
853                                 } while (++len <= match->length);
854                         } while (++match != cache_ptr);
855                 } else {
856                         /* Some lengths are big.  */
857                         do {
858                                 unsigned offset;
859                                 unsigned log2_offset;
860                                 u32 offset_cost;
861
862                                 offset = match->offset;
863                                 log2_offset = bsr32(offset);
864                                 offset_cost = log2_offset;
865                                 do {
866                                         unsigned adjusted_len;
867                                         unsigned len_hdr;
868                                         unsigned sym;
869                                         u32 cost_to_end;
870
871                                         adjusted_len = len - XPRESS_MIN_MATCH_LEN;
872                                         len_hdr = min(adjusted_len, 0xF);
873                                         sym = XPRESS_NUM_CHARS +
874                                               ((log2_offset << 4) | len_hdr);
875                                         cost_to_end =
876                                                 offset_cost + c->costs[sym] +
877                                                 (cur_node + len)->cost_to_end;
878                                         if (adjusted_len >= 0xF) {
879                                                 cost_to_end += 8;
880                                                 if (adjusted_len - 0xF >= 0xFF)
881                                                         cost_to_end += 16;
882                                         }
883                                         if (cost_to_end < best_cost_to_end) {
884                                                 best_cost_to_end = cost_to_end;
885                                                 best_item =
886                                                         ((u32)offset <<
887                                                          OPTIMUM_OFFSET_SHIFT) | len;
888                                         }
889                                 } while (++len <= match->length);
890                         } while (++match != cache_ptr);
891                 }
892                 cache_ptr -= num_matches;
893                 cur_node->cost_to_end = best_cost_to_end;
894                 cur_node->item = best_item;
895         } while (cur_node != c->optimum_nodes);
896 }
897
898 /*
899  * This routine finds matches at each position in the buffer in[0...in_nbytes].
900  * The matches are cached in the array c->match_cache, and the return value is a
901  * pointer past the last slot in this array that was filled.
902  */
903 static struct lz_match *
904 xpress_find_matches(struct xpress_compressor * restrict c,
905                     const void * restrict in, size_t in_nbytes)
906 {
907         const u8 * const in_begin = in;
908         const u8 *in_next = in_begin;
909         struct lz_match *cache_ptr = c->match_cache;
910         u32 next_hashes[2] = {};
911         u32 max_len = in_nbytes;
912         u32 nice_len = min(max_len, c->nice_match_length);
913
914         bt_matchfinder_init(&c->bt_mf);
915
916         for (;;) {
917                 struct lz_match *matches;
918                 u32 best_len;
919
920                 /* If we've found so many matches that the cache might overflow
921                  * if we keep finding more, then stop finding matches.  This
922                  * case is very unlikely.  */
923                 if (unlikely(cache_ptr >= c->cache_overflow_mark ||
924                              max_len < BT_MATCHFINDER_REQUIRED_NBYTES))
925                         break;
926
927                 matches = cache_ptr;
928
929                 /* Find matches with the current position using the binary tree
930                  * matchfinder and save them in the next available slots in
931                  * the match cache.  */
932                 cache_ptr =
933                         bt_matchfinder_get_matches(&c->bt_mf,
934                                                    in_begin,
935                                                    in_next - in_begin,
936                                                    max_len,
937                                                    nice_len,
938                                                    c->max_search_depth,
939                                                    next_hashes,
940                                                    &best_len,
941                                                    cache_ptr);
942                 cache_ptr->length = cache_ptr - matches;
943                 cache_ptr->offset = *in_next++;
944                 cache_ptr++;
945                 max_len--;
946                 nice_len = min(nice_len, max_len);
947
948                 /*
949                  * If there was a very long match found, then don't cache any
950                  * matches for the bytes covered by that match.  This avoids
951                  * degenerate behavior when compressing highly redundant data,
952                  * where the number of matches can be very large.
953                  *
954                  * This heuristic doesn't actually hurt the compression ratio
955                  * very much.  If there's a long match, then the data must be
956                  * highly compressible, so it doesn't matter as much what we do.
957                  */
958                 if (best_len >= nice_len) {
959                         if (unlikely(best_len +
960                                      BT_MATCHFINDER_REQUIRED_NBYTES >= max_len))
961                                 break;
962                         --best_len;
963                         do {
964                                 bt_matchfinder_skip_position(&c->bt_mf,
965                                                              in_begin,
966                                                              in_next - in_begin,
967                                                              nice_len,
968                                                              c->max_search_depth,
969                                                              next_hashes);
970                                 cache_ptr->length = 0;
971                                 cache_ptr->offset = *in_next++;
972                                 cache_ptr++;
973                                 max_len--;
974                                 nice_len = min(nice_len, max_len);
975                         } while (--best_len);
976                 }
977         }
978
979         while (max_len--) {
980                 cache_ptr->length = 0;
981                 cache_ptr->offset = *in_next++;
982                 cache_ptr++;
983         }
984
985         return cache_ptr;
986 }
987
988 /*
989  * This is the "near-optimal" XPRESS compressor.  It computes a compressed
990  * representation of the input buffer by executing a minimum cost path search
991  * over the graph of possible match/literal choices, assuming a certain cost for
992  * each Huffman symbol.  The result is usually close to optimal, but it is *not*
993  * guaranteed to be optimal because of (a) heuristic restrictions in which
994  * matches are considered, and (b) symbol costs are unknown until those symbols
995  * have already been chosen --- so iterative optimization must be used, and the
996  * algorithm might converge on a local optimum rather than a global optimum.
997  */
998 static size_t
999 xpress_compress_near_optimal(struct xpress_compressor * restrict c,
1000                              const void * restrict in, size_t in_nbytes,
1001                              void * restrict out, size_t out_nbytes_avail)
1002 {
1003         struct lz_match *end_cache_ptr;
1004         unsigned num_passes_remaining = c->num_optim_passes;
1005
1006         /* Run the input buffer through the matchfinder and save the results. */
1007         end_cache_ptr = xpress_find_matches(c, in, in_nbytes);
1008
1009         /* The first optimization pass uses a default cost model.  Each
1010          * additional optimization pass uses a cost model derived from the
1011          * Huffman code computed in the previous pass.  */
1012         xpress_set_default_costs(c);
1013         do {
1014                 xpress_find_min_cost_path(c, in_nbytes, end_cache_ptr);
1015                 xpress_tally_item_list(c, c->optimum_nodes + in_nbytes);
1016                 if (num_passes_remaining > 1) {
1017                         c->freqs[XPRESS_END_OF_DATA]++;
1018                         xpress_make_huffman_code(c);
1019                         xpress_update_costs(c);
1020                         xpress_reset_symbol_frequencies(c);
1021                 }
1022         } while (--num_passes_remaining);
1023
1024         return xpress_write(c, out, out_nbytes_avail, in_nbytes, true);
1025 }
1026
1027 #endif /* SUPPORT_NEAR_OPTIMAL_PARSING */
1028
1029 static size_t
1030 xpress_get_compressor_size(size_t max_bufsize, unsigned compression_level)
1031 {
1032 #if SUPPORT_NEAR_OPTIMAL_PARSING
1033         if (compression_level >= MIN_LEVEL_FOR_NEAR_OPTIMAL)
1034                 return offsetof(struct xpress_compressor, bt_mf) +
1035                         bt_matchfinder_size(max_bufsize);
1036 #endif
1037
1038         return offsetof(struct xpress_compressor, hc_mf) +
1039                 hc_matchfinder_size(max_bufsize);
1040 }
1041
1042 static u64
1043 xpress_get_needed_memory(size_t max_bufsize, unsigned compression_level,
1044                          bool destructive)
1045 {
1046         u64 size = 0;
1047
1048         if (max_bufsize > XPRESS_MAX_BUFSIZE)
1049                 return 0;
1050
1051         size += xpress_get_compressor_size(max_bufsize, compression_level);
1052
1053         if (compression_level < MIN_LEVEL_FOR_NEAR_OPTIMAL ||
1054             !SUPPORT_NEAR_OPTIMAL_PARSING) {
1055                 /* chosen_items  */
1056                 size += max_bufsize * sizeof(struct xpress_item);
1057         }
1058 #if SUPPORT_NEAR_OPTIMAL_PARSING
1059         else {
1060                 /* optimum_nodes  */
1061                 size += (max_bufsize + 1) * sizeof(struct xpress_optimum_node);
1062                 /* match_cache */
1063                 size += ((max_bufsize * CACHE_RESERVE_PER_POS) +
1064                          XPRESS_MAX_MATCH_LEN + max_bufsize) *
1065                                 sizeof(struct lz_match);
1066         }
1067 #endif
1068         return size;
1069 }
1070
1071 static int
1072 xpress_create_compressor(size_t max_bufsize, unsigned compression_level,
1073                          bool destructive, void **c_ret)
1074 {
1075         struct xpress_compressor *c;
1076
1077         if (max_bufsize > XPRESS_MAX_BUFSIZE)
1078                 return WIMLIB_ERR_INVALID_PARAM;
1079
1080         c = MALLOC(xpress_get_compressor_size(max_bufsize, compression_level));
1081         if (!c)
1082                 goto oom0;
1083
1084         if (compression_level < MIN_LEVEL_FOR_NEAR_OPTIMAL ||
1085             !SUPPORT_NEAR_OPTIMAL_PARSING)
1086         {
1087
1088                 c->chosen_items = MALLOC(max_bufsize * sizeof(struct xpress_item));
1089                 if (!c->chosen_items)
1090                         goto oom1;
1091
1092                 if (compression_level < 30) {
1093                         c->impl = xpress_compress_greedy;
1094                         c->max_search_depth = (compression_level * 30) / 16;
1095                         c->nice_match_length = (compression_level * 60) / 16;
1096                 } else {
1097                         c->impl = xpress_compress_lazy;
1098                         c->max_search_depth = (compression_level * 30) / 32;
1099                         c->nice_match_length = (compression_level * 60) / 32;
1100
1101                         /* xpress_compress_lazy() needs max_search_depth >= 2
1102                          * because it halves the max_search_depth when
1103                          * attempting a lazy match, and max_search_depth cannot
1104                          * be 0.  */
1105                         if (c->max_search_depth < 2)
1106                                 c->max_search_depth = 2;
1107                 }
1108         }
1109 #if SUPPORT_NEAR_OPTIMAL_PARSING
1110         else {
1111
1112                 c->optimum_nodes = MALLOC((max_bufsize + 1) *
1113                                           sizeof(struct xpress_optimum_node));
1114                 c->match_cache = MALLOC(((max_bufsize * CACHE_RESERVE_PER_POS) +
1115                                          XPRESS_MAX_MATCH_LEN + max_bufsize) *
1116                                         sizeof(struct lz_match));
1117                 if (!c->optimum_nodes || !c->match_cache) {
1118                         FREE(c->optimum_nodes);
1119                         FREE(c->match_cache);
1120                         goto oom1;
1121                 }
1122                 c->cache_overflow_mark =
1123                         &c->match_cache[max_bufsize * CACHE_RESERVE_PER_POS];
1124
1125                 c->impl = xpress_compress_near_optimal;
1126                 c->max_search_depth = (compression_level * 28) / 100;
1127                 c->nice_match_length = (compression_level * 56) / 100;
1128                 c->num_optim_passes = compression_level / 40;
1129         }
1130 #endif /* SUPPORT_NEAR_OPTIMAL_PARSING */
1131
1132         /* max_search_depth == 0 is invalid.  */
1133         if (c->max_search_depth < 1)
1134                 c->max_search_depth = 1;
1135
1136         *c_ret = c;
1137         return 0;
1138
1139 oom1:
1140         FREE(c);
1141 oom0:
1142         return WIMLIB_ERR_NOMEM;
1143 }
1144
1145 static size_t
1146 xpress_compress(const void *restrict in, size_t in_nbytes,
1147                 void *restrict out, size_t out_nbytes_avail, void *restrict _c)
1148 {
1149         struct xpress_compressor *c = _c;
1150
1151         /* Don't bother trying to compress very small inputs.  */
1152         if (in_nbytes < 25)
1153                 return 0;
1154
1155         if (out_nbytes_avail <= XPRESS_NUM_SYMBOLS / 2 + 4)
1156                 return 0;
1157
1158         xpress_reset_symbol_frequencies(c);
1159
1160         return (*c->impl)(c, in, in_nbytes, out, out_nbytes_avail);
1161 }
1162
1163 static void
1164 xpress_free_compressor(void *_c)
1165 {
1166         struct xpress_compressor *c = _c;
1167
1168 #if SUPPORT_NEAR_OPTIMAL_PARSING
1169         if (c->impl == xpress_compress_near_optimal) {
1170                 FREE(c->optimum_nodes);
1171                 FREE(c->match_cache);
1172         } else
1173 #endif
1174                 FREE(c->chosen_items);
1175         FREE(c);
1176 }
1177
1178 const struct compressor_ops xpress_compressor_ops = {
1179         .get_needed_memory  = xpress_get_needed_memory,
1180         .create_compressor  = xpress_create_compressor,
1181         .compress           = xpress_compress,
1182         .free_compressor    = xpress_free_compressor,
1183 };