@@ -28,20 +28,20 @@ # import "OFInflateStream.h" #else # import "OFInflate64Stream.h" # define OFInflateStream OFInflate64Stream #endif + +#import "huffman_tree.h" #import "OFInitializationFailedException.h" #import "OFInvalidFormatException.h" #import "OFNotOpenException.h" #import "OFOutOfMemoryException.h" #define BUFFER_SIZE OF_INFLATE_STREAM_BUFFER_SIZE -#define MAX_BITS 15 - enum state { BLOCK_HEADER, UNCOMPRESSED_BLOCK_HEADER, UNCOMPRESSED_BLOCK, HUFFMAN_TREE, @@ -55,15 +55,10 @@ AWAIT_DISTANCE, AWAIT_DISTANCE_EXTRA_BITS, PROCESS_PAIR }; -struct huffman_tree { - struct huffman_tree *leafs[2]; - uint16_t value; -}; - #ifndef INFLATE64 static const uint8_t numDistanceCodes = 30; static const uint8_t lengthCodes[29] = { /* indices are -257, values -3 */ 0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32, 40, 48, 56, @@ -103,20 +98,20 @@ }; #endif static const uint8_t codeLengthsOrder[19] = { 16, 17, 18, 0, 8, 7, 9, 6, 10, 5, 11, 4, 12, 3, 13, 2, 14, 1, 15 }; -static struct huffman_tree *fixedLitLenTree, *fixedDistTree; +static struct of_huffman_tree *fixedLitLenTree, *fixedDistTree; static bool tryReadBits(OFInflateStream *stream, uint16_t *bits, uint8_t count) { uint16_t ret = stream->_savedBits; assert(stream->_savedBitsLength < count); - for (uint8_t i = stream->_savedBitsLength; i < count; i++) { + for (uint_fast8_t i = stream->_savedBitsLength; i < count; i++) { if OF_UNLIKELY (stream->_bitIndex == 8) { if (stream->_bufferIndex < stream->_bufferLength) stream->_byte = stream->_buffer[stream->_bufferIndex++]; else { @@ -146,113 +141,10 @@ *bits = ret; return true; } -static struct huffman_tree * -newTree(void) -{ - struct huffman_tree *tree; - - if ((tree = malloc(sizeof(*tree))) == NULL) - @throw [OFOutOfMemoryException - exceptionWithRequestedSize: sizeof(*tree)]; - - tree->leafs[0] = tree->leafs[1] = NULL; - tree->value = 0xFFFF; - - return tree; -} - -static void -treeInsert(struct huffman_tree *tree, uint16_t code, uint8_t length, - uint16_t value) -{ - while (length > 0) { - uint8_t bit; - - length--; - bit = (code & (1 << length)) >> length; - - if (tree->leafs[bit] == NULL) - tree->leafs[bit] = newTree(); - - tree = tree->leafs[bit]; - } - - tree->value = value; -} - -static struct huffman_tree * -constructTree(uint8_t lengths[], uint16_t count) -{ - struct huffman_tree *tree; - uint16_t lengthCount[MAX_BITS + 1] = { 0 }; - uint16_t code, maxCode = 0, nextCode[MAX_BITS + 1]; - - for (uint16_t i = 0; i < count; i++) { - uint8_t length = lengths[i]; - - if OF_UNLIKELY (length > MAX_BITS) - @throw [OFInvalidFormatException exception]; - - if (length > 0) { - lengthCount[length]++; - maxCode = i; - } - } - - code = 0; - for (size_t i = 1; i <= MAX_BITS; i++) { - code = (code + lengthCount[i - 1]) << 1; - nextCode[i] = code; - } - - tree = newTree(); - - for (uint16_t i = 0; i <= maxCode; i++) { - uint8_t length = lengths[i]; - - if (length > 0) - treeInsert(tree, nextCode[length]++, length, i); - } - - return tree; -} - -static bool -walkTree(OFInflateStream *stream, struct huffman_tree **tree, uint16_t *value) -{ - struct huffman_tree *iter = *tree; - uint16_t bits; - - while (iter->value == 0xFFFF) { - if OF_UNLIKELY (!tryReadBits(stream, &bits, 1)) { - *tree = iter; - return false; - } - - if OF_UNLIKELY (iter->leafs[bits] == NULL) - @throw [OFInvalidFormatException exception]; - - iter = iter->leafs[bits]; - } - - *value = iter->value; - return true; -} - -static void -releaseTree(struct huffman_tree *tree) -{ - for (uint8_t i = 0; i < 2; i++) - if OF_LIKELY (tree->leafs[i] != NULL) - releaseTree(tree->leafs[i]); - - free(tree); -} - @implementation OFInflateStream + (void)initialize { uint8_t lengths[288]; @@ -266,16 +158,16 @@ for (uint16_t i = 256; i <= 279; i++) lengths[i] = 7; for (uint16_t i = 280; i <= 287; i++) lengths[i] = 8; - fixedLitLenTree = constructTree(lengths, 288); + fixedLitLenTree = of_huffman_tree_construct(lengths, 288); for (uint16_t i = 0; i <= 31; i++) lengths[i] = 5; - fixedDistTree = constructTree(lengths, 32); + fixedDistTree = of_huffman_tree_construct(lengths, 32); } #ifndef INFLATE64 + (instancetype)streamWithStream: (OFStream *)stream { @@ -289,19 +181,29 @@ - (instancetype)initWithStream: (OFStream *)stream { self = [super init]; - _stream = [stream retain]; + @try { + _stream = [stream retain]; - /* 0-7 address the bit, 8 means fetch next byte */ - _bitIndex = 8; + /* 0-7 address the bit, 8 means fetch next byte */ + _bitIndex = 8; + #ifdef INFLATE64 - _slidingWindowMask = 0xFFFF; + _slidingWindowMask = 0xFFFF; #else - _slidingWindowMask = 0x7FFF; + _slidingWindowMask = 0x7FFF; #endif + _slidingWindow = [self allocMemoryWithSize: + _slidingWindowMask + 1]; + /* Avoid leaking data */ + memset(_slidingWindow, 0, _slidingWindowMask + 1); + } @catch (id e) { + [self release]; + @throw e; + } return self; } - (void)dealloc @@ -308,31 +210,31 @@ { [self close]; if (_state == HUFFMAN_TREE) if (_context.huffmanTree.codeLenTree != NULL) - releaseTree(_context.huffmanTree.codeLenTree); + of_huffman_tree_release( + _context.huffmanTree.codeLenTree); if (_state == HUFFMAN_TREE || _state == HUFFMAN_BLOCK) { if (_context.huffman.litLenTree != fixedLitLenTree) - releaseTree(_context.huffman.litLenTree); + of_huffman_tree_release(_context.huffman.litLenTree); if (_context.huffman.distTree != fixedDistTree) - releaseTree(_context.huffman.distTree); + of_huffman_tree_release(_context.huffman.distTree); } [super dealloc]; } #endif - (size_t)lowlevelReadIntoBuffer: (void *)buffer_ length: (size_t)length { - uint8_t *buffer = buffer_; - uint16_t bits, tmp; - uint16_t value; + unsigned char *buffer = buffer_; + uint16_t bits, tmp, value; size_t bytesWritten = 0; - uint8_t *slidingWindow; + unsigned char *slidingWindow; uint16_t slidingWindowIndex; if (_stream == nil) @throw [OFNotOpenException exceptionWithObject: self]; @@ -423,20 +325,13 @@ ? (uint16_t)length : CTX.length - CTX.position); tmp = (uint16_t)[_stream readIntoBuffer: buffer + bytesWritten length: tmp]; - if OF_UNLIKELY (_slidingWindow == NULL) { - _slidingWindow = - [self allocMemoryWithSize: _slidingWindowMask + 1]; - /* Avoid leaking data */ - memset(_slidingWindow, 0, _slidingWindowMask + 1); - } - slidingWindow = _slidingWindow; slidingWindowIndex = _slidingWindowIndex; - for (uint16_t i = 0; i < tmp; i++) { + for (uint_fast16_t i = 0; i < tmp; i++) { slidingWindow[slidingWindowIndex] = buffer[bytesWritten + i]; slidingWindowIndex = (slidingWindowIndex + 1) & _slidingWindowMask; } @@ -482,21 +377,22 @@ if OF_LIKELY (CTX.lengths == NULL) { CTX.lengths = [self allocMemoryWithSize: 19]; memset(CTX.lengths, 0, 19); } - for (uint16_t i = CTX.receivedCount; + for (uint_fast16_t i = CTX.receivedCount; i < CTX.codeLenCodesCount + 4; i++) { if OF_UNLIKELY (!tryReadBits(self, &bits, 3)) { CTX.receivedCount = i; return bytesWritten; } CTX.lengths[codeLengthsOrder[i]] = bits; } - CTX.codeLenTree = constructTree(CTX.lengths, 19); + CTX.codeLenTree = of_huffman_tree_construct( + CTX.lengths, 19); CTX.treeIter = CTX.codeLenTree; [self freeMemory: CTX.lengths]; CTX.lengths = NULL; CTX.receivedCount = 0; @@ -505,17 +401,17 @@ if OF_LIKELY (CTX.lengths == NULL) CTX.lengths = [self allocMemoryWithSize: CTX.litLenCodesCount + CTX.distCodesCount + 258]; - for (uint16_t i = CTX.receivedCount; + for (uint_fast16_t i = CTX.receivedCount; i < CTX.litLenCodesCount + CTX.distCodesCount + 258;) { uint8_t j, count; if OF_LIKELY (CTX.value == 0xFF) { - if OF_UNLIKELY (!walkTree(self, &CTX.treeIter, - &value)) { + if OF_UNLIKELY (!of_huffman_tree_walk(self, + tryReadBits, &CTX.treeIter, &value)) { CTX.receivedCount = i; return bytesWritten; } CTX.treeIter = CTX.codeLenTree; @@ -577,16 +473,16 @@ CTX.lengths[i++] = value; CTX.value = 0xFF; } - releaseTree(CTX.codeLenTree); + of_huffman_tree_release(CTX.codeLenTree); CTX.codeLenTree = NULL; - CTX.litLenTree = constructTree(CTX.lengths, + CTX.litLenTree = of_huffman_tree_construct(CTX.lengths, CTX.litLenCodesCount + 257); - CTX.distTree = constructTree( + CTX.distTree = of_huffman_tree_construct( CTX.lengths + CTX.litLenCodesCount + 257, CTX.distCodesCount + 1); [self freeMemory: CTX.lengths]; @@ -611,19 +507,10 @@ return bytesWritten; buffer[bytesWritten++] = CTX.value; length--; - if (_slidingWindow == NULL) { - _slidingWindow = [self - allocMemoryWithSize: - _slidingWindowMask + 1]; - /* Avoid leaking data */ - memset(_slidingWindow, 0, - _slidingWindowMask + 1); - } - _slidingWindow[_slidingWindowIndex] = CTX.value; _slidingWindowIndex = (_slidingWindowIndex + 1) & _slidingWindowMask; @@ -642,12 +529,12 @@ CTX.treeIter = CTX.distTree; } /* Distance of length distance pair */ if (CTX.state == AWAIT_DISTANCE) { - if OF_UNLIKELY (!walkTree(self, &CTX.treeIter, - &value)) + if OF_UNLIKELY (!of_huffman_tree_walk(self, + tryReadBits, &CTX.treeIter, &value)) return bytesWritten; if OF_UNLIKELY (value >= numDistanceCodes) @throw [OFInvalidFormatException exception]; @@ -678,17 +565,11 @@ CTX.state = PROCESS_PAIR; } /* Length distance pair */ if (CTX.state == PROCESS_PAIR) { - uint16_t j; - - if OF_UNLIKELY (_slidingWindow == NULL) - @throw [OFInvalidFormatException - exception]; - - for (j = 0; j < CTX.length; j++) { + for (uint_fast16_t j = 0; j < CTX.length; j++) { uint16_t idx; if OF_UNLIKELY (length == 0) { CTX.length -= j; return bytesWritten; @@ -710,44 +591,36 @@ CTX.state = AWAIT_CODE; CTX.treeIter = CTX.litLenTree; } - if OF_UNLIKELY (!walkTree(self, &CTX.treeIter, &value)) + if OF_UNLIKELY (!of_huffman_tree_walk(self, tryReadBits, + &CTX.treeIter, &value)) return bytesWritten; /* End of block */ if OF_UNLIKELY (value == 256) { if (CTX.litLenTree != fixedLitLenTree) - releaseTree(CTX.litLenTree); + of_huffman_tree_release(CTX.litLenTree); if (CTX.distTree != fixedDistTree) - releaseTree(CTX.distTree); + of_huffman_tree_release(CTX.distTree); _state = BLOCK_HEADER; goto start; } /* Literal byte */ - if (value < 256) { + if OF_LIKELY (value < 256) { if OF_UNLIKELY (length == 0) { CTX.state = WRITE_VALUE; CTX.value = value; return bytesWritten; } buffer[bytesWritten++] = value; length--; - if (_slidingWindow == NULL) { - _slidingWindow = [self - allocMemoryWithSize: - _slidingWindowMask + 1]; - /* Avoid leaking data */ - memset(_slidingWindow, 0, - _slidingWindowMask + 1); - } - _slidingWindow[_slidingWindowIndex] = value; _slidingWindowIndex = (_slidingWindowIndex + 1) & _slidingWindowMask;