Merge pull request #975 from wjp/bitstream

Rework and optimize Common::BitStream

3 files changed, 273 insertions, 105 deletions

diff --git a/common/bitstream.h b/common/bitstream.h
index 325118bbee..1fbc0065ab 100644
--- a/common/bitstream.h
+++ b/common/bitstream.h
@@ -29,53 +29,10 @@
 #include "common/textconsole.h"
 #include "common/stream.h"
 #include "common/types.h"
+#include "common/util.h"
 
 namespace Common {
 
-/** A bit stream. */
-class BitStream {
-public:
-	virtual ~BitStream() {
-	}
-
-	/** Return the stream position in bits. */
-	virtual uint32 pos() const = 0;
-
-	/** Return the stream size in bits. */
-	virtual uint32 size() const = 0;
-
-	/** Has the end of the stream been reached? */
-	virtual bool eos() const = 0;
-
-	/** Rewind the bit stream back to the start. */
-	virtual void rewind() = 0;
-
-	/** Skip the specified amount of bits. */
-	virtual void skip(uint32 n) = 0;
-
-	/** Skip the bits to closest data value border. */
-	virtual void align() = 0;
-
-	/** Read a bit from the bit stream. */
-	virtual uint32 getBit() = 0;
-
-	/** Read a multi-bit value from the bit stream. */
-	virtual uint32 getBits(uint8 n) = 0;
-
-	/** Read a bit from the bit stream, without changing the stream's position. */
-	virtual uint32 peekBit() = 0;
-
-	/** Read a multi-bit value from the bit stream, without changing the stream's position. */
-	virtual uint32 peekBits(uint8 n) = 0;
-
-	/** Add a bit to the value x, making it an n+1-bit value. */
-	virtual void addBit(uint32 &x, uint32 n) = 0;
-
-protected:
-	BitStream() {
-	}
-};
-
 /**
  * A template implementing a bit stream for different data memory layouts.
  *
@@ -86,14 +43,16 @@ protected:
  * for valueBits, isLE and isMSB2LSB, reads 32bit little-endian values
  * from the data stream and hands out the bits in the order of LSB to MSB.
  */
-template<int valueBits, bool isLE, bool isMSB2LSB>
-class BitStreamImpl : public BitStream {
+template<class STREAM, int valueBits, bool isLE, bool isMSB2LSB>
+class BitStreamImpl {
 private:
-	SeekableReadStream *_stream;			///< The input stream.
+	STREAM *_stream;			///< The input stream.
 	DisposeAfterUse::Flag _disposeAfterUse; ///< Should we delete the stream on destruction?
 
 	uint32 _value;   ///< Current value.
 	uint8  _inValue; ///< Position within the current value.
+	uint32 _size;    ///< Total bitstream size (in bits)
+	uint32 _pos;     ///< Current bitstream position (in bits)
 
 	/** Read a data value. */
 	inline uint32 readData() {
@@ -119,7 +78,7 @@ private:
 
 	/** Read the next data value. */
 	inline void readValue() {
-		if ((size() - pos()) < valueBits)
+		if (_size - _pos < valueBits)
 			error("BitStreamImpl::readValue(): End of bit stream reached");
 
 		_value = readData();
@@ -133,19 +92,23 @@ private:
 
 public:
 	/** Create a bit stream using this input data stream and optionally delete it on destruction. */
-	BitStreamImpl(SeekableReadStream *stream, DisposeAfterUse::Flag disposeAfterUse = DisposeAfterUse::NO) :
-		_stream(stream), _disposeAfterUse(disposeAfterUse), _value(0), _inValue(0) {
+	BitStreamImpl(STREAM *stream, DisposeAfterUse::Flag disposeAfterUse = DisposeAfterUse::NO) :
+		_stream(stream), _disposeAfterUse(disposeAfterUse), _value(0), _inValue(0), _pos(0) {
 
 		if ((valueBits != 8) && (valueBits != 16) && (valueBits != 32))
 			error("BitStreamImpl: Invalid memory layout %d, %d, %d", valueBits, isLE, isMSB2LSB);
+
+		_size = (_stream->size() & ~((uint32) ((valueBits >> 3) - 1))) * 8;
 	}
 
 	/** Create a bit stream using this input data stream. */
-	BitStreamImpl(SeekableReadStream &stream) :
-		_stream(&stream), _disposeAfterUse(DisposeAfterUse::NO), _value(0), _inValue(0) {
+	BitStreamImpl(STREAM &stream) :
+		_stream(&stream), _disposeAfterUse(DisposeAfterUse::NO), _value(0), _inValue(0), _pos(0) {
 
 		if ((valueBits != 8) && (valueBits != 16) && (valueBits != 32))
 			error("BitStreamImpl: Invalid memory layout %d, %d, %d", valueBits, isLE, isMSB2LSB);
+
+		_size = (_stream->size() & ~((uint32) ((valueBits >> 3) - 1))) * 8;
 	}
 
 	~BitStreamImpl() {
@@ -153,14 +116,10 @@ public:
 			delete _stream;
 	}
 
-	/** Read a bit from the bit stream. */
-	uint32 getBit() {
-		// Check if we need the next value
-		if (_inValue == 0)
-			readValue();
-
+private:
+	uint32 getBit_internal() {
 		// Get the current bit
-		int b = 0;
+		uint32 b = 0;
 		if (isMSB2LSB)
 			b = ((_value & 0x80000000) == 0) ? 0 : 1;
 		else
@@ -172,8 +131,21 @@ public:
 		else
 			_value >>= 1;
 
+		return b;
+	}
+
+public:
+	/** Read a bit from the bit stream. */
+	uint32 getBit() {
+		// Check if we need the next value
+		if (_inValue == 0)
+			readValue();
+
+		uint32 b = getBit_internal();
+
 		// Increase the position within the current value
 		_inValue = (_inValue + 1) % valueBits;
+		_pos++;
 
 		return b;
 	}
@@ -198,16 +170,42 @@ public:
 		// Read the number of bits
 		uint32 v = 0;
 
-		if (isMSB2LSB) {
-			while (n-- > 0)
-				v = (v << 1) | getBit();
-		} else {
-			for (uint32 i = 0; i < n; i++)
-				v = (v >> 1) | (((uint32) getBit()) << 31);
+		uint8 nOrig = n;
+		if (_inValue) {
+			int count = MIN((int)n, valueBits - _inValue);
+			for (int i = 0; i < count; ++i) {
+				if (isMSB2LSB) {
+					v = (v << 1) | getBit_internal();
+				} else {
+					v = (v >> 1) | (getBit_internal() << 31);
+				}
+			}
+
+			n -= count;
+		}
+
+		while (n > 0) {
+			// NB: readValue doesn't care that _inValue is incorrect here
+			readValue();
 
-			v >>= (32 - n);
+			int count = MIN((int)n, valueBits);
+			for (int i = 0; i < count; ++i) {
+				if (isMSB2LSB) {
+					v = (v << 1) | getBit_internal();
+				} else {
+					v = (v >> 1) | (getBit_internal() << 31);
+				}
+			}
+
+			n -= count;
 		}
 
+		_inValue = (_inValue + nOrig) % valueBits;
+		_pos += nOrig;
+
+		if (!isMSB2LSB)
+			v >>= (32 - nOrig);
+
 		return v;
 	}
 
@@ -215,11 +213,13 @@ public:
 	uint32 peekBit() {
 		uint32 value   = _value;
 		uint8  inValue = _inValue;
-		uint32 curPos  = _stream->pos();
+		uint32 curStreamPos  = _stream->pos();
+		uint32 curPos = _pos;
 
 		uint32 v = getBit();
 
-		_stream->seek(curPos);
+		_pos     = curPos;
+		_stream->seek(curStreamPos);
 		_inValue = inValue;
 		_value   = value;
 
@@ -234,11 +234,13 @@ public:
 	uint32 peekBits(uint8 n) {
 		uint32 value   = _value;
 		uint8  inValue = _inValue;
-		uint32 curPos  = _stream->pos();
+		uint32 curStreamPos  = _stream->pos();
+		uint32 curPos = _pos;
 
 		uint32 v = getBits(n);
 
-		_stream->seek(curPos);
+		_pos     = curPos;
+		_stream->seek(curStreamPos);
 		_inValue = inValue;
 		_value   = value;
 
@@ -272,6 +274,7 @@ public:
 
 		_value   = 0;
 		_inValue = 0;
+		_pos     = 0;
 	}
 
 	/** Skip the specified amount of bits. */
@@ -288,47 +291,215 @@ public:
 
 	/** Return the stream position in bits. */
 	uint32 pos() const {
-		if (_stream->pos() == 0)
-			return 0;
-
-		uint32 p = (_inValue == 0) ? _stream->pos() : ((_stream->pos() - 1) & ~((uint32) ((valueBits >> 3) - 1)));
-		return p * 8 + _inValue;
+		return _pos;
 	}
 
 	/** Return the stream size in bits. */
 	uint32 size() const {
-		return (_stream->size() & ~((uint32) ((valueBits >> 3) - 1))) * 8;
+		return _size;
 	}
 
 	bool eos() const {
-		return _stream->eos() || (pos() >= size());
+		return _stream->eos() || (_pos >= _size);
 	}
 };
 
+
+
+/**
+ * A cut-down version of MemoryReadStream specifically for use with BitStream.
+ * It removes the virtual call overhead for reading bytes from a memory buffer,
+ * and allows directly inlining this access.
+ *
+ * The code duplication with MemoryReadStream is not ideal.
+ * It might be possible to avoid this by making this a final subclass of
+ * MemoryReadStream, but that is a C++11 feature.
+ */
+class BitStreamMemoryStream {
+private:
+	const byte * const _ptrOrig;
+	const byte *_ptr;
+	const uint32 _size;
+	uint32 _pos;
+	DisposeAfterUse::Flag _disposeMemory;
+	bool _eos;
+
+public:
+	BitStreamMemoryStream(const byte *dataPtr, uint32 dataSize, DisposeAfterUse::Flag disposeMemory = DisposeAfterUse::NO) :
+		_ptrOrig(dataPtr),
+		_ptr(dataPtr),
+		_size(dataSize),
+		_pos(0),
+		_disposeMemory(disposeMemory),
+		_eos(false) {}
+
+	~BitStreamMemoryStream() {
+		if (_disposeMemory)
+			free(const_cast<byte *>(_ptrOrig));
+	}
+
+	bool eos() const {
+		return _eos;
+	}
+
+	bool err() const {
+		return false;
+	}
+
+	int32 pos() const {
+		return _pos;
+	}
+
+	int32 size() const {
+		return _size;
+	}
+
+	bool seek(uint32 offset) {
+		assert(offset <= _size);
+
+		_eos = false;
+		_pos = offset;
+		_ptr = _ptrOrig + _pos;
+		return true;
+	}
+
+	byte readByte() {
+		if (_pos >= _size) {
+			_eos = true;
+			return 0;
+		}
+
+		_pos++;
+		return *_ptr++;
+	}
+
+	uint16 readUint16LE() {
+		if (_pos + 2 > _size) {
+			_eos = true;
+			if (_pos < _size) {
+				_pos++;
+				return *_ptr++;
+			} else {
+				return 0;
+			}
+		}
+
+		uint16 val = READ_LE_UINT16(_ptr);
+
+		_pos += 2;
+		_ptr += 2;
+
+		return val;
+	}
+
+	uint16 readUint16BE() {
+		if (_pos + 2 > _size) {
+			_eos = true;
+			if (_pos < _size) {
+				_pos++;
+				return (*_ptr++) << 8;
+			} else {
+				return 0;
+			}
+		}
+
+		uint16 val = READ_LE_UINT16(_ptr);
+
+		_pos += 2;
+		_ptr += 2;
+
+		return val;
+	}
+
+	uint32 readUint32LE() {
+		if (_pos + 4 > _size) {
+			uint32 val = readByte();
+			val |= (uint32)readByte() << 8;
+			val |= (uint32)readByte() << 16;
+			val |= (uint32)readByte() << 24;
+
+			return val;
+		}
+
+		uint32 val = READ_LE_UINT32(_ptr);
+
+		_pos += 4;
+		_ptr += 4;
+
+		return val;
+	}
+
+	uint32 readUint32BE() {
+		if (_pos + 4 > _size) {
+			uint32 val = (uint32)readByte() << 24;
+			val |= (uint32)readByte() << 16;
+			val |= (uint32)readByte() << 8;
+			val |= (uint32)readByte();
+
+			return val;
+		}
+
+		uint32 val = READ_BE_UINT32(_ptr);
+
+		_pos += 4;
+		_ptr += 4;
+
+		return val;
+	}
+
+};
+
+
 // typedefs for various memory layouts.
 
 /** 8-bit data, MSB to LSB. */
-typedef BitStreamImpl<8, false, true > BitStream8MSB;
+typedef BitStreamImpl<SeekableReadStream, 8, false, true > BitStream8MSB;
 /** 8-bit data, LSB to MSB. */
-typedef BitStreamImpl<8, false, false> BitStream8LSB;
+typedef BitStreamImpl<SeekableReadStream, 8, false, false> BitStream8LSB;
 
 /** 16-bit little-endian data, MSB to LSB. */
-typedef BitStreamImpl<16, true , true > BitStream16LEMSB;
+typedef BitStreamImpl<SeekableReadStream, 16, true , true > BitStream16LEMSB;
 /** 16-bit little-endian data, LSB to MSB. */
-typedef BitStreamImpl<16, true , false> BitStream16LELSB;
+typedef BitStreamImpl<SeekableReadStream, 16, true , false> BitStream16LELSB;
 /** 16-bit big-endian data, MSB to LSB. */
-typedef BitStreamImpl<16, false, true > BitStream16BEMSB;
+typedef BitStreamImpl<SeekableReadStream, 16, false, true > BitStream16BEMSB;
 /** 16-bit big-endian data, LSB to MSB. */
-typedef BitStreamImpl<16, false, false> BitStream16BELSB;
+typedef BitStreamImpl<SeekableReadStream, 16, false, false> BitStream16BELSB;
 
 /** 32-bit little-endian data, MSB to LSB. */
-typedef BitStreamImpl<32, true , true > BitStream32LEMSB;
+typedef BitStreamImpl<SeekableReadStream, 32, true , true > BitStream32LEMSB;
 /** 32-bit little-endian data, LSB to MSB. */
-typedef BitStreamImpl<32, true , false> BitStream32LELSB;
+typedef BitStreamImpl<SeekableReadStream, 32, true , false> BitStream32LELSB;
 /** 32-bit big-endian data, MSB to LSB. */
-typedef BitStreamImpl<32, false, true > BitStream32BEMSB;
+typedef BitStreamImpl<SeekableReadStream, 32, false, true > BitStream32BEMSB;
 /** 32-bit big-endian data, LSB to MSB. */
-typedef BitStreamImpl<32, false, false> BitStream32BELSB;
+typedef BitStreamImpl<SeekableReadStream, 32, false, false> BitStream32BELSB;
+
+
+
+/** 8-bit data, MSB to LSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 8, false, true > BitStreamMemory8MSB;
+/** 8-bit data, LSB to MSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 8, false, false> BitStreamMemory8LSB;
+
+/** 16-bit little-endian data, MSB to LSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 16, true , true > BitStreamMemory16LEMSB;
+/** 16-bit little-endian data, LSB to MSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 16, true , false> BitStreamMemory16LELSB;
+/** 16-bit big-endian data, MSB to LSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 16, false, true > BitStreamMemory16BEMSB;
+/** 16-bit big-endian data, LSB to MSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 16, false, false> BitStreamMemory16BELSB;
+
+/** 32-bit little-endian data, MSB to LSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 32, true , true > BitStreamMemory32LEMSB;
+/** 32-bit little-endian data, LSB to MSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 32, true , false> BitStreamMemory32LELSB;
+/** 32-bit big-endian data, MSB to LSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 32, false, true > BitStreamMemory32BEMSB;
+/** 32-bit big-endian data, LSB to MSB. */
+typedef BitStreamImpl<BitStreamMemoryStream, 32, false, false> BitStreamMemory32BELSB;
+
 
 } // End of namespace Common
 
diff --git a/common/huffman.cpp b/common/huffman.cpp
index afb4fa00b6..3268ddf251 100644
--- a/common/huffman.cpp
+++ b/common/huffman.cpp
@@ -68,19 +68,4 @@ void Huffman::setSymbols(const uint32 *symbols) {
 		_symbols[i]->symbol = symbols ? *symbols++ : i;
 }
 
-uint32 Huffman::getSymbol(BitStream &bits) const {
-	uint32 code = 0;
-
-	for (uint32 i = 0; i < _codes.size(); i++) {
-		bits.addBit(code, i);
-
-		for (CodeList::const_iterator cCode = _codes[i].begin(); cCode != _codes[i].end(); ++cCode)
-			if (code == cCode->code)
-				return cCode->symbol;
-	}
-
-	error("Unknown Huffman code");
-	return 0;
-}
-
 } // End of namespace Common
diff --git a/common/huffman.h b/common/huffman.h
index 5ba7deca87..f703d078dc 100644
--- a/common/huffman.h
+++ b/common/huffman.h
@@ -31,8 +31,6 @@
 
 namespace Common {
 
-class BitStream;
-
 /**
  * Huffman bitstream decoding
  *
@@ -56,7 +54,21 @@ public:
 	void setSymbols(const uint32 *symbols = 0);
 
 	/** Return the next symbol in the bitstream. */
-	uint32 getSymbol(BitStream &bits) const;
+	template<class BITSTREAM>
+	uint32 getSymbol(BITSTREAM &bits) const {
+		uint32 code = 0;
+
+		for (uint32 i = 0; i < _codes.size(); i++) {
+			bits.addBit(code, i);
+
+			for (CodeList::const_iterator cCode = _codes[i].begin(); cCode != _codes[i].end(); ++cCode)
+				if (code == cCode->code)
+					return cCode->symbol;
+		}
+
+		error("Unknown Huffman code");
+		return 0;
+	}
 
 private:
 	struct Symbol {