ObjFW  Diff

#import "OFKernelEventObserver.h"

#import "OFHuffmanTree.h"

#import "OFInvalidFormatException.h"
#import "OFNotOpenException.h"

enum state {
	StateBlockHeader,
	StateCodeLenCodesCount,
	StateCodeLenTree,
	StateCodeLenTreeSingle,
	StateLitLenCodesCount,
	StateLitLenTree,
	StateLitLenTreeSingle,
	StateDistCodesCount,
	StateDistTree,
	StateDistTreeSingle,
	StateBlockLitLen,
	StateBlockDistLength,
	StateBlockDistLengthExtra,
	StateBlockLenDistPair
};

@implementation OFLHADecompressingStream
@synthesize bytesConsumed = _bytesConsumed;

static OF_INLINE bool
tryReadBits(OFLHADecompressingStream *stream, uint16_t *bits, uint8_t count)

	uint16_t bits = 0, value = 0;
	size_t bytesWritten = 0;

	if (_stream == nil)
		@throw [OFNotOpenException exceptionWithObject: self];

	if (_stream.atEndOfStream && _bufferLength - _bufferIndex == 0 &&
	    _state == StateBlockHeader)
		return 0;

start:
	switch ((enum state)_state) {
	case StateBlockHeader:
		if OF_UNLIKELY (!tryReadBits(self, &bits, 16))
			return bytesWritten;

		_symbolsLeft = bits;

		_state = StateCodeLenCodesCount;
		goto start;
	case StateCodeLenCodesCount:
		if OF_UNLIKELY (!tryReadBits(self, &bits, 5))
			return bytesWritten;

		if OF_UNLIKELY (bits > 20)
			@throw [OFInvalidFormatException exception];

		if OF_UNLIKELY (bits == 0) {
			_state = StateCodeLenTreeSingle;
			goto start;
		}

		_codesCount = bits;
		_codesReceived = 0;
		_codesLengths = OFAllocZeroedMemory(bits, 1);
		_skip = true;

		_state = StateCodeLenTree;
		goto start;
	case StateCodeLenTree:
		while (_codesReceived < _codesCount) {
			if OF_UNLIKELY (_currentIsExtendedLength) {
				if OF_UNLIKELY (!tryReadBits(self, &bits, 1))
					return bytesWritten;

				if OF_UNLIKELY (bits == 0) {
					_codesReceived++;

				_codesReceived++;
		}

		_codeLenTree = OFHuffmanTreeNew(_codesLengths, _codesCount);
		OFFreeMemory(_codesLengths);
		_codesLengths = NULL;

		_state = StateLitLenCodesCount;
		goto start;
	case StateCodeLenTreeSingle:
		if OF_UNLIKELY (!tryReadBits(self, &bits, 5))
			return bytesWritten;

		_codeLenTree = OFHuffmanTreeNewSingle(bits);

		_state = StateLitLenCodesCount;
		goto start;
	case StateLitLenCodesCount:
		if OF_UNLIKELY (!tryReadBits(self, &bits, 9))
			return bytesWritten;

		if OF_UNLIKELY (bits > 510)
			@throw [OFInvalidFormatException exception];

		if OF_UNLIKELY (bits == 0) {
			OFHuffmanTreeFree(_codeLenTree);
			_codeLenTree = NULL;

			_state = StateLitLenTreeSingle;
			goto start;
		}

		_codesCount = bits;
		_codesReceived = 0;
		_codesLengths = OFAllocZeroedMemory(bits, 1);
		_skip = false;

		_treeIter = _codeLenTree;
		_state = StateLitLenTree;
		goto start;
	case StateLitLenTree:
		while (_codesReceived < _codesCount) {
			if OF_UNLIKELY (_skip) {
				uint16_t skipCount;

				switch (_codesLengths[_codesReceived]) {
				case 0:
					skipCount = 1;

		_litLenTree = OFHuffmanTreeNew(_codesLengths, _codesCount);
		OFFreeMemory(_codesLengths);
		_codesLengths = NULL;

		OFHuffmanTreeFree(_codeLenTree);
		_codeLenTree = NULL;

		_state = StateDistCodesCount;
		goto start;
	case StateLitLenTreeSingle:
		if OF_UNLIKELY (!tryReadBits(self, &bits, 9))
			return bytesWritten;

		_litLenTree = OFHuffmanTreeNewSingle(bits);

		_state = StateDistCodesCount;
		goto start;
	case StateDistCodesCount:
		if OF_UNLIKELY (!tryReadBits(self, &bits, _distanceBits))
			return bytesWritten;

		if OF_UNLIKELY (bits > _dictionaryBits)
			@throw [OFInvalidFormatException exception];

		if OF_UNLIKELY (bits == 0) {
			_state = StateDistTreeSingle;
			goto start;
		}

		_codesCount = bits;
		_codesReceived = 0;
		_codesLengths = OFAllocZeroedMemory(bits, 1);

		_treeIter = _codeLenTree;
		_state = StateDistTree;
		goto start;
	case StateDistTree:
		while (_codesReceived < _codesCount) {
			if OF_UNLIKELY (_currentIsExtendedLength) {
				if OF_UNLIKELY (!tryReadBits(self, &bits, 1))
					return bytesWritten;

				if OF_UNLIKELY (bits == 0) {
					_codesReceived++;

		}

		_distTree = OFHuffmanTreeNew(_codesLengths, _codesCount);
		OFFreeMemory(_codesLengths);
		_codesLengths = NULL;

		_treeIter = _litLenTree;
		_state = StateBlockLitLen;
		goto start;
	case StateDistTreeSingle:
		if OF_UNLIKELY (!tryReadBits(self, &bits, _distanceBits))
			return bytesWritten;

		_distTree = OFHuffmanTreeNewSingle(bits);

		_treeIter = _litLenTree;
		_state = StateBlockLitLen;
		goto start;
	case StateBlockLitLen:
		if OF_UNLIKELY (_symbolsLeft == 0) {
			OFHuffmanTreeFree(_litLenTree);
			OFHuffmanTreeFree(_distTree);
			_litLenTree = _distTree = NULL;

			_state = StateBlockHeader;

			/*
			 * We must return here, as there is no indication
			 * whether this was the last block. Whoever called this
			 * method needs to check if everything has been read
			 * already and only call read again if that is not the
			 * case.

			    _slidingWindowMask;

			_symbolsLeft--;
			_treeIter = _litLenTree;
		} else {
			_length = value - 253;
			_treeIter = _distTree;
			_state = StateBlockDistLength;
		}

		goto start;
	case StateBlockDistLength:
		if OF_UNLIKELY (!OFHuffmanTreeWalk(self, tryReadBits,
		    &_treeIter, &value))
			return bytesWritten;

		_distance = value;

		_state = (value < 2
		    ? StateBlockLenDistPair : StateBlockDistLengthExtra);
		goto start;
	case StateBlockDistLengthExtra:
		if OF_UNLIKELY (!tryReadBits(self, &bits, _distance - 1))
			return bytesWritten;

		_distance = bits + (1u << (_distance - 1));

		_state = StateBlockLenDistPair;
		goto start;
	case StateBlockLenDistPair:
		for (uint_fast16_t i = 0; i < _length; i++) {
			uint32_t idx;

			if OF_UNLIKELY (length == 0) {
				_length -= i;
				return bytesWritten;
			}

			_slidingWindowIndex = (_slidingWindowIndex + 1) &
			    _slidingWindowMask;
		}

		_symbolsLeft--;

		_treeIter = _litLenTree;
		_state = StateBlockLitLen;
		goto start;
	}

	OF_UNREACHABLE
}

- (bool)lowlevelIsAtEndOfStream
{
	if (_stream == nil)
		@throw [OFNotOpenException exceptionWithObject: self];

	return (_stream.atEndOfStream &&
	    _bufferLength - _bufferIndex == 0 && _state == StateBlockHeader);
}

- (int)fileDescriptorForReading
{
	return ((id <OFReadyForReadingObserving>)_stream)
	    .fileDescriptorForReading;
}