Added functions for bit-level reading.

[fjl] / input.h

diff --git a/input.h b/input.h
new file mode 100644 (file)
index 0000000..28c5f0b
--- /dev/null
+++ b/input.h
@@ -0,0 +1,88 @@
+#ifndef _INPUT_H
+#define _INPUT_H 1
+
+#include <assert.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <sys/types.h>
+#include <arpa/inet.h>
+
+// Optimize for 64 bits. We might want to replace this for 32-bit machines
+// (benchmark later).
+typedef uint64_t bitreservoir_t;
+typedef uint32_t bitreservoir_fill_t;
+
+static inline bitreservoir_fill_t read_bitreservoir_fill(uint8_t* source)
+{
+       return ntohl(*(bitreservoir_fill_t*)(source));
+}
+
+static const unsigned bitreservoir_size = 8 * sizeof(bitreservoir_t);
+static const unsigned bitreservoir_fill_size = 8 * sizeof(bitreservoir_fill_t);
+static const unsigned bytereservoir_size = 4096;
+
+// A function to read bytes from some input source.
+// A return value of -1 indicates error, a return value of 0 indicates EOF.
+typedef ssize_t (input_func_t)(void*, uint8_t*, size_t);
+
+struct data_source {
+       // Short-term bit reservoir; holds up to 64 bits. When it's empty,
+       // it needs to get refilled from the medium-term bit reservoir.
+       bitreservoir_t bits;
+       unsigned bits_available;
+       
+       // Medium-term bit reservoir; holds a few kilobytes of spare data.
+       // When this is empty, it needs to be refilled from the input
+       // stream.
+       uint8_t* bytes;
+       unsigned bytes_available;
+
+       // Data source.
+       input_func_t* input_func;
+       void* userdata;
+};
+       
+void init_data_source(struct data_source* source, input_func_t* input_func, void* userdata);
+
+// Internal function. Do not use.
+void possibly_refill_slow_path(struct data_source* source, unsigned num_bits);
+
+// Make sure there's at least NUM_BITS available in the short-term bit reservoir.
+// You usually want to call this before read_bits(). The reason it's separate
+// is that if you want two reads and you know the size of both, it's faster to
+// refill A+B, read A, read B than refill A, read A, refill B, read B.
+static inline void possibly_refill(struct data_source* source, unsigned num_bits)
+{
+       assert(num_bits <= bitreservoir_fill_size + 1);
+
+       if (source->bits_available >= num_bits) {
+               // Fast path (~90% of invocations?)
+               return;
+       }
+
+       // Slower path (~99% of remaining invocations?)
+       assert(source->bits_available + bitreservoir_fill_size < bitreservoir_size);
+       if (source->bytes_available >= sizeof(bitreservoir_fill_t)) {
+               bitreservoir_fill_t fill = read_bitreservoir_fill(source->bytes);
+               source->bytes += sizeof(bitreservoir_fill_t);
+               source->bytes_available -= sizeof(bitreservoir_fill_t);
+               source->bits |= (bitreservoir_t)fill << (bitreservoir_size - bitreservoir_fill_size - source->bits_available);
+               source->bits_available += bitreservoir_fill_size;
+               return;
+       }
+
+       // Slow path: Refill from data source.
+       // Should not be inlined, so split into a separate function.
+       possibly_refill_slow_path(source, num_bits);
+}
+
+static inline unsigned read_bits(struct data_source* source, unsigned num_bits)
+{
+       assert(source->bits_available >= num_bits);
+       unsigned ret = (source->bits >> (bitreservoir_size - num_bits));
+       source->bits <<= num_bits;
+       source->bits_available -= num_bits;
+       return ret;
+}
+
+#endif /* !defined(_INPUT_H) */