X-Git-Url: https://wimlib.net/git/?p=wimlib;a=blobdiff_plain;f=src%2Fcompress.c;h=3a936347df6de044ca48d68a2ef10ac04e740d90;hp=f98a5a524854c112eb35906d22cc3f0b779cdf70;hb=60523d25f34692d6f3a7c8bbda88eead17f23b12;hpb=4757f17833c96b8c83a7e17cbc6f374c449d60db

diff --git a/src/compress.c b/src/compress.c
index f98a5a52..3a936347 100644
--- a/src/compress.c
+++ b/src/compress.c
@@ -28,6 +28,7 @@
 #endif
 
 #include "wimlib/assert.h"
+#include "wimlib/endianness.h"
 #include "wimlib/compiler.h"
 #include "wimlib/compress.h"
 #include "wimlib/util.h"
@@ -38,45 +39,44 @@
 /* Writes @num_bits bits, given by the @num_bits least significant bits of
  * @bits, to the output @ostream. */
 void
-bitstream_put_bits(struct output_bitstream *ostream, output_bitbuf_t bits,
+bitstream_put_bits(struct output_bitstream *ostream, u32 bits,
 		   unsigned num_bits)
 {
-	unsigned rem_bits;
+	bits &= (1U << num_bits) - 1;
+	while (num_bits > ostream->free_bits) {
+		/* Buffer variable does not have space for the new bits.  It
+		 * needs to be flushed as a 16-bit integer.  Bits in the second
+		 * byte logically precede those in the first byte
+		 * (little-endian), but within each byte the bits are ordered
+		 * from high to low.  This is true for both XPRESS and LZX
+		 * compression.  */
+
+		/* There must be at least 2 bytes of space remaining.  */
+		if (unlikely(ostream->bytes_remaining < 2)) {
+			ostream->overrun = true;
+			return;
+		}
 
-	if (num_bits <= ostream->free_bits) {
-		/* Buffer variable had space for the new bits.  */
-		ostream->bitbuf = (ostream->bitbuf << num_bits) | bits;
-		ostream->free_bits -= num_bits;
-		return;
-	}
+		/* Fill the buffer with as many bits that fit.  */
+		unsigned fill_bits = ostream->free_bits;
 
-	/* Buffer variable does not have space for the new bits.  It needs to be
-	 * flushed as a 16-bit integer.  Bits in the second byte logically
-	 * precede those in the first byte (little-endian), but within each byte
-	 * the bits are ordered from high to low.  This is true for both XPRESS
-	 * and LZX compression.  */
+		ostream->bitbuf <<= fill_bits;
+		ostream->bitbuf |= bits >> (num_bits - fill_bits);
 
-	wimlib_assert(num_bits <= 16);
+		*(le16*)ostream->bit_output = cpu_to_le16(ostream->bitbuf);
+		ostream->bit_output = ostream->next_bit_output;
+		ostream->next_bit_output = ostream->output;
+		ostream->output += 2;
+		ostream->bytes_remaining -= 2;
 
-	/* There must be at least 2 bytes of space remaining.  */
-	if (unlikely(ostream->bytes_remaining < 2)) {
-		ostream->overrun = true;
-		return;
+		ostream->free_bits = 16;
+		num_bits -= fill_bits;
+		bits &= (1U << num_bits) - 1;
 	}
 
-	/* Fill the buffer with as many bits that fit.  */
-	rem_bits = num_bits - ostream->free_bits;
-	ostream->bitbuf <<= ostream->free_bits;
-	ostream->bitbuf |= bits >> rem_bits;
-
-	*(le16*)ostream->bit_output = cpu_to_le16(ostream->bitbuf);
-	ostream->bit_output = ostream->next_bit_output;
-	ostream->next_bit_output = ostream->output;
-	ostream->output += 2;
-	ostream->bytes_remaining -= 2;
-
-	ostream->free_bits = 16 - rem_bits;
-	ostream->bitbuf = bits;
+	/* Buffer variable has space for the new bits.  */
+	ostream->bitbuf = (ostream->bitbuf << num_bits) | bits;
+	ostream->free_bits -= num_bits;
 }
 
 void
@@ -127,7 +127,7 @@ init_output_bitstream(struct output_bitstream *ostream,
 }
 
 typedef struct {
-	u32 freq;
+	input_idx_t freq;
 	u16 sym;
 	union {
 		u16 path_len;
@@ -150,12 +150,12 @@ cmp_nodes_by_freq(const void *_leaf1, const void *_leaf2)
 	const HuffmanNode *leaf1 = _leaf1;
 	const HuffmanNode *leaf2 = _leaf2;
 
-	int freq_diff = (int)leaf1->freq - (int)leaf2->freq;
-
-	if (freq_diff == 0)
-		return (int)leaf1->sym - (int)leaf2->sym;
+	if (leaf1->freq > leaf2->freq)
+		return 1;
+	else if (leaf1->freq < leaf2->freq)
+		return -1;
 	else
-		return freq_diff;
+		return (int)leaf1->sym - (int)leaf2->sym;
 }
 
 /* Comparator function for HuffmanNodes.  Sorts primarily by code length and
@@ -251,14 +251,14 @@ huffman_tree_compute_path_lengths(HuffmanNode *base_node, u16 cur_len)
 void
 make_canonical_huffman_code(unsigned num_syms,
 			    unsigned max_codeword_len,
-			    const freq_t freq_tab[restrict],
+			    const input_idx_t freq_tab[restrict],
 			    u8 lens[restrict],
 			    u16 codewords[restrict])
 {
 	/* We require at least 2 possible symbols in the alphabet to produce a
 	 * valid Huffman decoding table. It is allowed that fewer than 2 symbols
 	 * are actually used, though. */
-	wimlib_assert(num_syms >= 2);
+	wimlib_assert(num_syms >= 2 && num_syms < INVALID_SYMBOL);
 
 	/* Initialize the lengths and codewords to 0 */
 	memset(lens, 0, num_syms * sizeof(lens[0]));