X-Git-Url: https://git.sesse.net/?p=fjl;a=blobdiff_plain;f=bitsource.c;h=b1593aa6ec554959b4bba6bbe3aceb75d4e04863;hp=72b143bf7092b666669b3eafb19baf1b7d3bd670;hb=HEAD;hpb=0b02847989970a190c2cfaec4d1abaa1f616284a

diff --git a/bitsource.c b/bitsource.c
index 72b143b..b1593aa 100644
--- a/bitsource.c
+++ b/bitsource.c
@@ -8,19 +8,26 @@
 #define MARKER_CHAR 0xff
 #define STUFF_MARKER 0x00
 
-void init_bit_source(struct bit_source* source, input_func_t* input_func, void* userdata)
+void init_bit_source(struct bit_source* source, input_func_t* input_func,
+                     unsigned padding_bytes, void* userdata)
 {
 	memset(source, 0, sizeof(*source));
 	source->bytes = (uint8_t*)malloc(BYTERESERVOIR_SIZE);
+	source->byte_read_ptr = source->bytes;
 	source->input_func = input_func;
+	source->padding_bytes_available = padding_bytes;
 	source->userdata = userdata;
 }
 
 void possibly_refill_slow_path(struct bit_source* source, unsigned num_bits)
 {
-	// First, make sure there's stuff in the byte reservoir if we can.
+	// First, move out the data we already read.
 	assert(source->bytes_available <= BYTERESERVOIR_SIZE);
+	assert(source->byte_read_ptr >= source->bytes);
+	memmove(source->bytes, source->byte_read_ptr, source->bytes_available);
+	source->byte_read_ptr = source->bytes;
 
+	// Then, make sure there's stuff in the byte reservoir if we can.
 	// Read data from the source until we have enough to satisfy the request.
 	while (source->bits_available + 8 * source->bytes_available < num_bits) {
 		const size_t bytes_to_read = BYTERESERVOIR_SIZE - source->bytes_available;
@@ -39,19 +46,30 @@ void possibly_refill_slow_path(struct bit_source* source, unsigned num_bits)
 			exit(1);
 		}
 		if (bytes_read == 0) {
-			fprintf(stderr, "Premature EOF\n");
-			exit(1);
+			source->source_eof = true;
+			if (source->padding_bytes_available > 0) {
+				unsigned padding_to_add = source->padding_bytes_available;
+				if (padding_to_add > bytes_to_read) {
+					padding_to_add = bytes_to_read;
+				}
+				memset(source->bytes + source->bytes_available, 0, padding_to_add);
+				source->padding_bytes_available -= padding_to_add;
+				source->bytes_available += padding_to_add;
+			} else {
+				fprintf(stderr, "Premature EOF\n");
+				exit(1);
+			}
+		} else {
+			source->bytes_available += bytes_read;
 		}
-		
-		source->bytes_available += bytes_read;
 	}
 
 	// Fill the bit reservoir one by one byte until we have enough.
 	while (source->bits_available < num_bits) {
 		assert(source->bytes_available > 0);
 		assert(source->bits_available + 8 <= BITRESERVOIR_SIZE);
-		uint8_t byte = *(source->bytes);
-		++source->bytes;
+		uint8_t byte = *(source->byte_read_ptr);
+		++source->byte_read_ptr;
 		--source->bytes_available;
 		source->bits |= ((bitreservoir_t)byte << (BITRESERVOIR_SIZE - source->bits_available - 8));
 		source->bits_available += 8;