Add optional padding data at the end to the bit source (is that the right place?...

diff --git a/bitsource.c b/bitsource.c
index 6cc627c70b08b462b0d0c9c7e0a4ec3d25cdceed..7a4ea4a0ca3c553e276be095932061313e833441 100644 (file)
--- a/bitsource.c
+++ b/bitsource.c
@@ -8,12 +8,14 @@
 #define MARKER_CHAR 0xff
 #define STUFF_MARKER 0x00
 
 #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;
 {
        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;
 }
 
        source->userdata = userdata;
 }
 


@@ -43,12 +45,20 @@ void possibly_refill_slow_path(struct bit_source* source, unsigned num_bits)
                        fprintf(stderr, "Input function returned error\n");
                        exit(1);
                }
                        fprintf(stderr, "Input function returned error\n");
                        exit(1);
                }

-               if (bytes_read == 0) {
+               if (bytes_read == 0 && 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 if (bytes_read == 0) {

                        fprintf(stderr, "Premature EOF\n");
                        exit(1);
                        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.
        }
 
        // Fill the bit reservoir one by one byte until we have enough.

diff --git a/bitsource.h b/bitsource.h
index bb3a6d246e62016665e5560ac048397b15f7f8e9..ae29b4c4a2200abce210242a541be2c289ea8163 100644 (file)
--- a/bitsource.h
+++ b/bitsource.h
@@ -37,12 +37,18 @@ struct bit_source {
        uint8_t* byte_read_ptr;
        unsigned bytes_available;
 
        uint8_t* byte_read_ptr;
        unsigned bytes_available;
 

+       // Some clients will purposedly read a bit ahead of the stream, causing
+       // problems at EOF. Thus, the client is allowed to request that we pad
+       // the end stream with a few bytes after the source reports EOF.
+       int padding_bytes_available;
+

        // Data source.
        input_func_t* input_func;
        void* userdata;
 };
 
        // Data source.
        input_func_t* input_func;
        void* userdata;
 };
 

-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);

 
 // Internal function. Do not use.
 void possibly_refill_slow_path(struct bit_source* source, unsigned num_bits);
 
 // Internal function. Do not use.
 void possibly_refill_slow_path(struct bit_source* source, unsigned num_bits);

diff --git a/bitsource_test.c b/bitsource_test.c
index 2b5217ff91a283fee9c544430bf73271cb565cfc..173e499b169c469414ce6303ccbeced15ca07c2e 100644 (file)
--- a/bitsource_test.c
+++ b/bitsource_test.c
@@ -40,7 +40,7 @@ void test_basic_reading()
        ud.bytes_left = sizeof(bytes);
 
        struct bit_source source;
        ud.bytes_left = sizeof(bytes);
 
        struct bit_source source;

-       init_bit_source(&source, custom_read, &ud);
+       init_bit_source(&source, custom_read, 0, &ud);

 
        for (unsigned i = 0; i < sizeof(bytes) * 8 / 6; ++i) {
                possibly_refill(&source, 6);
 
        for (unsigned i = 0; i < sizeof(bytes) * 8 / 6; ++i) {
                possibly_refill(&source, 6);


@@ -60,7 +60,7 @@ void test_slow_source()
        ud.bytes_left = sizeof(bytes);
 
        struct bit_source source;
        ud.bytes_left = sizeof(bytes);
 
        struct bit_source source;

-       init_bit_source(&source, custom_read_slow, &ud);
+       init_bit_source(&source, custom_read_slow, 0, &ud);

 
        for (unsigned i = 0; i < sizeof(bytes) * 8 / 6; ++i) {
                possibly_refill(&source, 6);
 
        for (unsigned i = 0; i < sizeof(bytes) * 8 / 6; ++i) {
                possibly_refill(&source, 6);


@@ -80,7 +80,7 @@ void test_variable_size()
        ud.bytes_left = sizeof(bytes);
 
        struct bit_source source;
        ud.bytes_left = sizeof(bytes);
 
        struct bit_source source;

-       init_bit_source(&source, custom_read, &ud);
+       init_bit_source(&source, custom_read, 0, &ud);

 
        {
                possibly_refill(&source, 4);
 
        {
                possibly_refill(&source, 4);


@@ -111,6 +111,32 @@ void test_variable_size()
                assert(ret == 0xf);
        }
 
                assert(ret == 0xf);
        }
 

+       assert(ud.bytes_left == 0);
+} 
+
+// Tests that padding bytes are properly added at the end, even crossing reads.
+void test_padding_reads()
+{
+       uint8_t bytes[] = { 0x01, 0x02 };
+
+       struct custom_read_userdata ud;
+       ud.bytes = bytes;
+       ud.bytes_left = sizeof(bytes);
+
+       struct bit_source source;
+       init_bit_source(&source, custom_read, 2, &ud);
+       
+       {
+               possibly_refill(&source, 8);
+               unsigned ret = read_bits(&source, 8);
+               assert(ret == 0x01);
+       }
+       {
+               possibly_refill(&source, 12);
+               unsigned ret = read_bits(&source, 12);
+               assert(ret == 0x020);
+       }
+

        assert(ud.bytes_left == 0);
 }
 
        assert(ud.bytes_left == 0);
 }
 


@@ -125,6 +151,9 @@ int main(void)
        printf("test_variable_size()\n");
        test_variable_size();
        
        printf("test_variable_size()\n");
        test_variable_size();
        

+       printf("test_padding_reads()\n");
+       test_padding_reads();
+       

        printf("All tests pass.\n");
        return 0;
 }
        printf("All tests pass.\n");
        return 0;
 }

diff --git a/dehuff.h b/dehuff.h
index 224005f84ad3dcb2797c9086a11485c8269c825c..a6a431a8139440a7052ce1dadd39b182ee5566b4 100644 (file)
--- a/dehuff.h
+++ b/dehuff.h
@@ -51,9 +51,8 @@ unsigned read_huffman_symbol_slow_path(const struct huffman_table* table,
 static inline unsigned read_huffman_symbol(const struct huffman_table* table,
                                            struct bit_source* source)
 {
 static inline unsigned read_huffman_symbol(const struct huffman_table* table,
                                            struct bit_source* source)
 {

-       // FIXME: We can read past the end of the stream here in some edge
-       // cases. We need to define some guarantees in the layers above.

        possibly_refill(source, DEHUF_TABLE_BITS);
        possibly_refill(source, DEHUF_TABLE_BITS);

+       assert(source->bits_available >= DEHUF_TABLE_BITS);

        unsigned lookup = peek_bits(source, DEHUF_TABLE_BITS);
        int code = table->lookup_table_codes[lookup];
        int length = table->lookup_table_length[lookup];
        unsigned lookup = peek_bits(source, DEHUF_TABLE_BITS);
        int code = table->lookup_table_codes[lookup];
        int length = table->lookup_table_length[lookup];


@@ -61,7 +60,7 @@ static inline unsigned read_huffman_symbol(const struct huffman_table* table,
        if (code == DEHUF_SLOW_PATH) {
                return read_huffman_symbol_slow_path(table, source);
        }
        if (code == DEHUF_SLOW_PATH) {
                return read_huffman_symbol_slow_path(table, source);
        }

-
+               

        read_bits(source, length);
        return code;
 }
        read_bits(source, length);
        return code;
 }

diff --git a/dehuff_test.c b/dehuff_test.c
index 0212244bdf4594be6cf4b54280a5eb80003d0a45..eed09520dab288214cf9a57236d4a24197045902 100644 (file)
--- a/dehuff_test.c
+++ b/dehuff_test.c
@@ -78,7 +78,7 @@ void test_table_gen()
 
 // Test that we can decode a simple bit stream.
 // Note that since we end on a long code, we won't crash into
 
 // Test that we can decode a simple bit stream.
 // Note that since we end on a long code, we won't crash into

-// the end-of-stream problems we currently have.
+// the end-of-stream problems we have without padding.

 void test_decoding()
 {
        huffman_tables_t tables;
 void test_decoding()
 {
        huffman_tables_t tables;


@@ -97,7 +97,7 @@ void test_decoding()
        ud.bytes_left = sizeof(bytes);
 
        struct bit_source source;
        ud.bytes_left = sizeof(bytes);
 
        struct bit_source source;

-       init_bit_source(&source, custom_read, &ud);
+       init_bit_source(&source, custom_read, 0, &ud);

                
        struct huffman_table* tbl = &tables[DC_CLASS][0];       
        for (unsigned i = 0; i < 12; ++i) {
                
        struct huffman_table* tbl = &tables[DC_CLASS][0];       
        for (unsigned i = 0; i < 12; ++i) {


@@ -108,6 +108,35 @@ void test_decoding()
        assert(source.bits_available == 0);
 }
 
        assert(source.bits_available == 0);
 }
 

+// Test that we can decode a bit stream that ends on a short code,
+// if we've got padding.
+void test_padded_decoding()
+{
+       huffman_tables_t tables;
+       read_example_tables(&tables);
+       
+       // Our stream looks like this:
+       //
+       //  0   1   2   3   4
+       // 00 010 011 100 101
+       uint8_t bytes[] = {
+               0x13, 0x94
+       };
+
+       struct custom_read_userdata ud;
+       ud.bytes = bytes;
+       ud.bytes_left = sizeof(bytes);
+
+       struct bit_source source;
+       init_bit_source(&source, custom_read, 2, &ud);
+               
+       struct huffman_table* tbl = &tables[DC_CLASS][0];       
+       for (unsigned i = 0; i < 5; ++i) {
+               unsigned symbol = read_huffman_symbol(tbl, &source);
+               assert(symbol == i);
+       }
+}
+

 int main(void)
 {
        printf("test_table_gen()\n");
 int main(void)
 {
        printf("test_table_gen()\n");


@@ -116,6 +145,9 @@ int main(void)
        printf("test_decoding()\n");
        test_decoding();
        
        printf("test_decoding()\n");
        test_decoding();
        

+       printf("test_padded_decoding()\n");
+       test_padded_decoding();
+       

        printf("All tests pass.\n");
        return 0;
 }
        printf("All tests pass.\n");
        return 0;
 }

diff --git a/driver.c b/driver.c
index 06d94d99d553e6aa4e784e22cfc99a4d19216df5..aa803d6ed00aa67266495b65df93e6582bc89151 100644 (file)
--- a/driver.c
+++ b/driver.c
@@ -147,7 +147,7 @@ void read_scan(struct byte_source* source, struct jpeg_image* image, huffman_tab
        }
 
        struct bit_source bits;
        }
 
        struct bit_source bits;

-       init_bit_source(&bits, byte_source_input_func, source);
+       init_bit_source(&bits, byte_source_input_func, 8, source);

                
        unsigned mcu_x = 0, mcu_y = 0;
 
                
        unsigned mcu_x = 0, mcu_y = 0;