Fix a buffer overflow in the LZMA encoder. It was due to my

misunderstanding of the code. There's no tiny fix for this
problem, so I also cleaned up the code in general.

This reduces the speed of the encoder 2-5 % in the fastest
compression mode ("lzma -1"). High compression modes should
have no noticeable performance difference.

This commit breaks things (especially LZMA_SYNC_FLUSH) but I
will fix them once the new format and LZMA2 has been roughly
implemented. Plain LZMA won't support LZMA_SYNC_FLUSH at all
and won't be supported in the new .lzma format. This may
change still but this is what it looks like now.

Support for known uncompressed size (that is, LZMA or LZMA2
without EOPM) is likely to go away. This means there will
be API changes.

1 files changed, 205 insertions, 178 deletions

diff --git a/src/liblzma/rangecoder/range_encoder.h b/src/liblzma/rangecoder/range_encoder.h
index b216e648..b156ee7f 100644
--- a/src/liblzma/rangecoder/range_encoder.h
+++ b/src/liblzma/rangecoder/range_encoder.h
@@ -4,7 +4,7 @@
 /// \brief      Range Encoder
 //
 //  Copyright (C) 1999-2006 Igor Pavlov
-//  Copyright (C) 2007 Lasse Collin
+//  Copyright (C) 2007-2008 Lasse Collin
 //
 //  This library is free software; you can redistribute it and/or
 //  modify it under the terms of the GNU Lesser General Public
@@ -24,14 +24,218 @@
 #include "range_common.h"
 
 
+/// Maximum number of symbols that can be put pending into lzma_range_encoder
+/// structure between calls to lzma_rc_encode(). For LZMA, 52+5 is enough
+/// (match with big distance and length followed by range encoder flush).
+#define RC_SYMBOLS_MAX 58
+
+
 typedef struct {
 	uint64_t low;
 	uint64_t cache_size;
 	uint32_t range;
 	uint8_t cache;
+
+	/// Number of symbols in the tables
+	size_t count;
+
+	/// rc_encode()'s position in the tables
+	size_t pos;
+
+	/// Symbols to encode
+	enum {
+		RC_BIT_0,
+		RC_BIT_1,
+		RC_DIRECT_0,
+		RC_DIRECT_1,
+		RC_FLUSH,
+	} symbols[RC_SYMBOLS_MAX];
+
+	/// Probabilities associated with RC_BIT_0 or RC_BIT_1
+	probability *probs[RC_SYMBOLS_MAX];
+
 } lzma_range_encoder;
 
 
+static inline void
+rc_reset(lzma_range_encoder *rc)
+{
+	rc->low = 0;
+	rc->cache_size = 1;
+	rc->range = UINT32_MAX;
+	rc->cache = 0;
+	rc->count = 0;
+	rc->pos = 0;
+}
+
+
+static inline void
+rc_bit(lzma_range_encoder *rc, probability *prob, uint32_t bit)
+{
+	rc->symbols[rc->count] = bit;
+	rc->probs[rc->count] = prob;
+	++rc->count;
+}
+
+
+static inline void
+rc_bittree(lzma_range_encoder *rc, probability *probs,
+		uint32_t bit_count, uint32_t symbol)
+{
+	uint32_t model_index = 1;
+
+	do {
+		const uint32_t bit = (symbol >> --bit_count) & 1;
+		rc_bit(rc, &probs[model_index], bit);
+		model_index = (model_index << 1) | bit;
+	} while (bit_count != 0);
+}
+
+
+static inline void
+rc_bittree_reverse(lzma_range_encoder *rc, probability *probs,
+		uint32_t bit_count, uint32_t symbol)
+{
+	uint32_t model_index = 1;
+
+	do {
+		const uint32_t bit = symbol & 1;
+		symbol >>= 1;
+		rc_bit(rc, &probs[model_index], bit);
+		model_index = (model_index << 1) | bit;
+	} while (--bit_count != 0);
+}
+
+
+static inline void
+rc_direct(lzma_range_encoder *rc,
+		uint32_t value, uint32_t bit_count)
+{
+	do {
+		rc->symbols[rc->count++]
+				= RC_DIRECT_0 + ((value >> --bit_count) & 1);
+	} while (bit_count != 0);
+}
+
+
+static inline void
+rc_flush(lzma_range_encoder *rc)
+{
+	for (size_t i = 0; i < 5; ++i)
+		rc->symbols[rc->count++] = RC_FLUSH;
+}
+
+
+static inline bool
+rc_shift_low(lzma_range_encoder *rc,
+		uint8_t *out, size_t *out_pos, size_t out_size)
+{
+	if ((uint32_t)(rc->low) < (uint32_t)(0xFF000000)
+			|| (uint32_t)(rc->low >> 32) != 0) {
+		do {
+			if (*out_pos == out_size)
+				return true;
+
+			out[*out_pos] = rc->cache + (uint8_t)(rc->low >> 32);
+			++*out_pos;
+			rc->cache = 0xFF;
+
+		} while (--rc->cache_size != 0);
+
+		rc->cache = (rc->low >> 24) & 0xFF;
+	}
+
+	++rc->cache_size;
+	rc->low = (rc->low & 0x00FFFFFF) << SHIFT_BITS;
+
+	return false;
+}
+
+
+static inline bool
+rc_encode(lzma_range_encoder *rc,
+		uint8_t *out, size_t *out_pos, size_t out_size)
+{
+	while (rc->pos < rc->count) {
+		// Normalize
+		if (rc->range < TOP_VALUE) {
+			if (rc_shift_low(rc, out, out_pos, out_size))
+				return true;
+
+			rc->range <<= SHIFT_BITS;
+		}
+
+		// Encode a bit
+		switch (rc->symbols[rc->pos]) {
+		case RC_BIT_0: {
+			probability prob = *rc->probs[rc->pos];
+			rc->range = (rc->range >> BIT_MODEL_TOTAL_BITS) * prob;
+			prob += (BIT_MODEL_TOTAL - prob) >> MOVE_BITS;
+			*rc->probs[rc->pos] = prob;
+			break;
+		}
+
+		case RC_BIT_1: {
+			probability prob = *rc->probs[rc->pos];
+			const uint32_t bound = prob
+					* (rc->range >> BIT_MODEL_TOTAL_BITS);
+			rc->low += bound;
+			rc->range -= bound;
+			prob -= prob >> MOVE_BITS;
+			*rc->probs[rc->pos] = prob;
+			break;
+		}
+
+		case RC_DIRECT_0:
+			rc->range >>= 1;
+			break;
+
+		case RC_DIRECT_1:
+			rc->range >>= 1;
+			rc->low += rc->range;
+			break;
+
+		case RC_FLUSH:
+			// Prevent further normalizations.
+			rc->range = UINT32_MAX;
+
+			// Flush the last five bytes (see rc_flush()).
+			do {
+				if (rc_shift_low(rc, out, out_pos, out_size))
+					return true;
+			} while (++rc->pos < rc->count);
+
+			// Reset the range encoder so we are ready to continue
+			// encoding if we weren't finishing the stream.
+			rc_reset(rc);
+			return false;
+
+		default:
+			assert(0);
+			break;
+		}
+
+		++rc->pos;
+	}
+
+	rc->count = 0;
+	rc->pos = 0;
+
+	return false;
+}
+
+
+static inline uint64_t
+rc_pending(const lzma_range_encoder *rc)
+{
+	return rc->cache_size + 5 - 1;
+}
+
+
+////////////
+// Prices //
+////////////
+
 #ifdef HAVE_SMALL
 /// Probability prices used by *_get_price() macros. This is initialized
 /// by lzma_rc_init() and is not modified later.
@@ -49,183 +253,6 @@ lzma_rc_prob_prices[BIT_MODEL_TOTAL >> MOVE_REDUCING_BITS];
 #endif
 
 
-/// Resets the range encoder structure.
-#define rc_reset(rc) \
-do { \
-	(rc).low = 0; \
-	(rc).range = UINT32_MAX; \
-	(rc).cache_size = 1; \
-	(rc).cache = 0; \
-} while (0)
-
-
-//////////////////
-// Bit encoding //
-//////////////////
-
-// These macros expect that the following variables are defined:
-//  - lzma_range_encoder rc;
-//  - uint8_t *out;
-//  - size_t out_pos_local;  // Local copy of *out_pos
-//  - size_t size_out;
-//
-// Macros pointing to these variables are also needed:
-//  - uint8_t rc_buffer[]; // Don't use a pointer, must be real array!
-//  - size_t rc_buffer_size;
-
-
-// Combined from NRangeCoder::CEncoder::Encode()
-// and NRangeCoder::CEncoder::UpdateModel().
-#define bit_encode(prob, symbol) \
-do { \
-	probability rc_prob = prob; \
-	const uint32_t rc_bound \
-			= (rc.range >> BIT_MODEL_TOTAL_BITS) * rc_prob; \
-	if ((symbol) == 0) { \
-		rc.range = rc_bound; \
-		rc_prob += (BIT_MODEL_TOTAL - rc_prob) >> MOVE_BITS; \
-	} else { \
-		rc.low += rc_bound; \
-		rc.range -= rc_bound; \
-		rc_prob -= rc_prob >> MOVE_BITS; \
-	} \
-	prob = rc_prob; \
-	rc_normalize(); \
-} while (0)
-
-
-// Optimized version of bit_encode(prob, 0)
-#define bit_encode_0(prob) \
-do { \
-	probability rc_prob = prob; \
-	rc.range = (rc.range >> BIT_MODEL_TOTAL_BITS) * rc_prob; \
-	rc_prob += (BIT_MODEL_TOTAL - rc_prob) >> MOVE_BITS; \
-	prob = rc_prob; \
-	rc_normalize(); \
-} while (0)
-
-
-// Optimized version of bit_encode(prob, 1)
-#define bit_encode_1(prob) \
-do { \
-	probability rc_prob = prob; \
-	const uint32_t rc_bound = (rc.range >> BIT_MODEL_TOTAL_BITS) \
-			* rc_prob; \
-	rc.low += rc_bound; \
-	rc.range -= rc_bound; \
-	rc_prob -= rc_prob >> MOVE_BITS; \
-	prob = rc_prob; \
-	rc_normalize(); \
-} while (0)
-
-
-///////////////////////
-// Bit tree encoding //
-///////////////////////
-
-#define bittree_encode(probs, bit_levels, symbol) \
-do { \
-	uint32_t model_index = 1; \
-	for (int32_t bit_index = bit_levels - 1; \
-			bit_index >= 0; --bit_index) { \
-		const uint32_t bit = ((symbol) >> bit_index) & 1; \
-		bit_encode((probs)[model_index], bit); \
-		model_index = (model_index << 1) | bit; \
-	} \
-} while (0)
-
-
-#define bittree_reverse_encode(probs, bit_levels, symbol) \
-do { \
-	uint32_t model_index = 1; \
-	for (uint32_t bit_index = 0; bit_index < bit_levels; ++bit_index) { \
-		const uint32_t bit = ((symbol) >> bit_index) & 1; \
-		bit_encode((probs)[model_index], bit); \
-		model_index = (model_index << 1) | bit; \
-	} \
-} while (0)
-
-
-/////////////////
-// Direct bits //
-/////////////////
-
-#define rc_encode_direct_bits(value, num_total_bits) \
-do { \
-	for (int32_t rc_i = (num_total_bits) - 1; rc_i >= 0; --rc_i) { \
-		rc.range >>= 1; \
-		if ((((value) >> rc_i) & 1) == 1) \
-			rc.low += rc.range; \
-		rc_normalize(); \
-	} \
-} while (0)
-
-
-//////////////////
-// Buffer "I/O" //
-//////////////////
-
-// Calls rc_shift_low() to write out a byte if needed.
-#define rc_normalize() \
-do { \
-	if (rc.range < TOP_VALUE) { \
-		rc.range <<= SHIFT_BITS; \
-		rc_shift_low(); \
-	} \
-} while (0)
-
-
-// Flushes all the pending output.
-#define rc_flush() \
-	for (int32_t rc_i = 0; rc_i < 5; ++rc_i) \
-		rc_shift_low()
-
-
-// Writes the compressed data to next_out.
-// TODO: Notation change?
-//       (uint32_t)(0xFF000000)  =>  ((uint32_t)(0xFF) << TOP_BITS)
-// TODO: Another notation change?
-//       rc.low = (uint32_t)(rc.low) << SHIFT_BITS;
-//       =>
-//       rc.low &= TOP_VALUE - 1;
-//       rc.low <<= SHIFT_BITS;
-#define rc_shift_low() \
-do { \
-	if ((uint32_t)(rc.low) < (uint32_t)(0xFF000000) \
-			|| (uint32_t)(rc.low >> 32) != 0) { \
-		uint8_t rc_temp = rc.cache; \
-		do { \
-			rc_write_byte(rc_temp + (uint8_t)(rc.low >> 32)); \
-			rc_temp = 0xFF; \
-		} while(--rc.cache_size != 0); \
-		rc.cache = (uint8_t)((uint32_t)(rc.low) >> 24); \
-	} \
-	++rc.cache_size; \
-	rc.low = (uint32_t)(rc.low) << SHIFT_BITS; \
-} while (0)
-
-
-// Write one byte of compressed data to *next_out. Updates out_pos_local.
-// If out_pos_local == out_size, the byte is appended to rc_buffer.
-#define rc_write_byte(b) \
-do { \
-	if (out_pos_local == out_size) { \
-		rc_buffer[rc_buffer_size++] = (uint8_t)(b); \
-		assert(rc_buffer_size < sizeof(rc_buffer)); \
-	} else { \
-		assert(rc_buffer_size == 0); \
-		out[out_pos_local++] = (uint8_t)(b); \
-	} \
-} while (0)
-
-
-//////////////////
-// Price macros //
-//////////////////
-
-// These macros expect that the following variables are defined:
-//  - uint32_t lzma_rc_prob_prices;
-
 #define bit_get_price(prob, symbol) \
 	lzma_rc_prob_prices[((((prob) - (symbol)) ^ (-(symbol))) \
 			& (BIT_MODEL_TOTAL - 1)) >> MOVE_REDUCING_BITS]