65c4b026627bdeaf912b29e04cdba435779f7cef
[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
535         if (in_nbytes <= 8192)
536                 len_3_too_far = 2048;
537         else
538                 len_3_too_far = 4096;
539
540         hc_matchfinder_init(&c->hc_mf);
541
542         do {
543                 unsigned length;
544                 unsigned offset;
545
546                 length = hc_matchfinder_longest_match(&c->hc_mf,
547                                                       in_begin,
548                                                       in_next,
549                                                       XPRESS_MIN_MATCH_LEN - 1,
550                                                       in_end - in_next,
551                                                       min(in_end - in_next, c->nice_match_length),
552                                                       c->max_search_depth,
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,
564                                                       in_end,
565                                                       length - 1);
566                         in_next += length - 1;
567                 } else {
568                         /* No match found  */
569                         *next_chosen_item++ =
570                                 xpress_record_literal(c, *in_next);
571                         in_next += 1;
572                 }
573         } while (in_next != in_end);
574
575         return xpress_write(c, out, out_nbytes_avail,
576                             next_chosen_item - c->chosen_items, false);
577 }
578
579 /*
580  * This is the "lazy" XPRESS compressor.  Before choosing a match, it checks to
581  * see if there's a longer match at the next position.  If yes, it outputs a
582  * literal and continues to the next position.  If no, it outputs the match.
583  */
584 static size_t
585 xpress_compress_lazy(struct xpress_compressor * restrict c,
586                      const void * restrict in, size_t in_nbytes,
587                      void * restrict out, size_t out_nbytes_avail)
588 {
589         const u8 * const in_begin = in;
590         const u8 *       in_next = in_begin;
591         const u8 * const in_end = in_begin + in_nbytes;
592         struct xpress_item *next_chosen_item = c->chosen_items;
593         unsigned len_3_too_far;
594
595         if (in_nbytes <= 8192)
596                 len_3_too_far = 2048;
597         else
598                 len_3_too_far = 4096;
599
600         hc_matchfinder_init(&c->hc_mf);
601
602         do {
603                 unsigned cur_len;
604                 unsigned cur_offset;
605                 unsigned next_len;
606                 unsigned next_offset;
607
608                 /* Find the longest match at the current position.  */
609                 cur_len = hc_matchfinder_longest_match(&c->hc_mf,
610                                                        in_begin,
611                                                        in_next,
612                                                        XPRESS_MIN_MATCH_LEN - 1,
613                                                        in_end - in_next,
614                                                        min(in_end - in_next, c->nice_match_length),
615                                                        c->max_search_depth,
616                                                        &cur_offset);
617                 in_next += 1;
618
619                 if (cur_len < XPRESS_MIN_MATCH_LEN ||
620                     (cur_len == XPRESS_MIN_MATCH_LEN &&
621                      cur_offset >= len_3_too_far))
622                 {
623                         /* No match found.  Choose a literal.  */
624                         *next_chosen_item++ =
625                                 xpress_record_literal(c, *(in_next - 1));
626                         continue;
627                 }
628
629         have_cur_match:
630                 /* We have a match at the current position.  */
631
632                 /* If the current match is very long, choose it immediately.  */
633                 if (cur_len >= c->nice_match_length) {
634
635                         *next_chosen_item++ =
636                                 xpress_record_match(c, cur_len, cur_offset);
637
638                         hc_matchfinder_skip_positions(&c->hc_mf,
639                                                       in_begin,
640                                                       in_next,
641                                                       in_end,
642                                                       cur_len - 1);
643                         in_next += cur_len - 1;
644                         continue;
645                 }
646
647                 /*
648                  * Try to find a match at the next position.
649                  *
650                  * Note: since we already have a match at the *current*
651                  * position, we use only half the 'max_search_depth' when
652                  * checking the *next* position.  This is a useful trade-off
653                  * because it's more worthwhile to use a greater search depth on
654                  * the initial match than on the next match (since a lot of the
655                  * time, that next match won't even be used).
656                  *
657                  * Note: it's possible to structure the code such that there's
658                  * only one call to longest_match(), which handles both the
659                  * "find the initial match" and "try to find a longer match"
660                  * cases.  However, it is faster to have two call sites, with
661                  * longest_match() inlined at each.
662                  */
663                 next_len = hc_matchfinder_longest_match(&c->hc_mf,
664                                                         in_begin,
665                                                         in_next,
666                                                         cur_len,
667                                                         in_end - in_next,
668                                                         min(in_end - in_next, c->nice_match_length),
669                                                         c->max_search_depth / 2,
670                                                         &next_offset);
671                 in_next += 1;
672
673                 if (next_len > cur_len) {
674                         /* Found a longer match at the next position, so output
675                          * a literal.  */
676                         *next_chosen_item++ =
677                                 xpress_record_literal(c, *(in_next - 2));
678                         cur_len = next_len;
679                         cur_offset = next_offset;
680                         goto have_cur_match;
681                 } else {
682                         /* Didn't find a longer match at the next position, so
683                          * output the current match.  */
684                         *next_chosen_item++ =
685                                 xpress_record_match(c, cur_len, cur_offset);
686                         hc_matchfinder_skip_positions(&c->hc_mf,
687                                                       in_begin,
688                                                       in_next,
689                                                       in_end,
690                                                       cur_len - 2);
691                         in_next += cur_len - 2;
692                         continue;
693                 }
694         } while (in_next != in_end);
695
696         return xpress_write(c, out, out_nbytes_avail,
697                             next_chosen_item - c->chosen_items, false);
698 }
699
700 #if SUPPORT_NEAR_OPTIMAL_PARSING
701
702 /*
703  * Set Huffman symbol costs for the first optimization pass.
704  *
705  * It works well to assume that each Huffman symbol is equally probable.  This
706  * results in each symbol being assigned a cost of -log2(1.0/num_syms) where
707  * 'num_syms' is the number of symbols in the alphabet.
708  */
709 static void
710 xpress_set_default_costs(struct xpress_compressor *c)
711 {
712         for (unsigned i = 0; i < XPRESS_NUM_SYMBOLS; i++)
713                 c->costs[i] = 9;
714 }
715
716 /* Update the cost model based on the codeword lengths @c->lens.  */
717 static void
718 xpress_update_costs(struct xpress_compressor *c)
719 {
720         for (unsigned i = 0; i < XPRESS_NUM_SYMBOLS; i++)
721                 c->costs[i] = c->lens[i] ? c->lens[i] : XPRESS_MAX_CODEWORD_LEN;
722 }
723
724 /*
725  * Follow the minimum cost path in the graph of possible match/literal choices
726  * and compute the frequencies of the Huffman symbols that are needed to output
727  * those matches and literals.
728  */
729 static void
730 xpress_tally_item_list(struct xpress_compressor *c,
731                        struct xpress_optimum_node *end_optimum_ptr)
732 {
733         struct xpress_optimum_node *cur_optimum_ptr = c->optimum_nodes;
734
735         do {
736                 unsigned length = cur_optimum_ptr->item & OPTIMUM_LEN_MASK;
737                 unsigned offset = cur_optimum_ptr->item >> OPTIMUM_OFFSET_SHIFT;
738
739                 if (length == 1) {
740                         /* Literal  */
741                         unsigned literal = offset;
742
743                         c->freqs[literal]++;
744                 } else {
745                         /* Match  */
746                         unsigned adjusted_len;
747                         unsigned offset_high_bit;
748                         unsigned len_hdr;
749                         unsigned sym;
750
751                         adjusted_len = length - XPRESS_MIN_MATCH_LEN;
752                         offset_high_bit = fls32(offset);
753                         len_hdr = min(0xF, adjusted_len);
754                         sym = XPRESS_NUM_CHARS + ((offset_high_bit << 4) | len_hdr);
755
756                         c->freqs[sym]++;
757                 }
758                 cur_optimum_ptr += length;
759         } while (cur_optimum_ptr != end_optimum_ptr);
760 }
761
762 /*
763  * Find a new minimum cost path through the graph of possible match/literal
764  * choices.  We find the minimum cost path from 'c->optimum_nodes[0]', which
765  * represents the node at the beginning of the input buffer, to
766  * 'c->optimum_nodes[in_nbytes]', which represents the node at the end of the
767  * input buffer.  Edge costs are evaluated using the cost model 'c->costs'.
768  *
769  * The algorithm works backward, starting at 'c->optimum_nodes[in_nbytes]' and
770  * proceeding backwards one position at a time.  At each position, the minimum
771  * cost to reach 'c->optimum_nodes[in_nbytes]' from that position is computed
772  * and the match/literal choice is saved.
773  */
774 static void
775 xpress_find_min_cost_path(struct xpress_compressor *c, size_t in_nbytes,
776                           struct lz_match *end_cache_ptr)
777 {
778         struct xpress_optimum_node *cur_optimum_ptr = c->optimum_nodes + in_nbytes;
779         struct lz_match *cache_ptr = end_cache_ptr;
780
781         cur_optimum_ptr->cost_to_end = 0;
782         do {
783                 unsigned literal;
784                 u32 best_item;
785                 u32 best_cost_to_end;
786                 unsigned num_matches;
787                 struct lz_match *match;
788                 unsigned len;
789
790                 cur_optimum_ptr--;
791                 cache_ptr--;
792
793                 literal = cache_ptr->offset;
794
795                 /* Consider coding a literal.  */
796                 best_item = ((u32)literal << OPTIMUM_OFFSET_SHIFT) | 1;
797                 best_cost_to_end = c->costs[literal] +
798                                    (cur_optimum_ptr + 1)->cost_to_end;
799
800                 num_matches = cache_ptr->length;
801
802                 if (num_matches == 0) {
803                         /* No matches; the only choice is the literal.  */
804                         cur_optimum_ptr->cost_to_end = best_cost_to_end;
805                         cur_optimum_ptr->item = best_item;
806                         continue;
807                 }
808
809                 /*
810                  * Consider each match length from the minimum
811                  * (XPRESS_MIN_MATCH_LEN) to the length of the longest match
812                  * found at this position.  For each length, consider only the
813                  * smallest offset for which that length is available.  Although
814                  * this is not guaranteed to be optimal due to the possibility
815                  * of a larger offset costing less than a smaller offset to
816                  * code, this is a very useful heuristic.
817                  */
818                 match = cache_ptr - num_matches;
819                 len = XPRESS_MIN_MATCH_LEN;
820                 if (cache_ptr[-1].length < 0xF + XPRESS_MIN_MATCH_LEN) {
821                         /* All lengths are small.  Optimize accordingly.  */
822                         do {
823                                 unsigned offset;
824                                 unsigned offset_high_bit;
825                                 u32 offset_cost;
826
827                                 offset = match->offset;
828                                 offset_high_bit = fls32(offset);
829                                 offset_cost = offset_high_bit;
830                                 do {
831                                         unsigned len_hdr;
832                                         unsigned sym;
833                                         u32 cost_to_end;
834
835                                         len_hdr = len - XPRESS_MIN_MATCH_LEN;
836                                         sym = XPRESS_NUM_CHARS +
837                                               ((offset_high_bit << 4) | len_hdr);
838                                         cost_to_end =
839                                                 offset_cost + c->costs[sym] +
840                                                 (cur_optimum_ptr + len)->cost_to_end;
841                                         if (cost_to_end < best_cost_to_end) {
842                                                 best_cost_to_end = cost_to_end;
843                                                 best_item =
844                                                         ((u32)offset <<
845                                                          OPTIMUM_OFFSET_SHIFT) | len;
846                                         }
847                                 } while (++len <= match->length);
848                         } while (++match != cache_ptr);
849                 } else {
850                         /* Some lengths are big.  */
851                         do {
852                                 unsigned offset;
853                                 unsigned offset_high_bit;
854                                 u32 offset_cost;
855
856                                 offset = match->offset;
857                                 offset_high_bit = fls32(offset);
858                                 offset_cost = offset_high_bit;
859                                 do {
860                                         unsigned adjusted_len;
861                                         unsigned len_hdr;
862                                         unsigned sym;
863                                         u32 cost_to_end;
864
865                                         adjusted_len = len - XPRESS_MIN_MATCH_LEN;
866                                         len_hdr = min(adjusted_len, 0xF);
867                                         sym = XPRESS_NUM_CHARS +
868                                               ((offset_high_bit << 4) | len_hdr);
869                                         cost_to_end =
870                                                 offset_cost + c->costs[sym] +
871                                                 (cur_optimum_ptr + len)->cost_to_end;
872                                         if (adjusted_len >= 0xF) {
873                                                 cost_to_end += 8;
874                                                 if (adjusted_len - 0xF >= 0xFF)
875                                                         cost_to_end += 16;
876                                         }
877                                         if (cost_to_end < best_cost_to_end) {
878                                                 best_cost_to_end = cost_to_end;
879                                                 best_item =
880                                                         ((u32)offset <<
881                                                          OPTIMUM_OFFSET_SHIFT) | len;
882                                         }
883                                 } while (++len <= match->length);
884                         } while (++match != cache_ptr);
885                 }
886                 cache_ptr -= num_matches;
887                 cur_optimum_ptr->cost_to_end = best_cost_to_end;
888                 cur_optimum_ptr->item = best_item;
889         } while (cur_optimum_ptr != c->optimum_nodes);
890 }
891
892 /*
893  * This routine finds matches at each position in the buffer in[0...in_nbytes].
894  * The matches are cached in the array c->match_cache, and the return value is a
895  * pointer past the last slot in this array that was filled.
896  */
897 static struct lz_match *
898 xpress_find_matches(struct xpress_compressor * restrict c,
899                     const void * restrict in, size_t in_nbytes)
900 {
901         const u8 * const in_begin = in;
902         const u8 *in_next = in_begin;
903         const u8 * const in_end = in_begin + in_nbytes;
904         struct lz_match *cache_ptr = c->match_cache;
905         u32 next_hash;
906
907         bt_matchfinder_init(&c->bt_mf);
908         next_hash = bt_matchfinder_hash_3_bytes(in_next);
909
910         do {
911                 struct lz_match *matches;
912                 unsigned best_len;
913
914                 /* If we've found so many matches that the cache might overflow
915                  * if we keep finding more, then stop finding matches.  This
916                  * case is very unlikely.  */
917                 if (unlikely(cache_ptr >= c->cache_overflow_mark)) {
918                         do {
919                                 cache_ptr->length = 0;
920                                 cache_ptr->offset = *in_next++;
921                                 cache_ptr++;
922                         } while (in_next != in_end);
923                         return cache_ptr;
924                 }
925
926                 matches = cache_ptr;
927
928                 /* Find matches with the current position using the binary tree
929                  * matchfinder and save them in the next available slots in
930                  * the match cache.  */
931                 cache_ptr =
932                         bt_matchfinder_get_matches(&c->bt_mf,
933                                                    in_begin,
934                                                    in_next,
935                                                    XPRESS_MIN_MATCH_LEN,
936                                                    in_end - in_next,
937                                                    min(in_end - in_next, c->nice_match_length),
938                                                    c->max_search_depth,
939                                                    &next_hash,
940                                                    &best_len,
941                                                    cache_ptr);
942                 cache_ptr->length = cache_ptr - matches;
943                 cache_ptr->offset = *in_next;
944                 in_next++;
945                 cache_ptr++;
946
947                 /*
948                  * If there was a very long match found, then don't cache any
949                  * matches for the bytes covered by that match.  This avoids
950                  * degenerate behavior when compressing highly redundant data,
951                  * where the number of matches can be very large.
952                  *
953                  * This heuristic doesn't actually hurt the compression ratio
954                  * very much.  If there's a long match, then the data must be
955                  * highly compressible, so it doesn't matter as much what we do.
956                  */
957                 if (best_len >= c->nice_match_length) {
958                         --best_len;
959                         do {
960                                 bt_matchfinder_skip_position(&c->bt_mf,
961                                                              in_begin,
962                                                              in_next,
963                                                              in_end,
964                                                              min(in_end - in_next,
965                                                                  c->nice_match_length),
966                                                              c->max_search_depth,
967                                                              &next_hash);
968
969                                 cache_ptr->length = 0;
970                                 cache_ptr->offset = *in_next++;
971                                 cache_ptr++;
972                         } while (--best_len);
973                 }
974         } while (in_next != in_end);
975
976         return cache_ptr;
977 }
978
979 /*
980  * This is the "near-optimal" XPRESS compressor.  It computes a compressed
981  * representation of the input buffer by executing a minimum cost path search
982  * over the graph of possible match/literal choices, assuming a certain cost for
983  * each Huffman symbol.  The result is usually close to optimal, but it is *not*
984  * guaranteed to be optimal because of (a) heuristic restrictions in which
985  * matches are considered, and (b) symbol costs are unknown until those symbols
986  * have already been chosen --- so iterative optimization must be used, and the
987  * algorithm might converge on a local optimum rather than a global optimum.
988  */
989 static size_t
990 xpress_compress_near_optimal(struct xpress_compressor * restrict c,
991                              const void * restrict in, size_t in_nbytes,
992                              void * restrict out, size_t out_nbytes_avail)
993 {
994         struct lz_match *end_cache_ptr;
995         unsigned num_passes_remaining = c->num_optim_passes;
996
997         /* Run the input buffer through the matchfinder and save the results. */
998         end_cache_ptr = xpress_find_matches(c, in, in_nbytes);
999
1000         /* The first optimization pass uses a default cost model.  Each
1001          * additional optimization pass uses a cost model derived from the
1002          * Huffman code computed in the previous pass.  */
1003         xpress_set_default_costs(c);
1004         do {
1005                 xpress_find_min_cost_path(c, in_nbytes, end_cache_ptr);
1006                 xpress_tally_item_list(c, c->optimum_nodes + in_nbytes);
1007                 if (num_passes_remaining > 1) {
1008                         c->freqs[XPRESS_END_OF_DATA]++;
1009                         xpress_make_huffman_code(c);
1010                         xpress_update_costs(c);
1011                         xpress_reset_symbol_frequencies(c);
1012                 }
1013         } while (--num_passes_remaining);
1014
1015         return xpress_write(c, out, out_nbytes_avail, in_nbytes, true);
1016 }
1017
1018 #endif /* SUPPORT_NEAR_OPTIMAL_PARSING */
1019
1020 static size_t
1021 xpress_get_compressor_size(size_t max_bufsize, unsigned compression_level)
1022 {
1023 #if SUPPORT_NEAR_OPTIMAL_PARSING
1024         if (compression_level >= MIN_LEVEL_FOR_NEAR_OPTIMAL)
1025                 return offsetof(struct xpress_compressor, bt_mf) +
1026                         bt_matchfinder_size(max_bufsize);
1027 #endif
1028
1029         return offsetof(struct xpress_compressor, hc_mf) +
1030                 hc_matchfinder_size(max_bufsize);
1031 }
1032
1033 static u64
1034 xpress_get_needed_memory(size_t max_bufsize, unsigned compression_level)
1035 {
1036         u64 size = 0;
1037
1038         if (max_bufsize > XPRESS_MAX_BUFSIZE)
1039                 return 0;
1040
1041         size += xpress_get_compressor_size(max_bufsize, compression_level);
1042
1043         if (compression_level < MIN_LEVEL_FOR_NEAR_OPTIMAL ||
1044             !SUPPORT_NEAR_OPTIMAL_PARSING) {
1045                 /* chosen_items  */
1046                 size += max_bufsize * sizeof(struct xpress_item);
1047         }
1048 #if SUPPORT_NEAR_OPTIMAL_PARSING
1049         else {
1050                 /* optimum_nodes  */
1051                 size += (max_bufsize + 1) * sizeof(struct xpress_optimum_node);
1052                 /* match_cache */
1053                 size += ((max_bufsize * CACHE_RESERVE_PER_POS) +
1054                          XPRESS_MAX_MATCH_LEN + max_bufsize) *
1055                                 sizeof(struct lz_match);
1056         }
1057 #endif
1058         return size;
1059 }
1060
1061 static int
1062 xpress_create_compressor(size_t max_bufsize, unsigned compression_level,
1063                          void **c_ret)
1064 {
1065         struct xpress_compressor *c;
1066
1067         if (max_bufsize > XPRESS_MAX_BUFSIZE)
1068                 return WIMLIB_ERR_INVALID_PARAM;
1069
1070         c = ALIGNED_MALLOC(xpress_get_compressor_size(max_bufsize, compression_level),
1071                            MATCHFINDER_ALIGNMENT);
1072         if (!c)
1073                 goto oom0;
1074
1075         if (compression_level < MIN_LEVEL_FOR_NEAR_OPTIMAL ||
1076             !SUPPORT_NEAR_OPTIMAL_PARSING)
1077         {
1078
1079                 c->chosen_items = MALLOC(max_bufsize * sizeof(struct xpress_item));
1080                 if (!c->chosen_items)
1081                         goto oom1;
1082
1083                 if (compression_level < 30) {
1084                         c->impl = xpress_compress_greedy;
1085                         c->max_search_depth = (compression_level * 24) / 16;
1086                         c->nice_match_length = (compression_level * 48) / 16;
1087                 } else {
1088                         c->impl = xpress_compress_lazy;
1089                         c->max_search_depth = (compression_level * 24) / 32;
1090                         c->nice_match_length = (compression_level * 48) / 32;
1091                 }
1092         }
1093 #if SUPPORT_NEAR_OPTIMAL_PARSING
1094         else {
1095
1096                 c->optimum_nodes = MALLOC((max_bufsize + 1) *
1097                                           sizeof(struct xpress_optimum_node));
1098                 c->match_cache = MALLOC(((max_bufsize * CACHE_RESERVE_PER_POS) +
1099                                          XPRESS_MAX_MATCH_LEN + max_bufsize) *
1100                                         sizeof(struct lz_match));
1101                 if (!c->optimum_nodes || !c->match_cache) {
1102                         FREE(c->optimum_nodes);
1103                         FREE(c->match_cache);
1104                         goto oom1;
1105                 }
1106                 c->cache_overflow_mark =
1107                         &c->match_cache[max_bufsize * CACHE_RESERVE_PER_POS];
1108
1109                 c->impl = xpress_compress_near_optimal;
1110                 c->max_search_depth = (compression_level * 32) / 100;
1111                 c->nice_match_length = (compression_level * 50) / 100;
1112                 c->num_optim_passes = compression_level / 40;
1113         }
1114 #endif /* SUPPORT_NEAR_OPTIMAL_PARSING */
1115
1116         *c_ret = c;
1117         return 0;
1118
1119 oom1:
1120         ALIGNED_FREE(c);
1121 oom0:
1122         return WIMLIB_ERR_NOMEM;
1123 }
1124
1125 static size_t
1126 xpress_compress(const void *in, size_t in_nbytes,
1127                 void *out, size_t out_nbytes_avail, void *_c)
1128 {
1129         struct xpress_compressor *c = _c;
1130
1131         /* Don't bother trying to compress very small inputs.  */
1132         if (in_nbytes < 25)
1133                 return 0;
1134
1135         if (out_nbytes_avail <= XPRESS_NUM_SYMBOLS / 2 + 4)
1136                 return 0;
1137
1138         xpress_reset_symbol_frequencies(c);
1139
1140         return (*c->impl)(c, in, in_nbytes, out, out_nbytes_avail);
1141 }
1142
1143 static void
1144 xpress_free_compressor(void *_c)
1145 {
1146         struct xpress_compressor *c = _c;
1147
1148 #if SUPPORT_NEAR_OPTIMAL_PARSING
1149         if (c->impl == xpress_compress_near_optimal) {
1150                 FREE(c->optimum_nodes);
1151                 FREE(c->match_cache);
1152         } else
1153 #endif
1154                 FREE(c->chosen_items);
1155         ALIGNED_FREE(c);
1156 }
1157
1158 const struct compressor_ops xpress_compressor_ops = {
1159         .get_needed_memory  = xpress_get_needed_memory,
1160         .create_compressor  = xpress_create_compressor,
1161         .compress           = xpress_compress,
1162         .free_compressor    = xpress_free_compressor,
1163 };