ai.vespa.airlift.zstd.ZstdInputStream Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of airlift-zstd Show documentation
Show all versions of airlift-zstd Show documentation
Fork of https://github.com/airlift/aircompressor (zstd only).
This module is temporary until we get an official release that includes the
ZstdInputStream API (which is already implemented by two different people
but neither PR shows any progress).
/*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package ai.vespa.airlift.zstd;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import static ai.vespa.airlift.zstd.Constants.COMPRESSED_BLOCK;
import static ai.vespa.airlift.zstd.Constants.MAGIC_NUMBER;
import static ai.vespa.airlift.zstd.Constants.MAGIC_SKIPFRAME_MAX;
import static ai.vespa.airlift.zstd.Constants.MAGIC_SKIPFRAME_MIN;
import static ai.vespa.airlift.zstd.Constants.MAX_BLOCK_SIZE;
import static ai.vespa.airlift.zstd.Constants.RAW_BLOCK;
import static ai.vespa.airlift.zstd.Constants.RLE_BLOCK;
import static ai.vespa.airlift.zstd.Constants.SIZE_OF_BLOCK_HEADER;
import static ai.vespa.airlift.zstd.Constants.SIZE_OF_BYTE;
import static ai.vespa.airlift.zstd.Constants.SIZE_OF_INT;
import static ai.vespa.airlift.zstd.Util.fail;
import static sun.misc.Unsafe.ARRAY_BYTE_BASE_OFFSET;
/**
* Take a compressed InputStream and decompress it as needed
* @author arnej27959
*/
public class ZstdInputStream
extends InputStream
{
private static final int DEFAULT_BUFFER_SIZE = 8 * 1024;
private static final int BUFFER_SIZE_MASK = ~(DEFAULT_BUFFER_SIZE - 1);
private static final int MAX_WINDOW_SIZE = 1 << 23;
private final InputStream inputStream;
private byte[] inputBuffer;
private int inputPosition;
private int inputEnd;
private byte[] outputBuffer;
private int outputPosition;
private int outputEnd;
private boolean isClosed;
private boolean seenEof;
private boolean lastBlock;
private boolean singleSegmentFlag;
private boolean contentChecksumFlag;
private long skipBytes;
private int windowSize;
private int blockMaximumSize = MAX_BLOCK_SIZE;
private int curBlockSize;
private int curBlockType = -1;
private FrameHeader curHeader;
private ZstdBlockDecompressor blockDecompressor;
private XxHash64 hasher;
private long evictedInput;
public ZstdInputStream(InputStream inp, int initialBufferSize)
{
this.inputStream = inp;
this.inputBuffer = new byte[initialBufferSize];
this.outputBuffer = new byte[initialBufferSize];
}
public ZstdInputStream(InputStream inp)
{
this(inp, DEFAULT_BUFFER_SIZE);
}
@Override
public int available()
{
return outputAvailable();
}
@Override
public int read() throws IOException
{
throwIfClosed();
if (ensureGotOutput()) {
int b = outputBuffer[outputPosition++];
return (b & 0xFF);
}
else {
return -1;
}
}
@Override
public int read(byte[] b) throws IOException
{
return read(b, 0, b.length);
}
@Override
public int read(byte[] b, int off, int len) throws IOException
{
throwIfClosed();
if (ensureGotOutput()) {
len = Math.min(outputAvailable(), len);
System.arraycopy(outputBuffer, outputPosition, b, off, len);
outputPosition += len;
return len;
}
else {
return -1;
}
}
@Override
public void close() throws IOException
{
throwIfClosed();
if (!seenEof) {
inputStream.close();
}
isClosed = true;
}
private void check(boolean condition, String reason)
{
Util.verify(condition, curInputFilePosition(), reason);
}
private boolean ensureGotOutput() throws IOException
{
while ((outputAvailable() == 0) && !seenEof) {
if (ensureGotFrameHeader() && ensureGotBlock()) {
decompressBlock();
}
}
if (outputAvailable() > 0) {
return true;
}
else {
check(seenEof, "unable to decode to EOF");
check(inputAvailable() == 0, "leftover input at end of file");
check(curHeader == null, "unfinished frame at end of file");
return false;
}
}
private void readMoreInput() throws IOException
{
ensureInputSpace(1024);
int got = inputStream.read(inputBuffer, inputEnd, inputSpace());
if (got == -1) {
seenEof = true;
inputStream.close();
}
else {
inputEnd += got;
}
}
private ByteBuffer inputBB()
{
ByteBuffer bb = ByteBuffer.wrap(inputBuffer, inputPosition, inputAvailable());
bb.order(ByteOrder.LITTLE_ENDIAN);
return bb;
}
private boolean ensureGotFrameHeader() throws IOException
{
if (curHeader != null) {
return true;
}
// a skip frame is minimum 8 bytes
// a data frame is minimum 4 + 2 + 3 = 9 bytes, but we only
// need 5 bytes to know the size of the frame header
if (inputAvailable() < 8) {
readMoreInput();
// retry from start
return false;
}
ByteBuffer bb = inputBB();
int magic = bb.getInt();
// skippable frame header magic
if ((magic >= MAGIC_SKIPFRAME_MIN) && (magic <= MAGIC_SKIPFRAME_MAX)) {
inputPosition += SIZE_OF_INT; // for magic
skipBytes = (bb.getInt() & 0xffff_ffffL) + SIZE_OF_INT;
inputPosition += SIZE_OF_INT; // for skipsize
while (skipBytes > 0) {
if (skipBytes <= inputAvailable()) {
inputPosition += skipBytes;
skipBytes = 0;
}
else {
skipBytes -= inputAvailable();
inputPosition = inputEnd;
readMoreInput();
if (seenEof) {
throw fail(curInputFilePosition(), "unfinished skip frame at end of file");
}
}
}
// entire frame skipped; retry from start
return false;
}
// zstd frame header magic
if (magic == MAGIC_NUMBER) {
int fhDesc = 0xFF & bb.get();
int frameContentSizeFlag = (fhDesc & 0b11000000) >> 6;
singleSegmentFlag = (fhDesc & 0b00100000) != 0;
contentChecksumFlag = (fhDesc & 0b00000100) != 0;
int dictionaryIdFlag = (fhDesc & 0b00000011);
// 4 byte magic + 1 byte fhDesc
int fhSize = SIZE_OF_INT + SIZE_OF_BYTE;
// add size of frameContentSize
if (frameContentSizeFlag == 0) {
fhSize += (singleSegmentFlag ? 1 : 0);
}
else {
fhSize += 1 << frameContentSizeFlag;
}
// add size of window descriptor
fhSize += (singleSegmentFlag ? 0 : 1);
// add size of dictionary id
fhSize += (1 << dictionaryIdFlag) >> 1;
if (fhSize > inputAvailable()) {
readMoreInput();
// retry from start
return false;
}
inputPosition += SIZE_OF_INT;
curHeader = readFrameHeader();
inputPosition += fhSize - SIZE_OF_INT;
startFrame();
return true;
}
else {
throw fail(curInputFilePosition(), "Invalid magic prefix: " + magic);
}
}
private void startFrame()
{
blockDecompressor = new ZstdBlockDecompressor(curHeader);
check(outputPosition == outputEnd, "orphan output present");
outputPosition = 0;
outputEnd = 0;
if (singleSegmentFlag) {
if (curHeader.contentSize > MAX_WINDOW_SIZE) {
throw fail(curInputFilePosition(), "Single segment too large: " + curHeader.contentSize);
}
windowSize = (int) curHeader.contentSize;
blockMaximumSize = windowSize;
ensureOutputSpace(windowSize);
}
else {
if (curHeader.windowSize > MAX_WINDOW_SIZE) {
throw fail(curInputFilePosition(), "Window size too large: " + curHeader.windowSize);
}
windowSize = curHeader.windowSize;
blockMaximumSize = Math.min(windowSize, MAX_BLOCK_SIZE);
ensureOutputSpace(blockMaximumSize + windowSize);
}
if (contentChecksumFlag) {
hasher = new XxHash64();
}
}
private boolean ensureGotBlock() throws IOException
{
check(curHeader != null, "no current frame");
if (curBlockType == -1) {
// must have a block now
if (inputAvailable() < SIZE_OF_BLOCK_HEADER) {
readMoreInput();
// retry from start
return false;
}
int blkHeader = nextByte() | nextByte() << 8 | nextByte() << 16;
lastBlock = (blkHeader & 0b001) != 0;
curBlockType = (blkHeader & 0b110) >> 1;
curBlockSize = blkHeader >> 3;
ensureInputSpace(curBlockSize + SIZE_OF_INT);
}
if (inputAvailable() < curBlockSize + (contentChecksumFlag ? SIZE_OF_INT : 0)) {
readMoreInput();
// retry from start
return false;
}
return true;
}
int nextByte()
{
int r = 0xFF & inputBuffer[inputPosition];
inputPosition++;
return r;
}
long inputAddress()
{
return ARRAY_BYTE_BASE_OFFSET + inputPosition;
}
long inputLimit()
{
return ARRAY_BYTE_BASE_OFFSET + inputEnd;
}
long outputAddress()
{
return ARRAY_BYTE_BASE_OFFSET + outputEnd;
}
long outputLimit()
{
return ARRAY_BYTE_BASE_OFFSET + outputBuffer.length;
}
int decodeRaw()
{
check(inputAddress() + curBlockSize <= inputLimit(), "Not enough input bytes");
check(outputAddress() + curBlockSize <= outputLimit(), "Not enough output space");
return ZstdBlockDecompressor.decodeRawBlock(inputBuffer, inputAddress(), curBlockSize, outputBuffer, outputAddress(), outputLimit());
}
int decodeRle()
{
check(inputAddress() + 1 <= inputLimit(), "Not enough input bytes");
check(outputAddress() + curBlockSize <= outputLimit(), "Not enough output space");
return ZstdBlockDecompressor.decodeRleBlock(curBlockSize, inputBuffer, inputAddress(), outputBuffer, outputAddress(), outputLimit());
}
int decodeCompressed()
{
check(inputAddress() + curBlockSize <= inputLimit(), "Not enough input bytes");
check(outputAddress() + blockMaximumSize <= outputLimit(), "Not enough output space");
return blockDecompressor.decodeCompressedBlock(
inputBuffer, inputAddress(),
curBlockSize,
outputBuffer, outputAddress(), outputLimit(),
windowSize, ARRAY_BYTE_BASE_OFFSET);
}
private void decompressBlock()
{
check(outputPosition == outputEnd, "orphan output present");
switch (curBlockType) {
case RAW_BLOCK:
ensureOutputSpace(curBlockSize);
outputEnd += decodeRaw();
inputPosition += curBlockSize;
break;
case RLE_BLOCK:
ensureOutputSpace(curBlockSize);
outputEnd += decodeRle();
inputPosition += 1;
break;
case COMPRESSED_BLOCK:
check(curBlockSize < blockMaximumSize, "compressed block must be smaller than Block_Maximum_Size");
ensureOutputSpace(blockMaximumSize);
outputEnd += decodeCompressed();
inputPosition += curBlockSize;
break;
default:
throw fail(curInputFilePosition(), "Invalid block type " + curBlockType);
}
if (contentChecksumFlag) {
hasher.update(outputBuffer, outputPosition, outputAvailable());
}
curBlockType = -1;
if (lastBlock) {
curHeader = null;
blockDecompressor = null;
if (contentChecksumFlag) {
check(inputAvailable() >= SIZE_OF_INT, "missing checksum data");
long hash = hasher.hash();
int checksum = inputBB().getInt();
if (checksum != (int) hash) {
throw fail(curInputFilePosition(), String.format("Bad checksum. Expected: %s, actual: %s", Integer.toHexString(checksum), Integer.toHexString((int) hash)));
}
inputPosition += SIZE_OF_INT;
hasher = null;
}
}
}
private int inputAvailable()
{
return inputEnd - inputPosition;
}
private int inputSpace()
{
return inputBuffer.length - inputEnd;
}
private long curInputFilePosition()
{
return evictedInput + inputPosition;
}
private void ensureInputSpace(int size)
{
if (inputSpace() < size) {
if (size < inputPosition) {
System.arraycopy(inputBuffer, inputPosition, inputBuffer, 0, inputAvailable());
}
else {
int newSize = (inputBuffer.length + size + DEFAULT_BUFFER_SIZE) & BUFFER_SIZE_MASK;
byte[] newBuf = new byte[newSize];
System.arraycopy(inputBuffer, inputPosition, newBuf, 0, inputAvailable());
inputBuffer = newBuf;
}
evictedInput += inputPosition;
inputEnd = inputAvailable();
inputPosition = 0;
}
}
private int outputAvailable()
{
return outputEnd - outputPosition;
}
private int outputSpace()
{
return outputBuffer.length - outputEnd;
}
private void ensureOutputSpace(int size)
{
if (outputSpace() < size) {
check(outputAvailable() == 0, "logic error");
byte[] newBuf;
if (windowSize * 4 + size < outputPosition) {
// plenty space in old buffer
newBuf = outputBuffer;
}
else {
int newSize = (outputBuffer.length
+ windowSize * 4
+ size
+ DEFAULT_BUFFER_SIZE) & BUFFER_SIZE_MASK;
newBuf = new byte[newSize];
}
// keep up to one window of old data
int sizeToKeep = Math.min(outputPosition, windowSize);
System.arraycopy(outputBuffer, outputPosition - sizeToKeep, newBuf, 0, sizeToKeep);
outputBuffer = newBuf;
outputEnd = sizeToKeep;
outputPosition = sizeToKeep;
}
}
private void throwIfClosed() throws IOException
{
if (isClosed) {
throw new IOException("Input stream is already closed");
}
}
private FrameHeader readFrameHeader()
{
long base = ARRAY_BYTE_BASE_OFFSET + inputPosition;
long limit = ARRAY_BYTE_BASE_OFFSET + inputEnd;
return ZstdFrameDecompressor.readFrameHeader(inputBuffer, base, limit);
}
}