A few cleanups and fixes from recent changes

[wimlib] / src / lzms_compress.c

diff --git a/src/lzms_compress.c b/src/lzms_compress.c
index abf39f5e89f194d5ff0fb087b28a3c4c73b57100..8db49668d296bf8d0340235ae78a4bc78aa372cf 100644 (file)
--- a/src/lzms_compress.c
+++ b/src/lzms_compress.c
@@ -287,6 +287,10 @@ struct lzms_compressor {
         */
        bool use_delta_matches;
 
+       /* If true, the compressor need not preserve the input buffer if it
+        * compresses the data successfully.  */
+       bool destructive;
+
        /* 'last_target_usages' is a large array that is only needed for
         * preprocessing, so it is in union with fields that don't need to be
         * initialized until after preprocessing.  */
@@ -303,7 +307,8 @@ struct lzms_compressor {
        u32 next_delta_hashes[NUM_POWERS_TO_CONSIDER];
 
        /* The per-byte graph nodes for near-optimal parsing  */
-       struct lzms_optimum_node optimum_nodes[NUM_OPTIM_NODES + MAX_FAST_LENGTH];
+       struct lzms_optimum_node optimum_nodes[NUM_OPTIM_NODES + MAX_FAST_LENGTH +
+                                              1 + MAX_FAST_LENGTH];
 
        /* Table: length => current cost for small match lengths  */
        u32 fast_length_cost_tab[MAX_FAST_LENGTH + 1];
@@ -942,7 +947,7 @@ lzms_encode_item_list(struct lzms_compressor *c,
 }
 
 /******************************************************************************
- *                             Cost evalution                                 *
+ *                             Cost evaluation                                *
  ******************************************************************************/
 
 /*
@@ -1760,6 +1765,11 @@ begin:
                                                                        span);
 
                                const u32 raw_offset = offset >> power;
+
+                               if (unlikely(raw_offset > DELTA_SOURCE_RAW_OFFSET_MASK -
+                                                         (LZMS_NUM_DELTA_REPS - 1)))
+                                       continue;
+
                                const u32 pair = (power << DELTA_SOURCE_POWER_SHIFT) |
                                                 raw_offset;
                                const u32 source = DELTA_SOURCE_TAG |
@@ -2100,7 +2110,8 @@ lzms_finalize(struct lzms_compressor *c)
 }
 
 static u64
-lzms_get_needed_memory(size_t max_bufsize, unsigned compression_level)
+lzms_get_needed_memory(size_t max_bufsize, unsigned compression_level,
+                      bool destructive)
 {
        u64 size = 0;
 
@@ -2109,8 +2120,8 @@ lzms_get_needed_memory(size_t max_bufsize, unsigned compression_level)
 
        size += sizeof(struct lzms_compressor);
 
-       /* in_buffer */
-       size += max_bufsize;
+       if (!destructive)
+               size += max_bufsize; /* in_buffer */
 
        /* mf */
        size += lcpit_matchfinder_get_needed_memory(max_bufsize);
@@ -2120,7 +2131,7 @@ lzms_get_needed_memory(size_t max_bufsize, unsigned compression_level)
 
 static int
 lzms_create_compressor(size_t max_bufsize, unsigned compression_level,
-                      void **c_ret)
+                      bool destructive, void **c_ret)
 {
        struct lzms_compressor *c;
        u32 nice_match_len;
@@ -2132,6 +2143,8 @@ lzms_create_compressor(size_t max_bufsize, unsigned compression_level,
        if (!c)
                goto oom0;
 
+       c->destructive = destructive;
+
        /* Scale nice_match_len with the compression level.  But to allow an
         * optimization for length cost calculations, don't allow nice_match_len
         * to exceed MAX_FAST_LENGTH.  */
@@ -2142,9 +2155,11 @@ lzms_create_compressor(size_t max_bufsize, unsigned compression_level,
        c->try_lit_lzrep0 = (compression_level >= 60);
        c->try_lzrep_lit_lzrep0 = (compression_level >= 60);
 
-       c->in_buffer = MALLOC(max_bufsize);
-       if (!c->in_buffer)
-               goto oom1;
+       if (!c->destructive) {
+               c->in_buffer = MALLOC(max_bufsize);
+               if (!c->in_buffer)
+                       goto oom1;
+       }
 
        if (!lcpit_matchfinder_init(&c->mf, max_bufsize, 2, nice_match_len))
                goto oom2;
@@ -2156,7 +2171,8 @@ lzms_create_compressor(size_t max_bufsize, unsigned compression_level,
        return 0;
 
 oom2:
-       FREE(c->in_buffer);
+       if (!c->destructive)
+               FREE(c->in_buffer);
 oom1:
        ALIGNED_FREE(c);
 oom0:
@@ -2168,13 +2184,17 @@ lzms_compress(const void *in, size_t in_nbytes,
              void *out, size_t out_nbytes_avail, void *_c)
 {
        struct lzms_compressor *c = _c;
+       size_t result;
 
        /* Don't bother trying to compress extremely small inputs.  */
        if (in_nbytes < 4)
                return 0;
 
        /* Copy the input data into the internal buffer and preprocess it.  */
-       memcpy(c->in_buffer, in, in_nbytes);
+       if (c->destructive)
+               c->in_buffer = (void *)in;
+       else
+               memcpy(c->in_buffer, in, in_nbytes);
        c->in_nbytes = in_nbytes;
        lzms_x86_filter(c->in_buffer, in_nbytes, c->last_target_usages, false);
 
@@ -2187,13 +2207,16 @@ lzms_compress(const void *in, size_t in_nbytes,
        lzms_range_encoder_init(&c->rc, out, out_nbytes_avail / sizeof(le16));
        lzms_output_bitstream_init(&c->os, out, out_nbytes_avail / sizeof(le16));
        lzms_init_states_and_probabilities(c);
-       lzms_init_huffman_codes(c, lzms_get_num_offset_slots(in_nbytes));
+       lzms_init_huffman_codes(c, lzms_get_num_offset_slots(c->in_nbytes));
 
        /* The main loop: parse and encode.  */
        lzms_near_optimal_parse(c);
 
        /* Return the compressed data size or 0.  */
-       return lzms_finalize(c);
+       result = lzms_finalize(c);
+       if (!result && c->destructive)
+               lzms_x86_filter(c->in_buffer, c->in_nbytes, c->last_target_usages, true);
+       return result;
 }
 
 static void
@@ -2201,7 +2224,8 @@ lzms_free_compressor(void *_c)
 {
        struct lzms_compressor *c = _c;
 
-       FREE(c->in_buffer);
+       if (!c->destructive)
+               FREE(c->in_buffer);
        lcpit_matchfinder_destroy(&c->mf);
        ALIGNED_FREE(c);
 }