/*
 * 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 io.prestosql.parquet;

import io.airlift.compress.Decompressor;
import io.airlift.compress.lz4.Lz4Decompressor;
import io.airlift.compress.lzo.LzoDecompressor;
import io.airlift.compress.snappy.SnappyDecompressor;
import io.airlift.compress.zstd.ZstdDecompressor;
import io.airlift.slice.DynamicSliceOutput;
import io.airlift.slice.Slice;
import org.apache.parquet.hadoop.metadata.CompressionCodecName;

import java.io.IOException;
import java.io.InputStream;
import java.util.zip.GZIPInputStream;

import static com.google.common.base.Preconditions.checkArgument;
import static io.airlift.slice.SizeOf.SIZE_OF_INT;
import static io.airlift.slice.SizeOf.SIZE_OF_LONG;
import static io.airlift.slice.Slices.EMPTY_SLICE;
import static io.airlift.slice.Slices.wrappedBuffer;
import static java.util.Objects.requireNonNull;
import static sun.misc.Unsafe.ARRAY_BYTE_BASE_OFFSET;

public final class ParquetCompressionUtils
{
    private static final int GZIP_BUFFER_SIZE = 8 * 1024;

    private ParquetCompressionUtils() {}

    public static Slice decompress(CompressionCodecName codec, Slice input, int uncompressedSize)
            throws IOException
    {
        requireNonNull(input, "input is null");

        if (input.length() == 0) {
            return EMPTY_SLICE;
        }

        switch (codec) {
            case GZIP:
                return decompressGzip(input, uncompressedSize);
            case SNAPPY:
                return decompressSnappy(input, uncompressedSize);
            case UNCOMPRESSED:
                return input;
            case LZO:
                return decompressLZO(input, uncompressedSize);
            case LZ4:
                return decompressLz4(input, uncompressedSize);
            case ZSTD:
                return decompressZstd(input, uncompressedSize);
            default:
                throw new ParquetCorruptionException("Codec not supported in Parquet: " + codec);
        }
    }

    private static Slice decompressSnappy(Slice input, int uncompressedSize)
    {
        byte[] buffer = new byte[uncompressedSize];
        decompress(new SnappyDecompressor(), input, 0, input.length(), buffer, 0);
        return wrappedBuffer(buffer);
    }

    private static Slice decompressZstd(Slice input, int uncompressedSize)
    {
        byte[] buffer = new byte[uncompressedSize];
        decompress(new ZstdDecompressor(), input, 0, input.length(), buffer, 0);
        return wrappedBuffer(buffer);
    }

    private static Slice decompressGzip(Slice input, int uncompressedSize)
            throws IOException
    {
        if (uncompressedSize == 0) {
            return EMPTY_SLICE;
        }

        DynamicSliceOutput sliceOutput = new DynamicSliceOutput(uncompressedSize);
        byte[] buffer = new byte[uncompressedSize];
        try (InputStream gzipInputStream = new GZIPInputStream(input.getInput(), GZIP_BUFFER_SIZE)) {
            int bytesRead;
            while ((bytesRead = gzipInputStream.read(buffer)) != -1) {
                sliceOutput.write(buffer, 0, bytesRead);
            }
            return sliceOutput.getUnderlyingSlice();
        }
    }

    private static Slice decompressLz4(Slice input, int uncompressedSize)
    {
        Lz4Decompressor decompressor = new Lz4Decompressor();
        long totalDecompressedCount = 0;
        // over allocate buffer which makes decompression easier
        byte[] output = new byte[uncompressedSize + SIZE_OF_LONG];
        int outputOffset = 0;
        int inputOffset = 0;
        int cumulativeUncompressedBlockLength = 0;

        while (totalDecompressedCount < uncompressedSize) {
            if (totalDecompressedCount == cumulativeUncompressedBlockLength) {
                cumulativeUncompressedBlockLength += Integer.reverseBytes(input.getInt(inputOffset));
                inputOffset += SIZE_OF_INT;
            }
            int compressedChunkLength = Integer.reverseBytes(input.getInt(inputOffset));
            inputOffset += SIZE_OF_INT;
            int decompressionSize = decompress(decompressor, input, inputOffset, compressedChunkLength, output, outputOffset);
            totalDecompressedCount += decompressionSize;
            outputOffset += decompressionSize;
            inputOffset += compressedChunkLength;
        }
        checkArgument(outputOffset == uncompressedSize);
        return wrappedBuffer(output, 0, uncompressedSize);
    }

    private static Slice decompressLZO(Slice input, int uncompressedSize)
    {
        LzoDecompressor lzoDecompressor = new LzoDecompressor();
        long totalDecompressedCount = 0;
        // over allocate buffer which makes decompression easier
        byte[] output = new byte[uncompressedSize + SIZE_OF_LONG];
        int outputOffset = 0;
        int inputOffset = 0;
        int cumulativeUncompressedBlockLength = 0;

        while (totalDecompressedCount < uncompressedSize) {
            if (totalDecompressedCount == cumulativeUncompressedBlockLength) {
                cumulativeUncompressedBlockLength += Integer.reverseBytes(input.getInt(inputOffset));
                inputOffset += SIZE_OF_INT;
            }
            int compressedChunkLength = Integer.reverseBytes(input.getInt(inputOffset));
            inputOffset += SIZE_OF_INT;
            int decompressionSize = decompress(lzoDecompressor, input, inputOffset, compressedChunkLength, output, outputOffset);
            totalDecompressedCount += decompressionSize;
            outputOffset += decompressionSize;
            inputOffset += compressedChunkLength;
        }
        checkArgument(outputOffset == uncompressedSize);
        return wrappedBuffer(output, 0, uncompressedSize);
    }

    private static int decompress(Decompressor decompressor, Slice input, int inputOffset, int inputLength, byte[] output, int outputOffset)
    {
        byte[] byteArray = (byte[]) input.getBase();
        int byteArrayOffset = inputOffset + (int) (input.getAddress() - ARRAY_BYTE_BASE_OFFSET);
        int size = decompressor.decompress(byteArray, byteArrayOffset, inputLength, output, outputOffset, output.length - outputOffset);
        return size;
    }
}

@@ -3781,31 +3781,22 @@
         
-Lz4Decompressor
+return
  decompr
@@ -3794,29 +3794,31 @@
 n decompress
-or = 
+Framed(
 new Lz4Decom
@@ -3830,770 +3830,163 @@
 or()
-;%0A        long totalDecompressedCount = 0;%0A        // over allocate buffer which makes decompression easier%0A        byte%5B%5D output = new byte%5BuncompressedSize + SIZE_OF_LONG%5D;%0A        int outputOffset = 0;%0A        int inputOffset = 0;%0A        int cumulativeUncompressedBlockLength = 0;%0A%0A        while (totalDecompressedCount %3C uncompressedSize) %7B%0A            if (totalDecompressedCount == cumulativeUncompressedBlockLength) %7B%0A                cumulativeUncompressedBlockLength += Integer.reverseBytes(input.getInt(inputOffset));%0A                inputOffset += SIZE_OF_INT;%0A            %7D%0A            int compressedChunkLength = Integer.reverseBytes(input.getInt(inputOffset));%0A            inputOffset += SIZE_OF_INT;%0A            int decompressionSize = decompress(d
+, input, uncompressedSize);%0A    %7D%0A%0A    private static Slice decompressLZO(Slice input, int uncompressedSize)%0A    %7B%0A        return decompressFramed(new LzoD
 ecom
@@ -3984,32 +3984,34 @@
  LzoDecompressor
+()
 , input, inputOf
@@ -4007,498 +4007,143 @@
 ut, 
-inputOffset, compressedChunkLength, output, outputOffset);%0A            totalDecompressedCount += decompression
+uncompressed
 Size
+)
 ;%0A    
-        outputOffset += decompressionSize;%0A            inputOffset += compressedChunkLength;%0A        %7D%0A        checkArgument(outputOffset == uncompressedSize);%0A        return wrappedBuffer(output, 0, uncompressedSize);%0A    %7D%0A%0A    private static Slice decompressLZO(Slice input, int uncompressedSize)%0A    %7B%0A        LzoDecompressor lzoDecompressor = new LzoDecompressor();
+%7D%0A%0A    private static Slice decompressFramed(Decompressor decompressor, Slice input, int uncompressedSize)%0A    %7B
 %0A   
@@ -4902,12 +4902,9 @@
 ess(
-lzoD
+d
 ecom

Remove duplicate code in ParquetCompressionUtils