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