]> wimlib.net Git - wimlib/blob - src/lzms_compress.c
Adjust naming of (de)compression files
[wimlib] / src / lzms_compress.c
1 /*
2  * lzms_compress.c
3  *
4  * A compressor for the LZMS compression format.
5  */
6
7 /*
8  * Copyright (C) 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 #include "wimlib/compress_common.h"
29 #include "wimlib/compressor_ops.h"
30 #include "wimlib/endianness.h"
31 #include "wimlib/error.h"
32 #include "wimlib/lz_mf.h"
33 #include "wimlib/lz_repsearch.h"
34 #include "wimlib/lzms_common.h"
35 #include "wimlib/unaligned.h"
36 #include "wimlib/util.h"
37
38 #include <string.h>
39 #include <limits.h>
40 #include <pthread.h>
41
42 /* Stucture used for writing raw bits as a series of 16-bit little endian coding
43  * units.  This starts at the *end* of the compressed data buffer and proceeds
44  * backwards.  */
45 struct lzms_output_bitstream {
46
47         /* Bits that haven't yet been written to the output buffer.  */
48         u64 bitbuf;
49
50         /* Number of bits currently held in @bitbuf.  */
51         unsigned bitcount;
52
53         /* Pointer to one past the next position in the compressed data buffer
54          * at which to output a 16-bit coding unit.  */
55         le16 *next;
56
57         /* Pointer to the beginning of the output buffer.  (The "end" when
58          * writing backwards!)  */
59         le16 *begin;
60 };
61
62 /* Stucture used for range encoding (raw version).  This starts at the
63  * *beginning* of the compressed data buffer and proceeds forward.  */
64 struct lzms_range_encoder_raw {
65
66         /* A 33-bit variable that holds the low boundary of the current range.
67          * The 33rd bit is needed to catch carries.  */
68         u64 low;
69
70         /* Size of the current range.  */
71         u32 range;
72
73         /* Next 16-bit coding unit to output.  */
74         u16 cache;
75
76         /* Number of 16-bit coding units whose output has been delayed due to
77          * possible carrying.  The first such coding unit is @cache; all
78          * subsequent such coding units are 0xffff.  */
79         u32 cache_size;
80
81         /* Pointer to the beginning of the output buffer.  */
82         le16 *begin;
83
84         /* Pointer to the position in the output buffer at which the next coding
85          * unit must be written.  */
86         le16 *next;
87
88         /* Pointer just past the end of the output buffer.  */
89         le16 *end;
90 };
91
92 /* Structure used for range encoding.  This wraps around `struct
93  * lzms_range_encoder_raw' to use and maintain probability entries.  */
94 struct lzms_range_encoder {
95
96         /* Pointer to the raw range encoder, which has no persistent knowledge
97          * of probabilities.  Multiple lzms_range_encoder's share the same
98          * lzms_range_encoder_raw.  */
99         struct lzms_range_encoder_raw *rc;
100
101         /* Bits recently encoded by this range encoder.  This is used as an
102          * index into @prob_entries.  */
103         u32 state;
104
105         /* Bitmask for @state to prevent its value from exceeding the number of
106          * probability entries.  */
107         u32 mask;
108
109         /* Probability entries being used for this range encoder.  */
110         struct lzms_probability_entry prob_entries[LZMS_MAX_NUM_STATES];
111 };
112
113 /* Structure used for Huffman encoding.  */
114 struct lzms_huffman_encoder {
115
116         /* Bitstream to write Huffman-encoded symbols and verbatim bits to.
117          * Multiple lzms_huffman_encoder's share the same lzms_output_bitstream.
118          */
119         struct lzms_output_bitstream *os;
120
121         /* Number of symbols that have been written using this code far.  Reset
122          * to 0 whenever the code is rebuilt.  */
123         u32 num_syms_written;
124
125         /* When @num_syms_written reaches this number, the Huffman code must be
126          * rebuilt.  */
127         u32 rebuild_freq;
128
129         /* Number of symbols in the represented Huffman code.  */
130         unsigned num_syms;
131
132         /* Running totals of symbol frequencies.  These are diluted slightly
133          * whenever the code is rebuilt.  */
134         u32 sym_freqs[LZMS_MAX_NUM_SYMS];
135
136         /* The length, in bits, of each symbol in the Huffman code.  */
137         u8 lens[LZMS_MAX_NUM_SYMS];
138
139         /* The codeword of each symbol in the Huffman code.  */
140         u32 codewords[LZMS_MAX_NUM_SYMS];
141 };
142
143 /* Internal compression parameters  */
144 struct lzms_compressor_params {
145         u32 min_match_length;
146         u32 nice_match_length;
147         u32 max_search_depth;
148         u32 optim_array_length;
149 };
150
151 /* State of the LZMS compressor  */
152 struct lzms_compressor {
153
154         /* Internal compression parameters  */
155         struct lzms_compressor_params params;
156
157         /* Data currently being compressed  */
158         u8 *cur_window;
159         u32 cur_window_size;
160
161         /* Lempel-Ziv match-finder  */
162         struct lz_mf *mf;
163
164         /* Temporary space to store found matches  */
165         struct lz_match *matches;
166
167         /* Per-position data for near-optimal parsing  */
168         struct lzms_mc_pos_data *optimum;
169         struct lzms_mc_pos_data *optimum_end;
170
171         /* Raw range encoder which outputs to the beginning of the compressed
172          * data buffer, proceeding forwards  */
173         struct lzms_range_encoder_raw rc;
174
175         /* Bitstream which outputs to the end of the compressed data buffer,
176          * proceeding backwards  */
177         struct lzms_output_bitstream os;
178
179         /* Range encoders  */
180         struct lzms_range_encoder main_range_encoder;
181         struct lzms_range_encoder match_range_encoder;
182         struct lzms_range_encoder lz_match_range_encoder;
183         struct lzms_range_encoder lz_repeat_match_range_encoders[LZMS_NUM_RECENT_OFFSETS - 1];
184         struct lzms_range_encoder delta_match_range_encoder;
185         struct lzms_range_encoder delta_repeat_match_range_encoders[LZMS_NUM_RECENT_OFFSETS - 1];
186
187         /* Huffman encoders  */
188         struct lzms_huffman_encoder literal_encoder;
189         struct lzms_huffman_encoder lz_offset_encoder;
190         struct lzms_huffman_encoder length_encoder;
191         struct lzms_huffman_encoder delta_power_encoder;
192         struct lzms_huffman_encoder delta_offset_encoder;
193
194         /* Used for preprocessing  */
195         s32 last_target_usages[65536];
196
197 #define LZMS_NUM_FAST_LENGTHS 256
198         /* Table: length => length slot for small lengths  */
199         u8 length_slot_fast[LZMS_NUM_FAST_LENGTHS];
200
201         /* Table: length => current cost for small match lengths  */
202         u32 length_cost_fast[LZMS_NUM_FAST_LENGTHS];
203
204 #define LZMS_NUM_FAST_OFFSETS 32768
205         /* Table: offset => offset slot for small offsets  */
206         u8 offset_slot_fast[LZMS_NUM_FAST_OFFSETS];
207 };
208
209 struct lzms_lz_lru_queue {
210         u32 recent_offsets[LZMS_NUM_RECENT_OFFSETS + 1];
211         u32 prev_offset;
212         u32 upcoming_offset;
213 };
214
215 static void
216 lzms_init_lz_lru_queue(struct lzms_lz_lru_queue *queue)
217 {
218         for (int i = 0; i < LZMS_NUM_RECENT_OFFSETS + 1; i++)
219                 queue->recent_offsets[i] = i + 1;
220
221         queue->prev_offset = 0;
222         queue->upcoming_offset = 0;
223 }
224
225 static void
226 lzms_update_lz_lru_queue(struct lzms_lz_lru_queue *queue)
227 {
228         if (queue->prev_offset != 0) {
229                 for (int i = LZMS_NUM_RECENT_OFFSETS - 1; i >= 0; i--)
230                         queue->recent_offsets[i + 1] = queue->recent_offsets[i];
231                 queue->recent_offsets[0] = queue->prev_offset;
232         }
233         queue->prev_offset = queue->upcoming_offset;
234 }
235
236 /*
237  * Match chooser position data:
238  *
239  * An array of these structures is used during the near-optimal match-choosing
240  * algorithm.  They correspond to consecutive positions in the window and are
241  * used to keep track of the cost to reach each position, and the match/literal
242  * choices that need to be chosen to reach that position.
243  */
244 struct lzms_mc_pos_data {
245
246         /* The cost, in bits, of the lowest-cost path that has been found to
247          * reach this position.  This can change as progressively lower cost
248          * paths are found to reach this position.  */
249         u32 cost;
250 #define MC_INFINITE_COST UINT32_MAX
251
252         /* The match or literal that was taken to reach this position.  This can
253          * change as progressively lower cost paths are found to reach this
254          * position.
255          *
256          * This variable is divided into two bitfields.
257          *
258          * Literals:
259          *      Low bits are 1, high bits are the literal.
260          *
261          * Explicit offset matches:
262          *      Low bits are the match length, high bits are the offset plus 2.
263          *
264          * Repeat offset matches:
265          *      Low bits are the match length, high bits are the queue index.
266          */
267         u64 mc_item_data;
268 #define MC_OFFSET_SHIFT 32
269 #define MC_LEN_MASK (((u64)1 << MC_OFFSET_SHIFT) - 1)
270
271         /* The LZMS adaptive state that exists at this position.  This is filled
272          * in lazily, only after the minimum-cost path to this position is
273          * found.
274          *
275          * Note: the way we handle this adaptive state in the "minimum-cost"
276          * parse is actually only an approximation.  It's possible for the
277          * globally optimal, minimum cost path to contain a prefix, ending at a
278          * position, where that path prefix is *not* the minimum cost path to
279          * that position.  This can happen if such a path prefix results in a
280          * different adaptive state which results in lower costs later.  We do
281          * not solve this problem; we only consider the lowest cost to reach
282          * each position, which seems to be an acceptable approximation.
283          *
284          * Note: this adaptive state also does not include the probability
285          * entries or current Huffman codewords.  Those aren't maintained
286          * per-position and are only updated occassionally.  */
287         struct lzms_adaptive_state {
288                 struct lzms_lz_lru_queue lru;
289                 u8 main_state;
290                 u8 match_state;
291                 u8 lz_match_state;
292                 u8 lz_repeat_match_state[LZMS_NUM_RECENT_OFFSETS - 1];
293         } state;
294 };
295
296 static void
297 lzms_init_fast_slots(struct lzms_compressor *c)
298 {
299         /* Create table mapping small lengths to length slots.  */
300         for (unsigned slot = 0, i = 0; i < LZMS_NUM_FAST_LENGTHS; i++) {
301                 while (i >= lzms_length_slot_base[slot + 1])
302                         slot++;
303                 c->length_slot_fast[i] = slot;
304         }
305
306         /* Create table mapping small offsets to offset slots.  */
307         for (unsigned slot = 0, i = 0; i < LZMS_NUM_FAST_OFFSETS; i++) {
308                 while (i >= lzms_offset_slot_base[slot + 1])
309                         slot++;
310                 c->offset_slot_fast[i] = slot;
311         }
312 }
313
314 static inline unsigned
315 lzms_get_length_slot_fast(const struct lzms_compressor *c, u32 length)
316 {
317         if (likely(length < LZMS_NUM_FAST_LENGTHS))
318                 return c->length_slot_fast[length];
319         else
320                 return lzms_get_length_slot(length);
321 }
322
323 static inline unsigned
324 lzms_get_offset_slot_fast(const struct lzms_compressor *c, u32 offset)
325 {
326         if (offset < LZMS_NUM_FAST_OFFSETS)
327                 return c->offset_slot_fast[offset];
328         else
329                 return lzms_get_offset_slot(offset);
330 }
331
332 /* Initialize the output bitstream @os to write backwards to the specified
333  * compressed data buffer @out that is @out_limit 16-bit integers long.  */
334 static void
335 lzms_output_bitstream_init(struct lzms_output_bitstream *os,
336                            le16 *out, size_t out_limit)
337 {
338         os->bitbuf = 0;
339         os->bitcount = 0;
340         os->next = out + out_limit;
341         os->begin = out;
342 }
343
344 /*
345  * Write some bits, contained in the low @num_bits bits of @bits (ordered from
346  * high-order to low-order), to the output bitstream @os.
347  *
348  * @max_num_bits is a compile-time constant that specifies the maximum number of
349  * bits that can ever be written at this call site.
350  */
351 static inline void
352 lzms_output_bitstream_put_varbits(struct lzms_output_bitstream *os,
353                                   u32 bits, unsigned num_bits,
354                                   unsigned max_num_bits)
355 {
356         LZMS_ASSERT(num_bits <= 48);
357
358         /* Add the bits to the bit buffer variable.  */
359         os->bitcount += num_bits;
360         os->bitbuf = (os->bitbuf << num_bits) | bits;
361
362         /* Check whether any coding units need to be written.  */
363         while (os->bitcount >= 16) {
364
365                 os->bitcount -= 16;
366
367                 /* Write a coding unit, unless it would underflow the buffer. */
368                 if (os->next != os->begin)
369                         put_unaligned_u16_le(os->bitbuf >> os->bitcount, --os->next);
370
371                 /* Optimization for call sites that never write more than 16
372                  * bits at once.  */
373                 if (max_num_bits <= 16)
374                         break;
375         }
376 }
377
378 /* Flush the output bitstream, ensuring that all bits written to it have been
379  * written to memory.  Returns %true if all bits have been output successfully,
380  * or %false if an overrun occurred.  */
381 static bool
382 lzms_output_bitstream_flush(struct lzms_output_bitstream *os)
383 {
384         if (os->next == os->begin)
385                 return false;
386
387         if (os->bitcount != 0)
388                 put_unaligned_u16_le(os->bitbuf << (16 - os->bitcount), --os->next);
389
390         return true;
391 }
392
393 /* Initialize the range encoder @rc to write forwards to the specified
394  * compressed data buffer @out that is @out_limit 16-bit integers long.  */
395 static void
396 lzms_range_encoder_raw_init(struct lzms_range_encoder_raw *rc,
397                             le16 *out, size_t out_limit)
398 {
399         rc->low = 0;
400         rc->range = 0xffffffff;
401         rc->cache = 0;
402         rc->cache_size = 1;
403         rc->begin = out;
404         rc->next = out - 1;
405         rc->end = out + out_limit;
406 }
407
408 /*
409  * Attempt to flush bits from the range encoder.
410  *
411  * Note: this is based on the public domain code for LZMA written by Igor
412  * Pavlov.  The only differences in this function are that in LZMS the bits must
413  * be output in 16-bit coding units instead of 8-bit coding units, and that in
414  * LZMS the first coding unit is not ignored by the decompressor, so the encoder
415  * cannot output a dummy value to that position.
416  *
417  * The basic idea is that we're writing bits from @rc->low to the output.
418  * However, due to carrying, the writing of coding units with value 0xffff, as
419  * well as one prior coding unit, must be delayed until it is determined whether
420  * a carry is needed.
421  */
422 static void
423 lzms_range_encoder_raw_shift_low(struct lzms_range_encoder_raw *rc)
424 {
425         if ((u32)(rc->low) < 0xffff0000 ||
426             (u32)(rc->low >> 32) != 0)
427         {
428                 /* Carry not needed (rc->low < 0xffff0000), or carry occurred
429                  * ((rc->low >> 32) != 0, a.k.a. the carry bit is 1).  */
430                 do {
431                         if (likely(rc->next >= rc->begin)) {
432                                 if (rc->next != rc->end) {
433                                         put_unaligned_u16_le(rc->cache +
434                                                              (u16)(rc->low >> 32),
435                                                              rc->next++);
436                                 }
437                         } else {
438                                 rc->next++;
439                         }
440                         rc->cache = 0xffff;
441                 } while (--rc->cache_size != 0);
442
443                 rc->cache = (rc->low >> 16) & 0xffff;
444         }
445         ++rc->cache_size;
446         rc->low = (rc->low & 0xffff) << 16;
447 }
448
449 static void
450 lzms_range_encoder_raw_normalize(struct lzms_range_encoder_raw *rc)
451 {
452         if (rc->range <= 0xffff) {
453                 rc->range <<= 16;
454                 lzms_range_encoder_raw_shift_low(rc);
455         }
456 }
457
458 static bool
459 lzms_range_encoder_raw_flush(struct lzms_range_encoder_raw *rc)
460 {
461         for (unsigned i = 0; i < 4; i++)
462                 lzms_range_encoder_raw_shift_low(rc);
463         return rc->next != rc->end;
464 }
465
466 /* Encode the next bit using the range encoder (raw version).
467  *
468  * @prob is the chance out of LZMS_PROBABILITY_MAX that the next bit is 0.  */
469 static inline void
470 lzms_range_encoder_raw_encode_bit(struct lzms_range_encoder_raw *rc,
471                                   int bit, u32 prob)
472 {
473         lzms_range_encoder_raw_normalize(rc);
474
475         u32 bound = (rc->range >> LZMS_PROBABILITY_BITS) * prob;
476         if (bit == 0) {
477                 rc->range = bound;
478         } else {
479                 rc->low += bound;
480                 rc->range -= bound;
481         }
482 }
483
484 /* Encode a bit using the specified range encoder. This wraps around
485  * lzms_range_encoder_raw_encode_bit() to handle using and updating the
486  * appropriate state and probability entry.  */
487 static void
488 lzms_range_encode_bit(struct lzms_range_encoder *enc, int bit)
489 {
490         struct lzms_probability_entry *prob_entry;
491         u32 prob;
492
493         /* Load the probability entry corresponding to the current state.  */
494         prob_entry = &enc->prob_entries[enc->state];
495
496         /* Update the state based on the next bit.  */
497         enc->state = ((enc->state << 1) | bit) & enc->mask;
498
499         /* Get the probability that the bit is 0.  */
500         prob = lzms_get_probability(prob_entry);
501
502         /* Update the probability entry.  */
503         lzms_update_probability_entry(prob_entry, bit);
504
505         /* Encode the bit.  */
506         lzms_range_encoder_raw_encode_bit(enc->rc, bit, prob);
507 }
508
509 /* Called when an adaptive Huffman code needs to be rebuilt.  */
510 static void
511 lzms_rebuild_huffman_code(struct lzms_huffman_encoder *enc)
512 {
513         make_canonical_huffman_code(enc->num_syms,
514                                     LZMS_MAX_CODEWORD_LEN,
515                                     enc->sym_freqs,
516                                     enc->lens,
517                                     enc->codewords);
518
519         /* Dilute the frequencies.  */
520         for (unsigned i = 0; i < enc->num_syms; i++) {
521                 enc->sym_freqs[i] >>= 1;
522                 enc->sym_freqs[i] += 1;
523         }
524         enc->num_syms_written = 0;
525 }
526
527 /* Encode a symbol using the specified Huffman encoder.  */
528 static inline void
529 lzms_huffman_encode_symbol(struct lzms_huffman_encoder *enc, unsigned sym)
530 {
531         lzms_output_bitstream_put_varbits(enc->os,
532                                           enc->codewords[sym],
533                                           enc->lens[sym],
534                                           LZMS_MAX_CODEWORD_LEN);
535         ++enc->sym_freqs[sym];
536         if (++enc->num_syms_written == enc->rebuild_freq)
537                 lzms_rebuild_huffman_code(enc);
538 }
539
540 static void
541 lzms_update_fast_length_costs(struct lzms_compressor *c);
542
543 /* Encode a match length.  */
544 static void
545 lzms_encode_length(struct lzms_compressor *c, u32 length)
546 {
547         unsigned slot;
548         unsigned num_extra_bits;
549         u32 extra_bits;
550
551         slot = lzms_get_length_slot_fast(c, length);
552
553         extra_bits = length - lzms_length_slot_base[slot];
554         num_extra_bits = lzms_extra_length_bits[slot];
555
556         lzms_huffman_encode_symbol(&c->length_encoder, slot);
557         if (c->length_encoder.num_syms_written == 0)
558                 lzms_update_fast_length_costs(c);
559
560         lzms_output_bitstream_put_varbits(c->length_encoder.os,
561                                           extra_bits, num_extra_bits, 30);
562 }
563
564 /* Encode an LZ match offset.  */
565 static void
566 lzms_encode_lz_offset(struct lzms_compressor *c, u32 offset)
567 {
568         unsigned slot;
569         unsigned num_extra_bits;
570         u32 extra_bits;
571
572         slot = lzms_get_offset_slot_fast(c, offset);
573
574         extra_bits = offset - lzms_offset_slot_base[slot];
575         num_extra_bits = lzms_extra_offset_bits[slot];
576
577         lzms_huffman_encode_symbol(&c->lz_offset_encoder, slot);
578         lzms_output_bitstream_put_varbits(c->lz_offset_encoder.os,
579                                           extra_bits, num_extra_bits, 30);
580 }
581
582 /* Encode a literal byte.  */
583 static void
584 lzms_encode_literal(struct lzms_compressor *c, unsigned literal)
585 {
586         /* Main bit: 0 = a literal, not a match.  */
587         lzms_range_encode_bit(&c->main_range_encoder, 0);
588
589         /* Encode the literal using the current literal Huffman code.  */
590         lzms_huffman_encode_symbol(&c->literal_encoder, literal);
591 }
592
593 /* Encode an LZ repeat offset match.  */
594 static void
595 lzms_encode_lz_repeat_offset_match(struct lzms_compressor *c,
596                                    u32 length, unsigned rep_index)
597 {
598         unsigned i;
599
600         /* Main bit: 1 = a match, not a literal.  */
601         lzms_range_encode_bit(&c->main_range_encoder, 1);
602
603         /* Match bit: 0 = an LZ match, not a delta match.  */
604         lzms_range_encode_bit(&c->match_range_encoder, 0);
605
606         /* LZ match bit: 1 = repeat offset, not an explicit offset.  */
607         lzms_range_encode_bit(&c->lz_match_range_encoder, 1);
608
609         /* Encode the repeat offset index.  A 1 bit is encoded for each index
610          * passed up.  This sequence of 1 bits is terminated by a 0 bit, or
611          * automatically when (LZMS_NUM_RECENT_OFFSETS - 1) 1 bits have been
612          * encoded.  */
613         for (i = 0; i < rep_index; i++)
614                 lzms_range_encode_bit(&c->lz_repeat_match_range_encoders[i], 1);
615
616         if (i < LZMS_NUM_RECENT_OFFSETS - 1)
617                 lzms_range_encode_bit(&c->lz_repeat_match_range_encoders[i], 0);
618
619         /* Encode the match length.  */
620         lzms_encode_length(c, length);
621 }
622
623 /* Encode an LZ explicit offset match.  */
624 static void
625 lzms_encode_lz_explicit_offset_match(struct lzms_compressor *c,
626                                      u32 length, u32 offset)
627 {
628         /* Main bit: 1 = a match, not a literal.  */
629         lzms_range_encode_bit(&c->main_range_encoder, 1);
630
631         /* Match bit: 0 = an LZ match, not a delta match.  */
632         lzms_range_encode_bit(&c->match_range_encoder, 0);
633
634         /* LZ match bit: 0 = explicit offset, not a repeat offset.  */
635         lzms_range_encode_bit(&c->lz_match_range_encoder, 0);
636
637         /* Encode the match offset.  */
638         lzms_encode_lz_offset(c, offset);
639
640         /* Encode the match length.  */
641         lzms_encode_length(c, length);
642 }
643
644 static void
645 lzms_encode_item(struct lzms_compressor *c, u64 mc_item_data)
646 {
647         u32 len = mc_item_data & MC_LEN_MASK;
648         u32 offset_data = mc_item_data >> MC_OFFSET_SHIFT;
649
650         if (len == 1)
651                 lzms_encode_literal(c, offset_data);
652         else if (offset_data < LZMS_NUM_RECENT_OFFSETS)
653                 lzms_encode_lz_repeat_offset_match(c, len, offset_data);
654         else
655                 lzms_encode_lz_explicit_offset_match(c, len, offset_data - LZMS_OFFSET_OFFSET);
656 }
657
658 /* Encode a list of matches and literals chosen by the parsing algorithm.  */
659 static void
660 lzms_encode_item_list(struct lzms_compressor *c,
661                       struct lzms_mc_pos_data *cur_optimum_ptr)
662 {
663         struct lzms_mc_pos_data *end_optimum_ptr;
664         u64 saved_item;
665         u64 item;
666
667         /* The list is currently in reverse order (last item to first item).
668          * Reverse it.  */
669         end_optimum_ptr = cur_optimum_ptr;
670         saved_item = cur_optimum_ptr->mc_item_data;
671         do {
672                 item = saved_item;
673                 cur_optimum_ptr -= item & MC_LEN_MASK;
674                 saved_item = cur_optimum_ptr->mc_item_data;
675                 cur_optimum_ptr->mc_item_data = item;
676         } while (cur_optimum_ptr != c->optimum);
677
678         /* Walk the list of items from beginning to end, encoding each item.  */
679         do {
680                 lzms_encode_item(c, cur_optimum_ptr->mc_item_data);
681                 cur_optimum_ptr += (cur_optimum_ptr->mc_item_data) & MC_LEN_MASK;
682         } while (cur_optimum_ptr != end_optimum_ptr);
683 }
684
685 /* Each bit costs 1 << LZMS_COST_SHIFT units.  */
686 #define LZMS_COST_SHIFT 6
687
688 /*#define LZMS_RC_COSTS_USE_FLOATING_POINT*/
689
690 static u32
691 lzms_rc_costs[LZMS_PROBABILITY_MAX + 1];
692
693 #ifdef LZMS_RC_COSTS_USE_FLOATING_POINT
694 #  include <math.h>
695 #endif
696
697 static void
698 lzms_do_init_rc_costs(void)
699 {
700         /* Fill in a table that maps range coding probabilities needed to code a
701          * bit X (0 or 1) to the number of bits (scaled by a constant factor, to
702          * handle fractional costs) needed to code that bit X.
703          *
704          * Consider the range of the range decoder.  To eliminate exactly half
705          * the range (logical probability of 0.5), we need exactly 1 bit.  For
706          * lower probabilities we need more bits and for higher probabilities we
707          * need fewer bits.  In general, a logical probability of N will
708          * eliminate the proportion 1 - N of the range; this information takes
709          * log2(1 / N) bits to encode.
710          *
711          * The below loop is simply calculating this number of bits for each
712          * possible probability allowed by the LZMS compression format, but
713          * without using real numbers.  To handle fractional probabilities, each
714          * cost is multiplied by (1 << LZMS_COST_SHIFT).  These techniques are
715          * based on those used by LZMA.
716          *
717          * Note that in LZMS, a probability x really means x / 64, and 0 / 64 is
718          * really interpreted as 1 / 64 and 64 / 64 is really interpreted as
719          * 63 / 64.
720          */
721         for (u32 i = 0; i <= LZMS_PROBABILITY_MAX; i++) {
722                 u32 prob = i;
723
724                 if (prob == 0)
725                         prob = 1;
726                 else if (prob == LZMS_PROBABILITY_MAX)
727                         prob = LZMS_PROBABILITY_MAX - 1;
728
729         #ifdef LZMS_RC_COSTS_USE_FLOATING_POINT
730                 lzms_rc_costs[i] = log2((double)LZMS_PROBABILITY_MAX / prob) *
731                                         (1 << LZMS_COST_SHIFT);
732         #else
733                 u32 w = prob;
734                 u32 bit_count = 0;
735                 for (u32 j = 0; j < LZMS_COST_SHIFT; j++) {
736                         w *= w;
737                         bit_count <<= 1;
738                         while (w >= ((u32)1 << 16)) {
739                                 w >>= 1;
740                                 ++bit_count;
741                         }
742                 }
743                 lzms_rc_costs[i] = (LZMS_PROBABILITY_BITS << LZMS_COST_SHIFT) -
744                                    (15 + bit_count);
745         #endif
746         }
747 }
748
749 static void
750 lzms_init_rc_costs(void)
751 {
752         static pthread_once_t once = PTHREAD_ONCE_INIT;
753
754         pthread_once(&once, lzms_do_init_rc_costs);
755 }
756
757 /* Return the cost to range-encode the specified bit from the specified state.*/
758 static inline u32
759 lzms_rc_bit_cost(const struct lzms_range_encoder *enc, u8 cur_state, int bit)
760 {
761         u32 prob_zero;
762         u32 prob_correct;
763
764         prob_zero = enc->prob_entries[cur_state].num_recent_zero_bits;
765
766         if (bit == 0)
767                 prob_correct = prob_zero;
768         else
769                 prob_correct = LZMS_PROBABILITY_MAX - prob_zero;
770
771         return lzms_rc_costs[prob_correct];
772 }
773
774 /* Return the cost to Huffman-encode the specified symbol.  */
775 static inline u32
776 lzms_huffman_symbol_cost(const struct lzms_huffman_encoder *enc, unsigned sym)
777 {
778         return (u32)enc->lens[sym] << LZMS_COST_SHIFT;
779 }
780
781 /* Return the cost to encode the specified literal byte.  */
782 static inline u32
783 lzms_literal_cost(const struct lzms_compressor *c, unsigned literal,
784                   const struct lzms_adaptive_state *state)
785 {
786         return lzms_rc_bit_cost(&c->main_range_encoder, state->main_state, 0) +
787                lzms_huffman_symbol_cost(&c->literal_encoder, literal);
788 }
789
790 /* Update the table that directly provides the costs for small lengths.  */
791 static void
792 lzms_update_fast_length_costs(struct lzms_compressor *c)
793 {
794         u32 len;
795         int slot = -1;
796         u32 cost = 0;
797
798         for (len = 1; len < LZMS_NUM_FAST_LENGTHS; len++) {
799
800                 while (len >= lzms_length_slot_base[slot + 1]) {
801                         slot++;
802                         cost = (u32)(c->length_encoder.lens[slot] +
803                                      lzms_extra_length_bits[slot]) << LZMS_COST_SHIFT;
804                 }
805
806                 c->length_cost_fast[len] = cost;
807         }
808 }
809
810 /* Return the cost to encode the specified match length, which must be less than
811  * LZMS_NUM_FAST_LENGTHS.  */
812 static inline u32
813 lzms_fast_length_cost(const struct lzms_compressor *c, u32 length)
814 {
815         LZMS_ASSERT(length < LZMS_NUM_FAST_LENGTHS);
816         return c->length_cost_fast[length];
817 }
818
819 /* Return the cost to encode the specified LZ match offset.  */
820 static inline u32
821 lzms_lz_offset_cost(const struct lzms_compressor *c, u32 offset)
822 {
823         unsigned slot = lzms_get_offset_slot_fast(c, offset);
824
825         return (u32)(c->lz_offset_encoder.lens[slot] +
826                      lzms_extra_offset_bits[slot]) << LZMS_COST_SHIFT;
827 }
828
829 /*
830  * Consider coding the match at repeat offset index @rep_idx.  Consider each
831  * length from the minimum (2) to the full match length (@rep_len).
832  */
833 static inline void
834 lzms_consider_lz_repeat_offset_match(const struct lzms_compressor *c,
835                                      struct lzms_mc_pos_data *cur_optimum_ptr,
836                                      u32 rep_len, unsigned rep_idx)
837 {
838         u32 len;
839         u32 base_cost;
840         u32 cost;
841         unsigned i;
842
843         base_cost = cur_optimum_ptr->cost;
844
845         base_cost += lzms_rc_bit_cost(&c->main_range_encoder,
846                                       cur_optimum_ptr->state.main_state, 1);
847
848         base_cost += lzms_rc_bit_cost(&c->match_range_encoder,
849                                       cur_optimum_ptr->state.match_state, 0);
850
851         base_cost += lzms_rc_bit_cost(&c->lz_match_range_encoder,
852                                       cur_optimum_ptr->state.lz_match_state, 1);
853
854         for (i = 0; i < rep_idx; i++)
855                 base_cost += lzms_rc_bit_cost(&c->lz_repeat_match_range_encoders[i],
856                                               cur_optimum_ptr->state.lz_repeat_match_state[i], 1);
857
858         if (i < LZMS_NUM_RECENT_OFFSETS - 1)
859                 base_cost += lzms_rc_bit_cost(&c->lz_repeat_match_range_encoders[i],
860                                               cur_optimum_ptr->state.lz_repeat_match_state[i], 0);
861
862         len = 2;
863         do {
864                 cost = base_cost + lzms_fast_length_cost(c, len);
865                 if (cost < (cur_optimum_ptr + len)->cost) {
866                         (cur_optimum_ptr + len)->mc_item_data =
867                                 ((u64)rep_idx << MC_OFFSET_SHIFT) | len;
868                         (cur_optimum_ptr + len)->cost = cost;
869                 }
870         } while (++len <= rep_len);
871 }
872
873 /*
874  * Consider coding each match in @matches as an explicit offset match.
875  *
876  * @matches must be sorted by strictly increasing length and strictly increasing
877  * offset.  This is guaranteed by the match-finder.
878  *
879  * We consider each length from the minimum (2) to the longest
880  * (matches[num_matches - 1].len).  For each length, we consider only the
881  * smallest offset for which that length is available.  Although this is not
882  * guaranteed to be optimal due to the possibility of a larger offset costing
883  * less than a smaller offset to code, this is a very useful heuristic.
884  */
885 static inline void
886 lzms_consider_lz_explicit_offset_matches(const struct lzms_compressor *c,
887                                          struct lzms_mc_pos_data *cur_optimum_ptr,
888                                          const struct lz_match matches[],
889                                          u32 num_matches)
890 {
891         u32 len;
892         u32 i;
893         u32 base_cost;
894         u32 position_cost;
895         u32 cost;
896
897         base_cost = cur_optimum_ptr->cost;
898
899         base_cost += lzms_rc_bit_cost(&c->main_range_encoder,
900                                       cur_optimum_ptr->state.main_state, 1);
901
902         base_cost += lzms_rc_bit_cost(&c->match_range_encoder,
903                                       cur_optimum_ptr->state.match_state, 0);
904
905         base_cost += lzms_rc_bit_cost(&c->lz_match_range_encoder,
906                                       cur_optimum_ptr->state.lz_match_state, 0);
907         len = 2;
908         i = 0;
909         do {
910                 position_cost = base_cost + lzms_lz_offset_cost(c, matches[i].offset);
911                 do {
912                         cost = position_cost + lzms_fast_length_cost(c, len);
913                         if (cost < (cur_optimum_ptr + len)->cost) {
914                                 (cur_optimum_ptr + len)->mc_item_data =
915                                         ((u64)(matches[i].offset + LZMS_OFFSET_OFFSET)
916                                                 << MC_OFFSET_SHIFT) | len;
917                                 (cur_optimum_ptr + len)->cost = cost;
918                         }
919                 } while (++len <= matches[i].len);
920         } while (++i != num_matches);
921 }
922
923 static void
924 lzms_init_adaptive_state(struct lzms_adaptive_state *state)
925 {
926         unsigned i;
927
928         lzms_init_lz_lru_queue(&state->lru);
929         state->main_state = 0;
930         state->match_state = 0;
931         state->lz_match_state = 0;
932         for (i = 0; i < LZMS_NUM_RECENT_OFFSETS - 1; i++)
933                 state->lz_repeat_match_state[i] = 0;
934 }
935
936 static inline void
937 lzms_update_main_state(struct lzms_adaptive_state *state, int is_match)
938 {
939         state->main_state = ((state->main_state << 1) | is_match) % LZMS_NUM_MAIN_STATES;
940 }
941
942 static inline void
943 lzms_update_match_state(struct lzms_adaptive_state *state, int is_delta)
944 {
945         state->match_state = ((state->match_state << 1) | is_delta) % LZMS_NUM_MATCH_STATES;
946 }
947
948 static inline void
949 lzms_update_lz_match_state(struct lzms_adaptive_state *state, int is_repeat_offset)
950 {
951         state->lz_match_state = ((state->lz_match_state << 1) | is_repeat_offset) % LZMS_NUM_LZ_MATCH_STATES;
952 }
953
954 static inline void
955 lzms_update_lz_repeat_match_state(struct lzms_adaptive_state *state, int rep_idx)
956 {
957         int i;
958
959         for (i = 0; i < rep_idx; i++)
960                 state->lz_repeat_match_state[i] =
961                         ((state->lz_repeat_match_state[i] << 1) | 1) %
962                                 LZMS_NUM_LZ_REPEAT_MATCH_STATES;
963
964         if (i < LZMS_NUM_RECENT_OFFSETS - 1)
965                 state->lz_repeat_match_state[i] =
966                         ((state->lz_repeat_match_state[i] << 1) | 0) %
967                                 LZMS_NUM_LZ_REPEAT_MATCH_STATES;
968 }
969
970 /*
971  * The main near-optimal parsing routine.
972  *
973  * Briefly, the algorithm does an approximate minimum-cost path search to find a
974  * "near-optimal" sequence of matches and literals to output, based on the
975  * current cost model.  The algorithm steps forward, position by position (byte
976  * by byte), and updates the minimum cost path to reach each later position that
977  * can be reached using a match or literal from the current position.  This is
978  * essentially Dijkstra's algorithm in disguise: the graph nodes are positions,
979  * the graph edges are possible matches/literals to code, and the cost of each
980  * edge is the estimated number of bits that will be required to output the
981  * corresponding match or literal.  But one difference is that we actually
982  * compute the lowest-cost path in pieces, where each piece is terminated when
983  * there are no choices to be made.
984  *
985  * Notes:
986  *
987  * - This does not output any delta matches.
988  *
989  * - The costs of literals and matches are estimated using the range encoder
990  *   states and the semi-adaptive Huffman codes.  Except for range encoding
991  *   states, costs are assumed to be constant throughout a single run of the
992  *   parsing algorithm, which can parse up to @optim_array_length bytes of data.
993  *   This introduces a source of inaccuracy because the probabilities and
994  *   Huffman codes can change over this part of the data.
995  */
996 static void
997 lzms_near_optimal_parse(struct lzms_compressor *c)
998 {
999         const u8 *window_ptr;
1000         const u8 *window_end;
1001         struct lzms_mc_pos_data *cur_optimum_ptr;
1002         struct lzms_mc_pos_data *end_optimum_ptr;
1003         u32 num_matches;
1004         u32 longest_len;
1005         u32 rep_max_len;
1006         unsigned rep_max_idx;
1007         unsigned literal;
1008         unsigned i;
1009         u32 cost;
1010         u32 len;
1011         u32 offset_data;
1012
1013         window_ptr = c->cur_window;
1014         window_end = window_ptr + c->cur_window_size;
1015
1016         lzms_init_adaptive_state(&c->optimum[0].state);
1017
1018 begin:
1019         /* Start building a new list of items, which will correspond to the next
1020          * piece of the overall minimum-cost path.  */
1021
1022         cur_optimum_ptr = c->optimum;
1023         cur_optimum_ptr->cost = 0;
1024         end_optimum_ptr = cur_optimum_ptr;
1025
1026         /* States should currently be consistent with the encoders.  */
1027         LZMS_ASSERT(cur_optimum_ptr->state.main_state == c->main_range_encoder.state);
1028         LZMS_ASSERT(cur_optimum_ptr->state.match_state == c->match_range_encoder.state);
1029         LZMS_ASSERT(cur_optimum_ptr->state.lz_match_state == c->lz_match_range_encoder.state);
1030         for (i = 0; i < LZMS_NUM_RECENT_OFFSETS - 1; i++)
1031                 LZMS_ASSERT(cur_optimum_ptr->state.lz_repeat_match_state[i] ==
1032                             c->lz_repeat_match_range_encoders[i].state);
1033
1034         if (window_ptr == window_end)
1035                 return;
1036
1037         /* The following loop runs once for each per byte in the window, except
1038          * in a couple shortcut cases.  */
1039         for (;;) {
1040
1041                 /* Find explicit offset matches with the current position.  */
1042                 num_matches = lz_mf_get_matches(c->mf, c->matches);
1043
1044                 if (num_matches) {
1045                         /*
1046                          * Find the longest repeat offset match with the current
1047                          * position.
1048                          *
1049                          * Heuristics:
1050                          *
1051                          * - Only search for repeat offset matches if the
1052                          *   match-finder already found at least one match.
1053                          *
1054                          * - Only consider the longest repeat offset match.  It
1055                          *   seems to be rare for the optimal parse to include a
1056                          *   repeat offset match that doesn't have the longest
1057                          *   length (allowing for the possibility that not all
1058                          *   of that length is actually used).
1059                          */
1060                         if (likely(window_ptr - c->cur_window >= LZMS_MAX_INIT_RECENT_OFFSET)) {
1061                                 BUILD_BUG_ON(LZMS_NUM_RECENT_OFFSETS != 3);
1062                                 rep_max_len = lz_repsearch3(window_ptr,
1063                                                             window_end - window_ptr,
1064                                                             cur_optimum_ptr->state.lru.recent_offsets,
1065                                                             &rep_max_idx);
1066                         } else {
1067                                 rep_max_len = 0;
1068                         }
1069
1070                         if (rep_max_len) {
1071                                 /* If there's a very long repeat offset match,
1072                                  * choose it immediately.  */
1073                                 if (rep_max_len >= c->params.nice_match_length) {
1074
1075                                         lz_mf_skip_positions(c->mf, rep_max_len - 1);
1076                                         window_ptr += rep_max_len;
1077
1078                                         if (cur_optimum_ptr != c->optimum)
1079                                                 lzms_encode_item_list(c, cur_optimum_ptr);
1080
1081                                         lzms_encode_lz_repeat_offset_match(c, rep_max_len,
1082                                                                            rep_max_idx);
1083
1084                                         c->optimum[0].state = cur_optimum_ptr->state;
1085
1086                                         lzms_update_main_state(&c->optimum[0].state, 1);
1087                                         lzms_update_match_state(&c->optimum[0].state, 0);
1088                                         lzms_update_lz_match_state(&c->optimum[0].state, 1);
1089                                         lzms_update_lz_repeat_match_state(&c->optimum[0].state,
1090                                                                           rep_max_idx);
1091
1092                                         c->optimum[0].state.lru.upcoming_offset =
1093                                                 c->optimum[0].state.lru.recent_offsets[rep_max_idx];
1094
1095                                         for (i = rep_max_idx; i < LZMS_NUM_RECENT_OFFSETS; i++)
1096                                                 c->optimum[0].state.lru.recent_offsets[i] =
1097                                                         c->optimum[0].state.lru.recent_offsets[i + 1];
1098
1099                                         lzms_update_lz_lru_queue(&c->optimum[0].state.lru);
1100                                         goto begin;
1101                                 }
1102
1103                                 /* If reaching any positions for the first time,
1104                                  * initialize their costs to "infinity".  */
1105                                 while (end_optimum_ptr < cur_optimum_ptr + rep_max_len)
1106                                         (++end_optimum_ptr)->cost = MC_INFINITE_COST;
1107
1108                                 /* Consider coding a repeat offset match.  */
1109                                 lzms_consider_lz_repeat_offset_match(c, cur_optimum_ptr,
1110                                                                      rep_max_len, rep_max_idx);
1111                         }
1112
1113                         longest_len = c->matches[num_matches - 1].len;
1114
1115                         /* If there's a very long explicit offset match, choose
1116                          * it immediately.  */
1117                         if (longest_len >= c->params.nice_match_length) {
1118
1119                                 lz_mf_skip_positions(c->mf, longest_len - 1);
1120                                 window_ptr += longest_len;
1121
1122                                 if (cur_optimum_ptr != c->optimum)
1123                                         lzms_encode_item_list(c, cur_optimum_ptr);
1124
1125                                 lzms_encode_lz_explicit_offset_match(c, longest_len,
1126                                                                      c->matches[num_matches - 1].offset);
1127
1128                                 c->optimum[0].state = cur_optimum_ptr->state;
1129
1130                                 lzms_update_main_state(&c->optimum[0].state, 1);
1131                                 lzms_update_match_state(&c->optimum[0].state, 0);
1132                                 lzms_update_lz_match_state(&c->optimum[0].state, 0);
1133
1134                                 c->optimum[0].state.lru.upcoming_offset =
1135                                         c->matches[num_matches - 1].offset;
1136
1137                                 lzms_update_lz_lru_queue(&c->optimum[0].state.lru);
1138                                 goto begin;
1139                         }
1140
1141                         /* If reaching any positions for the first time,
1142                          * initialize their costs to "infinity".  */
1143                         while (end_optimum_ptr < cur_optimum_ptr + longest_len)
1144                                 (++end_optimum_ptr)->cost = MC_INFINITE_COST;
1145
1146                         /* Consider coding an explicit offset match.  */
1147                         lzms_consider_lz_explicit_offset_matches(c, cur_optimum_ptr,
1148                                                                  c->matches, num_matches);
1149                 } else {
1150                         /* No matches found.  The only choice at this position
1151                          * is to code a literal.  */
1152
1153                         if (end_optimum_ptr == cur_optimum_ptr)
1154                                 (++end_optimum_ptr)->cost = MC_INFINITE_COST;
1155                 }
1156
1157                 /* Consider coding a literal.
1158
1159                  * To avoid an extra unpredictable brench, actually checking the
1160                  * preferability of coding a literal is integrated into the
1161                  * adaptive state update code below.  */
1162                 literal = *window_ptr++;
1163                 cost = cur_optimum_ptr->cost +
1164                        lzms_literal_cost(c, literal, &cur_optimum_ptr->state);
1165
1166                 /* Advance to the next position.  */
1167                 cur_optimum_ptr++;
1168
1169                 /* The lowest-cost path to the current position is now known.
1170                  * Finalize the adaptive state that results from taking this
1171                  * lowest-cost path.  */
1172
1173                 if (cost < cur_optimum_ptr->cost) {
1174                         /* Literal  */
1175                         cur_optimum_ptr->cost = cost;
1176                         cur_optimum_ptr->mc_item_data = ((u64)literal << MC_OFFSET_SHIFT) | 1;
1177
1178                         cur_optimum_ptr->state = (cur_optimum_ptr - 1)->state;
1179
1180                         lzms_update_main_state(&cur_optimum_ptr->state, 0);
1181
1182                         cur_optimum_ptr->state.lru.upcoming_offset = 0;
1183                 } else {
1184                         /* LZ match  */
1185                         len = cur_optimum_ptr->mc_item_data & MC_LEN_MASK;
1186                         offset_data = cur_optimum_ptr->mc_item_data >> MC_OFFSET_SHIFT;
1187
1188                         cur_optimum_ptr->state = (cur_optimum_ptr - len)->state;
1189
1190                         lzms_update_main_state(&cur_optimum_ptr->state, 1);
1191                         lzms_update_match_state(&cur_optimum_ptr->state, 0);
1192
1193                         if (offset_data >= LZMS_NUM_RECENT_OFFSETS) {
1194
1195                                 /* Explicit offset LZ match  */
1196
1197                                 lzms_update_lz_match_state(&cur_optimum_ptr->state, 0);
1198
1199                                 cur_optimum_ptr->state.lru.upcoming_offset =
1200                                         offset_data - LZMS_OFFSET_OFFSET;
1201                         } else {
1202                                 /* Repeat offset LZ match  */
1203
1204                                 lzms_update_lz_match_state(&cur_optimum_ptr->state, 1);
1205                                 lzms_update_lz_repeat_match_state(&cur_optimum_ptr->state,
1206                                                                   offset_data);
1207
1208                                 cur_optimum_ptr->state.lru.upcoming_offset =
1209                                         cur_optimum_ptr->state.lru.recent_offsets[offset_data];
1210
1211                                 for (i = offset_data; i < LZMS_NUM_RECENT_OFFSETS; i++)
1212                                         cur_optimum_ptr->state.lru.recent_offsets[i] =
1213                                                 cur_optimum_ptr->state.lru.recent_offsets[i + 1];
1214                         }
1215                 }
1216
1217                 lzms_update_lz_lru_queue(&cur_optimum_ptr->state.lru);
1218
1219                 /*
1220                  * This loop will terminate when either of the following
1221                  * conditions is true:
1222                  *
1223                  * (1) cur_optimum_ptr == end_optimum_ptr
1224                  *
1225                  *      There are no paths that extend beyond the current
1226                  *      position.  In this case, any path to a later position
1227                  *      must pass through the current position, so we can go
1228                  *      ahead and choose the list of items that led to this
1229                  *      position.
1230                  *
1231                  * (2) cur_optimum_ptr == c->optimum_end
1232                  *
1233                  *      This bounds the number of times the algorithm can step
1234                  *      forward before it is guaranteed to start choosing items.
1235                  *      This limits the memory usage.  It also guarantees that
1236                  *      the parser will not go too long without updating the
1237                  *      probability tables.
1238                  *
1239                  * Note: no check for end-of-window is needed because
1240                  * end-of-window will trigger condition (1).
1241                  */
1242                 if (cur_optimum_ptr == end_optimum_ptr ||
1243                     cur_optimum_ptr == c->optimum_end)
1244                 {
1245                         c->optimum[0].state = cur_optimum_ptr->state;
1246                         break;
1247                 }
1248         }
1249
1250         /* Output the current list of items that constitute the minimum-cost
1251          * path to the current position.  */
1252         lzms_encode_item_list(c, cur_optimum_ptr);
1253         goto begin;
1254 }
1255
1256 static void
1257 lzms_init_range_encoder(struct lzms_range_encoder *enc,
1258                         struct lzms_range_encoder_raw *rc, u32 num_states)
1259 {
1260         enc->rc = rc;
1261         enc->state = 0;
1262         LZMS_ASSERT(is_power_of_2(num_states));
1263         enc->mask = num_states - 1;
1264         lzms_init_probability_entries(enc->prob_entries, num_states);
1265 }
1266
1267 static void
1268 lzms_init_huffman_encoder(struct lzms_huffman_encoder *enc,
1269                           struct lzms_output_bitstream *os,
1270                           unsigned num_syms,
1271                           unsigned rebuild_freq)
1272 {
1273         enc->os = os;
1274         enc->num_syms_written = 0;
1275         enc->rebuild_freq = rebuild_freq;
1276         enc->num_syms = num_syms;
1277         for (unsigned i = 0; i < num_syms; i++)
1278                 enc->sym_freqs[i] = 1;
1279
1280         make_canonical_huffman_code(enc->num_syms,
1281                                     LZMS_MAX_CODEWORD_LEN,
1282                                     enc->sym_freqs,
1283                                     enc->lens,
1284                                     enc->codewords);
1285 }
1286
1287 /* Prepare the LZMS compressor for compressing a block of data.  */
1288 static void
1289 lzms_prepare_compressor(struct lzms_compressor *c, const u8 *udata, u32 ulen,
1290                         le16 *cdata, u32 clen16)
1291 {
1292         unsigned num_offset_slots;
1293
1294         /* Copy the uncompressed data into the @c->cur_window buffer.  */
1295         memcpy(c->cur_window, udata, ulen);
1296         c->cur_window_size = ulen;
1297
1298         /* Initialize the raw range encoder (writing forwards).  */
1299         lzms_range_encoder_raw_init(&c->rc, cdata, clen16);
1300
1301         /* Initialize the output bitstream for Huffman symbols and verbatim bits
1302          * (writing backwards).  */
1303         lzms_output_bitstream_init(&c->os, cdata, clen16);
1304
1305         /* Calculate the number of offset slots required.  */
1306         num_offset_slots = lzms_get_offset_slot(ulen - 1) + 1;
1307
1308         /* Initialize a Huffman encoder for each alphabet.  */
1309         lzms_init_huffman_encoder(&c->literal_encoder, &c->os,
1310                                   LZMS_NUM_LITERAL_SYMS,
1311                                   LZMS_LITERAL_CODE_REBUILD_FREQ);
1312
1313         lzms_init_huffman_encoder(&c->lz_offset_encoder, &c->os,
1314                                   num_offset_slots,
1315                                   LZMS_LZ_OFFSET_CODE_REBUILD_FREQ);
1316
1317         lzms_init_huffman_encoder(&c->length_encoder, &c->os,
1318                                   LZMS_NUM_LENGTH_SYMS,
1319                                   LZMS_LENGTH_CODE_REBUILD_FREQ);
1320
1321         lzms_init_huffman_encoder(&c->delta_offset_encoder, &c->os,
1322                                   num_offset_slots,
1323                                   LZMS_DELTA_OFFSET_CODE_REBUILD_FREQ);
1324
1325         lzms_init_huffman_encoder(&c->delta_power_encoder, &c->os,
1326                                   LZMS_NUM_DELTA_POWER_SYMS,
1327                                   LZMS_DELTA_POWER_CODE_REBUILD_FREQ);
1328
1329         /* Initialize range encoders, all of which wrap around the same
1330          * lzms_range_encoder_raw.  */
1331         lzms_init_range_encoder(&c->main_range_encoder,
1332                                 &c->rc, LZMS_NUM_MAIN_STATES);
1333
1334         lzms_init_range_encoder(&c->match_range_encoder,
1335                                 &c->rc, LZMS_NUM_MATCH_STATES);
1336
1337         lzms_init_range_encoder(&c->lz_match_range_encoder,
1338                                 &c->rc, LZMS_NUM_LZ_MATCH_STATES);
1339
1340         for (unsigned i = 0; i < ARRAY_LEN(c->lz_repeat_match_range_encoders); i++)
1341                 lzms_init_range_encoder(&c->lz_repeat_match_range_encoders[i],
1342                                         &c->rc, LZMS_NUM_LZ_REPEAT_MATCH_STATES);
1343
1344         lzms_init_range_encoder(&c->delta_match_range_encoder,
1345                                 &c->rc, LZMS_NUM_DELTA_MATCH_STATES);
1346
1347         for (unsigned i = 0; i < ARRAY_LEN(c->delta_repeat_match_range_encoders); i++)
1348                 lzms_init_range_encoder(&c->delta_repeat_match_range_encoders[i],
1349                                         &c->rc, LZMS_NUM_DELTA_REPEAT_MATCH_STATES);
1350
1351         /* Set initial length costs for lengths < LZMS_NUM_FAST_LENGTHS.  */
1352         lzms_update_fast_length_costs(c);
1353 }
1354
1355 /* Flush the output streams, prepare the final compressed data, and return its
1356  * size in bytes.
1357  *
1358  * A return value of 0 indicates that the data could not be compressed to fit in
1359  * the available space.  */
1360 static size_t
1361 lzms_finalize(struct lzms_compressor *c, u8 *cdata, size_t csize_avail)
1362 {
1363         size_t num_forwards_bytes;
1364         size_t num_backwards_bytes;
1365
1366         /* Flush both the forwards and backwards streams, and make sure they
1367          * didn't cross each other and start overwriting each other's data.  */
1368         if (!lzms_output_bitstream_flush(&c->os))
1369                 return 0;
1370
1371         if (!lzms_range_encoder_raw_flush(&c->rc))
1372                 return 0;
1373
1374         if (c->rc.next > c->os.next)
1375                 return 0;
1376
1377         /* Now the compressed buffer contains the data output by the forwards
1378          * bitstream, then empty space, then data output by the backwards
1379          * bitstream.  Move the data output by the backwards bitstream to be
1380          * adjacent to the data output by the forward bitstream, and calculate
1381          * the compressed size that this results in.  */
1382         num_forwards_bytes = (u8*)c->rc.next - (u8*)cdata;
1383         num_backwards_bytes = ((u8*)cdata + csize_avail) - (u8*)c->os.next;
1384
1385         memmove(cdata + num_forwards_bytes, c->os.next, num_backwards_bytes);
1386
1387         return num_forwards_bytes + num_backwards_bytes;
1388 }
1389
1390 /* Set internal compression parameters for the specified compression level and
1391  * maximum window size.  */
1392 static void
1393 lzms_build_params(unsigned int compression_level,
1394                   struct lzms_compressor_params *params)
1395 {
1396         /* Allow length 2 matches if the compression level is sufficiently high.
1397          */
1398         if (compression_level >= 45)
1399                 params->min_match_length = 2;
1400         else
1401                 params->min_match_length = 3;
1402
1403         /* Scale nice_match_length and max_search_depth with the compression
1404          * level.  But to allow an optimization on length cost calculations,
1405          * don't allow nice_match_length to exceed LZMS_NUM_FAST_LENGTH.  */
1406         params->nice_match_length = ((u64)compression_level * 32) / 50;
1407         if (params->nice_match_length < params->min_match_length)
1408                 params->nice_match_length = params->min_match_length;
1409         if (params->nice_match_length > LZMS_NUM_FAST_LENGTHS)
1410                 params->nice_match_length = LZMS_NUM_FAST_LENGTHS;
1411         params->max_search_depth = compression_level;
1412
1413         params->optim_array_length = 1024;
1414 }
1415
1416 /* Given the internal compression parameters and maximum window size, build the
1417  * Lempel-Ziv match-finder parameters.  */
1418 static void
1419 lzms_build_mf_params(const struct lzms_compressor_params *lzms_params,
1420                      u32 max_window_size, struct lz_mf_params *mf_params)
1421 {
1422         memset(mf_params, 0, sizeof(*mf_params));
1423
1424         /* Choose an appropriate match-finding algorithm.  */
1425         if (max_window_size <= 2097152)
1426                 mf_params->algorithm = LZ_MF_BINARY_TREES;
1427         else if (max_window_size <= 33554432)
1428                 mf_params->algorithm = LZ_MF_LCP_INTERVAL_TREE;
1429         else
1430                 mf_params->algorithm = LZ_MF_LINKED_SUFFIX_ARRAY;
1431
1432         mf_params->max_window_size = max_window_size;
1433         mf_params->min_match_len = lzms_params->min_match_length;
1434         mf_params->max_search_depth = lzms_params->max_search_depth;
1435         mf_params->nice_match_len = lzms_params->nice_match_length;
1436 }
1437
1438 static void
1439 lzms_free_compressor(void *_c);
1440
1441 static u64
1442 lzms_get_needed_memory(size_t max_block_size, unsigned int compression_level)
1443 {
1444         struct lzms_compressor_params params;
1445         struct lz_mf_params mf_params;
1446         u64 size = 0;
1447
1448         if (max_block_size >= INT32_MAX)
1449                 return 0;
1450
1451         lzms_build_params(compression_level, &params);
1452         lzms_build_mf_params(&params, max_block_size, &mf_params);
1453
1454         size += sizeof(struct lzms_compressor);
1455
1456         /* cur_window */
1457         size += max_block_size;
1458
1459         /* mf */
1460         size += lz_mf_get_needed_memory(mf_params.algorithm, max_block_size);
1461
1462         /* matches */
1463         size += min(params.max_search_depth, params.nice_match_length) *
1464                 sizeof(struct lz_match);
1465
1466         /* optimum */
1467         size += (params.optim_array_length + params.nice_match_length) *
1468                 sizeof(struct lzms_mc_pos_data);
1469
1470         return size;
1471 }
1472
1473 static int
1474 lzms_create_compressor(size_t max_block_size, unsigned int compression_level,
1475                        void **ctx_ret)
1476 {
1477         struct lzms_compressor *c;
1478         struct lzms_compressor_params params;
1479         struct lz_mf_params mf_params;
1480
1481         if (max_block_size >= INT32_MAX)
1482                 return WIMLIB_ERR_INVALID_PARAM;
1483
1484         lzms_build_params(compression_level, &params);
1485         lzms_build_mf_params(&params, max_block_size, &mf_params);
1486         if (!lz_mf_params_valid(&mf_params))
1487                 return WIMLIB_ERR_INVALID_PARAM;
1488
1489         c = CALLOC(1, sizeof(struct lzms_compressor));
1490         if (!c)
1491                 goto oom;
1492
1493         c->params = params;
1494
1495         c->cur_window = MALLOC(max_block_size);
1496         if (!c->cur_window)
1497                 goto oom;
1498
1499         c->mf = lz_mf_alloc(&mf_params);
1500         if (!c->mf)
1501                 goto oom;
1502
1503         c->matches = MALLOC(min(params.max_search_depth,
1504                                 params.nice_match_length) *
1505                             sizeof(struct lz_match));
1506         if (!c->matches)
1507                 goto oom;
1508
1509         c->optimum = MALLOC((params.optim_array_length +
1510                              params.nice_match_length) *
1511                             sizeof(struct lzms_mc_pos_data));
1512         if (!c->optimum)
1513                 goto oom;
1514         c->optimum_end = &c->optimum[params.optim_array_length];
1515
1516         lzms_init_rc_costs();
1517
1518         lzms_init_fast_slots(c);
1519
1520         *ctx_ret = c;
1521         return 0;
1522
1523 oom:
1524         lzms_free_compressor(c);
1525         return WIMLIB_ERR_NOMEM;
1526 }
1527
1528 static size_t
1529 lzms_compress(const void *uncompressed_data, size_t uncompressed_size,
1530               void *compressed_data, size_t compressed_size_avail, void *_c)
1531 {
1532         struct lzms_compressor *c = _c;
1533
1534         /* Don't bother compressing extremely small inputs.  */
1535         if (uncompressed_size < 4)
1536                 return 0;
1537
1538         /* Cap the available compressed size to a 32-bit integer and round it
1539          * down to the nearest multiple of 2.  */
1540         if (compressed_size_avail > UINT32_MAX)
1541                 compressed_size_avail = UINT32_MAX;
1542         if (compressed_size_avail & 1)
1543                 compressed_size_avail--;
1544
1545         /* Initialize the compressor structures.  */
1546         lzms_prepare_compressor(c, uncompressed_data, uncompressed_size,
1547                                 compressed_data, compressed_size_avail / 2);
1548
1549         /* Preprocess the uncompressed data.  */
1550         lzms_x86_filter(c->cur_window, c->cur_window_size,
1551                         c->last_target_usages, false);
1552
1553         /* Load the window into the match-finder.  */
1554         lz_mf_load_window(c->mf, c->cur_window, c->cur_window_size);
1555
1556         /* Compute and encode a literal/match sequence that decompresses to the
1557          * preprocessed data.  */
1558         lzms_near_optimal_parse(c);
1559
1560         /* Return the compressed data size or 0.  */
1561         return lzms_finalize(c, compressed_data, compressed_size_avail);
1562 }
1563
1564 static void
1565 lzms_free_compressor(void *_c)
1566 {
1567         struct lzms_compressor *c = _c;
1568
1569         if (c) {
1570                 FREE(c->cur_window);
1571                 lz_mf_free(c->mf);
1572                 FREE(c->matches);
1573                 FREE(c->optimum);
1574                 FREE(c);
1575         }
1576 }
1577
1578 const struct compressor_ops lzms_compressor_ops = {
1579         .get_needed_memory  = lzms_get_needed_memory,
1580         .create_compressor  = lzms_create_compressor,
1581         .compress           = lzms_compress,
1582         .free_compressor    = lzms_free_compressor,
1583 };