BlockLZ4CompressorOutputStream.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.commons.compress.compressors.lz4;

import java.io.IOException;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Deque;
import java.util.Iterator;
import java.util.LinkedList;

import org.apache.commons.compress.compressors.CompressorOutputStream;
import org.apache.commons.compress.compressors.lz77support.LZ77Compressor;
import org.apache.commons.compress.compressors.lz77support.Parameters;
import org.apache.commons.compress.utils.ByteUtils;

/**
 * CompressorOutputStream for the LZ4 block format.
 *
 * @see <a href="https://lz4.github.io/lz4/lz4_Block_format.html">LZ4 Block Format Description</a>
 * @since 1.14
 * @NotThreadSafe
 */
public class BlockLZ4CompressorOutputStream extends CompressorOutputStream<OutputStream> {

    static final class Pair {

        private static int lengths(final int litLength, final int brLength) {
            final int l = Math.min(litLength, 15);
            final int br = brLength < 4 ? 0 : brLength < 19 ? brLength - 4 : 15;
            return l << BlockLZ4CompressorInputStream.SIZE_BITS | br;
        }

        private static void writeLength(int length, final OutputStream out) throws IOException {
            while (length >= 255) {
                out.write(255);
                length -= 255;
            }
            out.write(length);
        }

        private final Deque<byte[]> literals = new LinkedList<>();

        private int literalLength;

        private int brOffset, brLength;

        private boolean written;

        byte[] addLiteral(final LZ77Compressor.LiteralBlock block) {
            final byte[] copy = Arrays.copyOfRange(block.getData(), block.getOffset(), block.getOffset() + block.getLength());
            literals.add(copy);
            literalLength += copy.length;
            return copy;
        }

        private int backReferenceLength() {
            return brLength;
        }

        boolean canBeWritten(final int lengthOfBlocksAfterThisPair) {
            return hasBackReference() && lengthOfBlocksAfterThisPair >= MIN_OFFSET_OF_LAST_BACK_REFERENCE + MIN_BACK_REFERENCE_LENGTH;
        }

        boolean hasBackReference() {
            return brOffset > 0;
        }

        private boolean hasBeenWritten() {
            return written;
        }

        int length() {
            return literalLength() + brLength;
        }

        private int literalLength() {
            // This method is performance sensitive
            if (literalLength != 0) {
                return literalLength;
            }
            int length = 0;
            for (final byte[] b : literals) {
                length += b.length;
            }
            return literalLength = length;
        }

        private void prependLiteral(final byte[] data) {
            literals.addFirst(data);
            literalLength += data.length;
        }

        private void prependTo(final Pair other) {
            final Iterator<byte[]> listBackwards = literals.descendingIterator();
            while (listBackwards.hasNext()) {
                other.prependLiteral(listBackwards.next());
            }
        }

        void setBackReference(final LZ77Compressor.BackReference block) {
            if (hasBackReference()) {
                throw new IllegalStateException();
            }
            brOffset = block.getOffset();
            brLength = block.getLength();
        }

        private Pair splitWithNewBackReferenceLengthOf(final int newBackReferenceLength) {
            final Pair p = new Pair();
            p.literals.addAll(literals);
            p.brOffset = brOffset;
            p.brLength = newBackReferenceLength;
            return p;
        }

        void writeTo(final OutputStream out) throws IOException {
            final int litLength = literalLength();
            out.write(lengths(litLength, brLength));
            if (litLength >= BlockLZ4CompressorInputStream.BACK_REFERENCE_SIZE_MASK) {
                writeLength(litLength - BlockLZ4CompressorInputStream.BACK_REFERENCE_SIZE_MASK, out);
            }
            for (final byte[] b : literals) {
                out.write(b);
            }
            if (hasBackReference()) {
                ByteUtils.toLittleEndian(out, brOffset, 2);
                if (brLength - MIN_BACK_REFERENCE_LENGTH >= BlockLZ4CompressorInputStream.BACK_REFERENCE_SIZE_MASK) {
                    writeLength(brLength - MIN_BACK_REFERENCE_LENGTH - BlockLZ4CompressorInputStream.BACK_REFERENCE_SIZE_MASK, out);
                }
            }
            written = true;
        }
    }

    private static final int MIN_BACK_REFERENCE_LENGTH = 4;

    /*
     *
     * The LZ4 block format has a few properties that make it less straight-forward than one would hope:
     *
     * literal blocks and back-references must come in pairs (except for the very last literal block), so consecutive literal blocks created by the compressor
     * must be merged into a single block.
     *
     * the start of a literal/back-reference pair contains the length of the back-reference (at least some part of it) so we can't start writing the literal
     * before we know how long the next back-reference is going to be.
     *
     * there are special rules for the final blocks
     *
     * > There are specific parsing rules to respect in order to remain > compatible with assumptions made by the decoder : > > 1. The last 5 bytes are always
     * literals > > 2. The last match must start at least 12 bytes before end of > block. Consequently, a block with less than 13 bytes cannot be > compressed.
     *
     * which means any back-reference may need to get rewritten as a literal block unless we know the next block is at least of length 5 and the sum of this
     * block's length and offset and the next block's length is at least twelve.
     */

    private static final int MIN_OFFSET_OF_LAST_BACK_REFERENCE = 12;

    /**
     * Returns a builder correctly configured for the LZ4 algorithm.
     *
     * @return a builder correctly configured for the LZ4 algorithm
     */
    public static Parameters.Builder createParameterBuilder() {
        final int maxLen = BlockLZ4CompressorInputStream.WINDOW_SIZE - 1;
        return Parameters.builder(BlockLZ4CompressorInputStream.WINDOW_SIZE).withMinBackReferenceLength(MIN_BACK_REFERENCE_LENGTH)
                .withMaxBackReferenceLength(maxLen).withMaxOffset(maxLen).withMaxLiteralLength(maxLen);
    }

    private final LZ77Compressor compressor;

    // used in one-arg write method
    private final byte[] oneByte = new byte[1];
    private boolean finished;

    private final Deque<Pair> pairs = new LinkedList<>();

    // keeps track of the last window-size bytes (64k) in order to be
    // able to expand back-references when needed
    private final Deque<byte[]> expandedBlocks = new LinkedList<>();

    /**
     * Creates a new LZ4 output stream.
     *
     * @param out An OutputStream to read compressed data from
     */
    public BlockLZ4CompressorOutputStream(final OutputStream out) {
        this(out, createParameterBuilder().build());
    }

    /**
     * Creates a new LZ4 output stream.
     *
     * @param out     An OutputStream to read compressed data from
     * @param params The parameters to use for LZ77 compression.
     */
    public BlockLZ4CompressorOutputStream(final OutputStream out, final Parameters params) {
        super(out);
        compressor = new LZ77Compressor(params, block -> {
            switch (block.getType()) {
            case LITERAL:
                addLiteralBlock((LZ77Compressor.LiteralBlock) block);
                break;
            case BACK_REFERENCE:
                addBackReference((LZ77Compressor.BackReference) block);
                break;
            case EOD:
                writeFinalLiteralBlock();
                break;
            }
        });
    }

    private void addBackReference(final LZ77Compressor.BackReference block) throws IOException {
        final Pair last = writeBlocksAndReturnUnfinishedPair(block.getLength());
        last.setBackReference(block);
        recordBackReference(block);
        clearUnusedBlocksAndPairs();
    }

    private void addLiteralBlock(final LZ77Compressor.LiteralBlock block) throws IOException {
        final Pair last = writeBlocksAndReturnUnfinishedPair(block.getLength());
        recordLiteral(last.addLiteral(block));
        clearUnusedBlocksAndPairs();
    }

    private void clearUnusedBlocks() {
        int blockLengths = 0;
        int blocksToKeep = 0;
        for (final byte[] b : expandedBlocks) {
            blocksToKeep++;
            blockLengths += b.length;
            if (blockLengths >= BlockLZ4CompressorInputStream.WINDOW_SIZE) {
                break;
            }
        }
        final int size = expandedBlocks.size();
        for (int i = blocksToKeep; i < size; i++) {
            expandedBlocks.removeLast();
        }
    }

    private void clearUnusedBlocksAndPairs() {
        clearUnusedBlocks();
        clearUnusedPairs();
    }

    private void clearUnusedPairs() {
        int pairLengths = 0;
        int pairsToKeep = 0;
        for (final Iterator<Pair> it = pairs.descendingIterator(); it.hasNext();) {
            final Pair p = it.next();
            pairsToKeep++;
            pairLengths += p.length();
            if (pairLengths >= BlockLZ4CompressorInputStream.WINDOW_SIZE) {
                break;
            }
        }
        final int size = pairs.size();
        for (int i = pairsToKeep; i < size; i++) {
            final Pair p = pairs.peekFirst();
            if (!p.hasBeenWritten()) {
                break;
            }
            pairs.removeFirst();
        }
    }

    @Override
    public void close() throws IOException {
        try {
            finish();
        } finally {
            out.close();
        }
    }

    private byte[] expand(final int offset, final int length) {
        final byte[] expanded = new byte[length];
        if (offset == 1) { // surprisingly common special case
            final byte[] block = expandedBlocks.peekFirst();
            final byte b = block[block.length - 1];
            if (b != 0) { // the fresh array contains 0s anyway
                Arrays.fill(expanded, b);
            }
        } else {
            expandFromList(expanded, offset, length);
        }
        return expanded;
    }

    private void expandFromList(final byte[] expanded, final int offset, final int length) {
        int offsetRemaining = offset;
        int lengthRemaining = length;
        int writeOffset = 0;
        while (lengthRemaining > 0) {
            // find block that contains offsetRemaining
            byte[] block = null;
            final int copyLen;
            final int copyOffset;
            if (offsetRemaining > 0) {
                int blockOffset = 0;
                for (final byte[] b : expandedBlocks) {
                    if (b.length + blockOffset >= offsetRemaining) {
                        block = b;
                        break;
                    }
                    blockOffset += b.length;
                }
                if (block == null) {
                    // should not be possible
                    throw new IllegalStateException("Failed to find a block containing offset " + offset);
                }
                copyOffset = blockOffset + block.length - offsetRemaining;
                copyLen = Math.min(lengthRemaining, block.length - copyOffset);
            } else {
                // offsetRemaining is negative or 0 and points into the expanded bytes
                block = expanded;
                copyOffset = -offsetRemaining;
                copyLen = Math.min(lengthRemaining, writeOffset + offsetRemaining);
            }
            System.arraycopy(block, copyOffset, expanded, writeOffset, copyLen);
            offsetRemaining -= copyLen;
            lengthRemaining -= copyLen;
            writeOffset += copyLen;
        }
    }

    /**
     * Compresses all remaining data and writes it to the stream, doesn't close the underlying stream.
     *
     * @throws IOException if an error occurs
     */
    public void finish() throws IOException {
        if (!finished) {
            compressor.finish();
            finished = true;
        }
    }

    /**
     * Adds some initial data to fill the window with.
     *
     * @param data the data to fill the window with.
     * @param off  offset of real data into the array
     * @param len  amount of data
     * @throws IllegalStateException if the stream has already started to write data
     * @see LZ77Compressor#prefill
     */
    public void prefill(final byte[] data, final int off, final int len) {
        if (len > 0) {
            final byte[] b = Arrays.copyOfRange(data, off, off + len);
            compressor.prefill(b);
            recordLiteral(b);
        }
    }

    private void recordBackReference(final LZ77Compressor.BackReference block) {
        expandedBlocks.addFirst(expand(block.getOffset(), block.getLength()));
    }

    private void recordLiteral(final byte[] b) {
        expandedBlocks.addFirst(b);
    }

    private void rewriteLastPairs() {
        final LinkedList<Pair> lastPairs = new LinkedList<>();
        final LinkedList<Integer> pairLength = new LinkedList<>();
        int offset = 0;
        for (final Iterator<Pair> it = pairs.descendingIterator(); it.hasNext();) {
            final Pair p = it.next();
            if (p.hasBeenWritten()) {
                break;
            }
            final int len = p.length();
            pairLength.addFirst(len);
            lastPairs.addFirst(p);
            offset += len;
            if (offset >= MIN_OFFSET_OF_LAST_BACK_REFERENCE) {
                break;
            }
        }
        lastPairs.forEach(pairs::remove);
        // lastPairs may contain between one and four Pairs:
        // * the last pair may be a one byte literal
        // * all other Pairs contain a back-reference which must be four bytes long at minimum
        // we could merge them all into a single literal block but
        // this may harm compression. For example compressing
        // "bla.tar" from our tests yields a last block containing a
        // back-reference of length > 2k and we'd end up with a last
        // literal of that size rather than a 2k back-reference and a
        // 12 byte literal at the end.

        // Instead we merge all but the first of lastPairs into a new
        // literal-only Pair "replacement" and look at the
        // back-reference in the first of lastPairs and see if we can
        // split it. We can split it if it is longer than 16 -
        // replacement.length (i.e. the minimal length of four is kept
        // while making sure the last literal is at least twelve bytes
        // long). If we can't split it, we expand the first of the pairs
        // as well.

        // this is not optimal, we could get better compression
        // results with more complex approaches as the last literal
        // only needs to be five bytes long if the previous
        // back-reference has an offset big enough

        final int lastPairsSize = lastPairs.size();
        int toExpand = 0;
        for (int i = 1; i < lastPairsSize; i++) {
            toExpand += pairLength.get(i);
        }
        final Pair replacement = new Pair();
        if (toExpand > 0) {
            replacement.prependLiteral(expand(toExpand, toExpand));
        }
        final Pair splitCandidate = lastPairs.get(0);
        final int stillNeeded = MIN_OFFSET_OF_LAST_BACK_REFERENCE - toExpand;
        final int brLen = splitCandidate.hasBackReference() ? splitCandidate.backReferenceLength() : 0;
        if (splitCandidate.hasBackReference() && brLen >= MIN_BACK_REFERENCE_LENGTH + stillNeeded) {
            replacement.prependLiteral(expand(toExpand + stillNeeded, stillNeeded));
            pairs.add(splitCandidate.splitWithNewBackReferenceLengthOf(brLen - stillNeeded));
        } else {
            if (splitCandidate.hasBackReference()) {
                replacement.prependLiteral(expand(toExpand + brLen, brLen));
            }
            splitCandidate.prependTo(replacement);
        }
        pairs.add(replacement);
    }

    @Override
    public void write(final byte[] data, final int off, final int len) throws IOException {
        compressor.compress(data, off, len);
    }

    @Override
    public void write(final int b) throws IOException {
        oneByte[0] = (byte) (b & 0xff);
        write(oneByte);
    }

    private Pair writeBlocksAndReturnUnfinishedPair(final int length) throws IOException {
        writeWritablePairs(length);
        Pair last = pairs.peekLast();
        if (last == null || last.hasBackReference()) {
            last = new Pair();
            pairs.addLast(last);
        }
        return last;
    }

    private void writeFinalLiteralBlock() throws IOException {
        rewriteLastPairs();
        for (final Pair p : pairs) {
            if (!p.hasBeenWritten()) {
                p.writeTo(out);
            }
        }
        pairs.clear();
    }

    private void writeWritablePairs(final int lengthOfBlocksAfterLastPair) throws IOException {
        int unwrittenLength = lengthOfBlocksAfterLastPair;
        for (final Iterator<Pair> it = pairs.descendingIterator(); it.hasNext();) {
            final Pair p = it.next();
            if (p.hasBeenWritten()) {
                break;
            }
            unwrittenLength += p.length();
        }
        for (final Pair p : pairs) {
            if (p.hasBeenWritten()) {
                continue;
            }
            unwrittenLength -= p.length();
            if (!p.canBeWritten(unwrittenLength)) {
                break;
            }
            p.writeTo(out);
        }
    }
}