    001/*
    002 * Licensed to the Apache Software Foundation (ASF) under one
    003 * or more contributor license agreements.  See the NOTICE file
    004 * distributed with this work for additional information
    005 * regarding copyright ownership.  The ASF licenses this file
    006 * to you under the Apache License, Version 2.0 (the
    007 * "License"); you may not use this file except in compliance
    008 * with the License.  You may obtain a copy of the License at
    009 *
    010 * http://www.apache.org/licenses/LICENSE-2.0
    011 *
    012 * Unless required by applicable law or agreed to in writing,
    013 * software distributed under the License is distributed on an
    014 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    015 * KIND, either express or implied.  See the License for the
    016 * specific language governing permissions and limitations
    017 * under the License.
    018 */
    019package org.apache.commons.compress.compressors.snappy;
    020
    021import java.io.IOException;
    022import java.io.OutputStream;
    023
    024import org.apache.commons.compress.compressors.CompressorOutputStream;
    025import org.apache.commons.compress.compressors.lz77support.LZ77Compressor;
    026import org.apache.commons.compress.compressors.lz77support.Parameters;
    027import org.apache.commons.compress.utils.ByteUtils;
    028
    029/**
    030 * CompressorOutputStream for the raw Snappy format.
    031 *
    032 * <p>This implementation uses an internal buffer in order to handle
    033 * the back-references that are at the heart of the LZ77 algorithm.
    034 * The size of the buffer must be at least as big as the biggest
    035 * offset used in the compressed stream.  The current version of the
    036 * Snappy algorithm as defined by Google works on 32k blocks and
    037 * doesn't contain offsets bigger than 32k which is the default block
    038 * size used by this class.</p>
    039 *
    040 * <p>The raw Snappy format requires the uncompressed size to be
    041 * written at the beginning of the stream using a varint
    042 * representation, i.e. the number of bytes needed to write the
    043 * information is not known before the uncompressed size is
    044 * known. We've chosen to make the uncompressedSize a parameter of the
    045 * constructor in favor of buffering the whole output until the size
    046 * is known. When using the {@link FramedSnappyCompressorOutputStream}
    047 * this limitation is taken care of by the warpping framing
    048 * format.</p>
    049 *
    050 * @see <a href="https://github.com/google/snappy/blob/master/format_description.txt">Snappy compressed format description</a>
    051 * @since 1.14
    052 * @NotThreadSafe
    053 */
    054public class SnappyCompressorOutputStream extends CompressorOutputStream {
    055    private final LZ77Compressor compressor;
    056    private final OutputStream os;
    057    private final ByteUtils.ByteConsumer consumer;
    058
    059    // used in one-arg write method
    060    private final byte[] oneByte = new byte[1];
    061
    062    private boolean finished;
    063
    064    /**
    065     * Constructor using the default block size of 32k.
    066     *
    067     * @param os the outputstream to write compressed data to
    068     * @param uncompressedSize the uncompressed size of data
    069     * @throws IOException if writing of the size fails
    070     */
    071    public SnappyCompressorOutputStream(final OutputStream os, final long uncompressedSize) throws IOException {
    072        this(os, uncompressedSize, SnappyCompressorInputStream.DEFAULT_BLOCK_SIZE);
    073    }
    074
    075    /**
    076     * Constructor using a configurable block size.
    077     *
    078     * @param os the outputstream to write compressed data to
    079     * @param uncompressedSize the uncompressed size of data
    080     * @param blockSize the block size used - must be a power of two
    081     * @throws IOException if writing of the size fails
    082     */
    083    public SnappyCompressorOutputStream(final OutputStream os, final long uncompressedSize, final int blockSize)
    084        throws IOException {
    085        this(os, uncompressedSize, createParameterBuilder(blockSize).build());
    086    }
    087
    088    /**
    089     * Constructor providing full control over the underlying LZ77 compressor.
    090     *
    091     * @param os the outputstream to write compressed data to
    092     * @param uncompressedSize the uncompressed size of data
    093     * @param params the parameters to use by the compressor - note
    094     * that the format itself imposes some limits like a maximum match
    095     * length of 64 bytes
    096     * @throws IOException if writing of the size fails
    097     */
    098    public SnappyCompressorOutputStream(final OutputStream os, final long uncompressedSize, final Parameters params)
    099        throws IOException {
    100        this.os = os;
    101        consumer = new ByteUtils.OutputStreamByteConsumer(os);
    102        compressor = new LZ77Compressor(params, block -> {
    103            switch (block.getType()) {
    104            case LITERAL:
    105                writeLiteralBlock((LZ77Compressor.LiteralBlock) block);
    106                break;
    107            case BACK_REFERENCE:
    108                writeBackReference((LZ77Compressor.BackReference) block);
    109                break;
    110            case EOD:
    111                break;
    112            }
    113        });
    114        writeUncompressedSize(uncompressedSize);
    115    }
    116
    117    @Override
    118    public void write(final int b) throws IOException {
    119        oneByte[0] = (byte) (b & 0xff);
    120        write(oneByte);
    121    }
    122
    123    @Override
    124    public void write(final byte[] data, final int off, final int len) throws IOException {
    125        compressor.compress(data, off, len);
    126    }
    127
    128    @Override
    129    public void close() throws IOException {
    130        try {
    131            finish();
    132        } finally {
    133            os.close();
    134        }
    135    }
    136
    137    /**
    138     * Compresses all remaining data and writes it to the stream,
    139     * doesn't close the underlying stream.
    140     * @throws IOException if an error occurs
    141     */
    142    public void finish() throws IOException {
    143        if (!finished) {
    144            compressor.finish();
    145            finished = true;
    146        }
    147    }
    148
    149    private void writeUncompressedSize(long uncompressedSize) throws IOException {
    150        boolean more = false;
    151        do {
    152            int currentByte = (int) (uncompressedSize & 0x7F);
    153            more = uncompressedSize > currentByte;
    154            if (more) {
    155                currentByte |= 0x80;
    156            }
    157            os.write(currentByte);
    158            uncompressedSize >>= 7;
    159        } while (more);
    160    }
    161
    162    // literal length is stored as (len - 1) either inside the tag
    163    // (six bits minus four flags) or in 1 to 4 bytes after the tag
    164    private static final int MAX_LITERAL_SIZE_WITHOUT_SIZE_BYTES = 60;
    165    private static final int MAX_LITERAL_SIZE_WITH_ONE_SIZE_BYTE = 1 << 8;
    166    private static final int MAX_LITERAL_SIZE_WITH_TWO_SIZE_BYTES = 1 << 16;
    167    private static final int MAX_LITERAL_SIZE_WITH_THREE_SIZE_BYTES = 1 << 24;
    168
    169    private static final int ONE_SIZE_BYTE_MARKER = 60 << 2;
    170    private static final int TWO_SIZE_BYTE_MARKER = 61 << 2;
    171    private static final int THREE_SIZE_BYTE_MARKER = 62 << 2;
    172    private static final int FOUR_SIZE_BYTE_MARKER = 63 << 2;
    173
    174    private void writeLiteralBlock(final LZ77Compressor.LiteralBlock block) throws IOException {
    175        final int len = block.getLength();
    176        if (len <= MAX_LITERAL_SIZE_WITHOUT_SIZE_BYTES) {
    177            writeLiteralBlockNoSizeBytes(block, len);
    178        } else if (len <= MAX_LITERAL_SIZE_WITH_ONE_SIZE_BYTE) {
    179            writeLiteralBlockOneSizeByte(block, len);
    180        } else if (len <= MAX_LITERAL_SIZE_WITH_TWO_SIZE_BYTES) {
    181            writeLiteralBlockTwoSizeBytes(block, len);
    182        } else if (len <= MAX_LITERAL_SIZE_WITH_THREE_SIZE_BYTES) {
    183            writeLiteralBlockThreeSizeBytes(block, len);
    184        } else {
    185            writeLiteralBlockFourSizeBytes(block, len);
    186        }
    187    }
    188
    189    private void writeLiteralBlockNoSizeBytes(final LZ77Compressor.LiteralBlock block, final int len) throws IOException {
    190        writeLiteralBlockWithSize(len - 1 << 2, 0, len, block);
    191    }
    192
    193    private void writeLiteralBlockOneSizeByte(final LZ77Compressor.LiteralBlock block, final int len) throws IOException {
    194        writeLiteralBlockWithSize(ONE_SIZE_BYTE_MARKER, 1, len, block);
    195    }
    196
    197    private void writeLiteralBlockTwoSizeBytes(final LZ77Compressor.LiteralBlock block, final int len) throws IOException {
    198        writeLiteralBlockWithSize(TWO_SIZE_BYTE_MARKER, 2, len, block);
    199    }
    200
    201    private void writeLiteralBlockThreeSizeBytes(final LZ77Compressor.LiteralBlock block, final int len) throws IOException {
    202        writeLiteralBlockWithSize(THREE_SIZE_BYTE_MARKER, 3, len, block);
    203    }
    204
    205    private void writeLiteralBlockFourSizeBytes(final LZ77Compressor.LiteralBlock block, final int len) throws IOException {
    206        writeLiteralBlockWithSize(FOUR_SIZE_BYTE_MARKER, 4, len, block);
    207    }
    208
    209    private void writeLiteralBlockWithSize(final int tagByte, final int sizeBytes, final int len, final LZ77Compressor.LiteralBlock block)
    210        throws IOException {
    211        os.write(tagByte);
    212        writeLittleEndian(sizeBytes, len - 1);
    213        os.write(block.getData(), block.getOffset(), len);
    214    }
    215
    216    private void writeLittleEndian(final int numBytes, final int num) throws IOException {
    217        ByteUtils.toLittleEndian(consumer, num, numBytes);
    218    }
    219
    220    // Back-references ("copies") have their offset/size information
    221    // in two, three or five bytes.
    222    private static final int MIN_MATCH_LENGTH_WITH_ONE_OFFSET_BYTE = 4;
    223    private static final int MAX_MATCH_LENGTH_WITH_ONE_OFFSET_BYTE = 11;
    224    private static final int MAX_OFFSET_WITH_ONE_OFFSET_BYTE = 1 << 11 - 1;
    225    private static final int MAX_OFFSET_WITH_TWO_OFFSET_BYTES = 1 << 16 - 1;
    226
    227    private static final int ONE_BYTE_COPY_TAG = 1;
    228    private static final int TWO_BYTE_COPY_TAG = 2;
    229    private static final int FOUR_BYTE_COPY_TAG = 3;
    230
    231    private void writeBackReference(final LZ77Compressor.BackReference block) throws IOException {
    232        final int len = block.getLength();
    233        final int offset = block.getOffset();
    234        if (len >= MIN_MATCH_LENGTH_WITH_ONE_OFFSET_BYTE && len <= MAX_MATCH_LENGTH_WITH_ONE_OFFSET_BYTE
    235            && offset <= MAX_OFFSET_WITH_ONE_OFFSET_BYTE) {
    236            writeBackReferenceWithOneOffsetByte(len, offset);
    237        } else if (offset < MAX_OFFSET_WITH_TWO_OFFSET_BYTES) {
    238            writeBackReferenceWithTwoOffsetBytes(len, offset);
    239        } else {
    240            writeBackReferenceWithFourOffsetBytes(len, offset);
    241        }
    242    }
    243
    244    private void writeBackReferenceWithOneOffsetByte(final int len, final int offset) throws IOException {
    245        os.write(ONE_BYTE_COPY_TAG | ((len - 4) << 2) | ((offset & 0x700) >> 3));
    246        os.write(offset & 0xff);
    247    }
    248
    249    private void writeBackReferenceWithTwoOffsetBytes(final int len, final int offset) throws IOException {
    250        writeBackReferenceWithLittleEndianOffset(TWO_BYTE_COPY_TAG, 2, len, offset);
    251    }
    252
    253    private void writeBackReferenceWithFourOffsetBytes(final int len, final int offset) throws IOException {
    254        writeBackReferenceWithLittleEndianOffset(FOUR_BYTE_COPY_TAG, 4, len, offset);
    255    }
    256
    257    private void writeBackReferenceWithLittleEndianOffset(final int tag, final int offsetBytes, final int len, final int offset)
    258        throws IOException {
    259        os.write(tag | ((len - 1) << 2));
    260        writeLittleEndian(offsetBytes, offset);
    261    }
    262
    263    // technically the format could use shorter matches but with a
    264    // length of three the offset would be encoded as at least two
    265    // bytes in addition to the tag, so yield no compression at all
    266    private static final int MIN_MATCH_LENGTH = 4;
    267    // Snappy stores the match length in six bits of the tag
    268    private static final int MAX_MATCH_LENGTH = 64;
    269
    270    /**
    271     * Returns a builder correctly configured for the Snappy algorithm using the gven block size.
    272     * @param blockSize the block size.
    273     * @return a builder correctly configured for the Snappy algorithm using the gven block size
    274     */
    275    public static Parameters.Builder createParameterBuilder(final int blockSize) {
    276        // the max offset and max literal length defined by the format
    277        // are 2^32 - 1 and 2^32 respectively - with blockSize being
    278        // an integer we will never exceed that
    279        return Parameters.builder(blockSize)
    280            .withMinBackReferenceLength(MIN_MATCH_LENGTH)
    281            .withMaxBackReferenceLength(MAX_MATCH_LENGTH)
    282            .withMaxOffset(blockSize)
    283            .withMaxLiteralLength(blockSize);
    284    }
    285}