AssetRipper_AssetRipper/Source/AssetRipper.IO.Files/BundleFiles/LzmaCompression.cs

using AssetRipper.IO.Files.Exceptions;
using SharpCompress.Compressors.LZMA;
using System.Buffers;
using System.Buffers.Binary;

namespace AssetRipper.IO.Files.BundleFiles;

public static class LzmaCompression
{
	/// <summary>
	/// Read LZMA properties and decompress LZMA data
	/// </summary>
	/// <param name="compressedStream">LZMA compressed stream</param>
	/// <param name="compressedSize">Compressed data length</param>
	/// <param name="decompressedStream">Stream for decompressed output</param>
	/// <param name="decompressedSize">Decompressed data length</param>
	public static void DecompressLzmaStream(Stream compressedStream, long compressedSize, Stream decompressedStream, long decompressedSize)
	{
		byte[] properties = new byte[PropertiesSize];
		long basePosition = compressedStream.Position;

		compressedStream.ReadExactly(properties, 0, PropertiesSize);

		long headSize = compressedStream.Position - basePosition;
		long headlessSize = compressedSize - headSize;

		DecompressLzmaStream(properties, compressedStream, headlessSize, decompressedStream, decompressedSize);

		if (compressedStream.Position > basePosition + compressedSize)
		{
			DecompressionFailedException.ThrowReadMoreThanExpected(CompressionType.Lzma, compressedSize, compressedStream.Position - basePosition);
		}
		compressedStream.Position = basePosition + compressedSize;
	}

	/// <summary>
	/// Read LZMA properties and decompressed size and decompress LZMA data
	/// </summary>
	/// <param name="compressedStream">LZMA compressed stream</param>
	/// <param name="compressedSize">Compressed data length</param>
	/// <param name="decompressedStream">Stream for decompressed output</param>
	public static void DecompressLzmaSizeStream(Stream compressedStream, long compressedSize, Stream decompressedStream)
	{
		byte[] properties = new byte[PropertiesSize]; //GetBuffer();
		byte[] sizeBytes = new byte[UncompressedSize]; //GetBuffer();
		long basePosition = compressedStream.Position;

		compressedStream.ReadExactly(properties, 0, PropertiesSize);
		compressedStream.ReadExactly(sizeBytes, 0, UncompressedSize);
		long decompressedSize = BinaryPrimitives.ReadInt64LittleEndian(sizeBytes);

		long headSize = compressedStream.Position - basePosition;
		long headlessSize = compressedSize - headSize;

		DecompressLzmaStream(properties, compressedStream, headlessSize, decompressedStream, decompressedSize);

		if (compressedStream.Position > basePosition + compressedSize)
		{
			DecompressionFailedException.ThrowReadMoreThanExpected(CompressionType.Lzma, compressedSize, compressedStream.Position - basePosition);
		}
		compressedStream.Position = basePosition + compressedSize;
	}

	private static void DecompressLzmaStream(byte[] properties, Stream compressedStream, long headlessSize, Stream decompressedStream, long decompressedSize)
	{
		LzmaStream lzmaStream = new LzmaStream(properties, compressedStream, headlessSize, -1, null, false);

		byte[] buffer = ArrayPool<byte>.Shared.Rent(1024);
		long totalRead = 0;
		while (totalRead < decompressedSize)
		{
			int toRead = (int)Math.Min(buffer.Length, decompressedSize - totalRead);
			int read = lzmaStream.Read(buffer, 0, toRead);
			if (read > 0)
			{
				decompressedStream.Write(buffer, 0, read);
				totalRead += read;
			}
			else
			{
				break;
			}
		}
		ArrayPool<byte>.Shared.Return(buffer);
	}

	/// <summary>
	/// Compress some data with LZMA.
	/// </summary>
	/// <param name="uncompressedStream">The source stream with uncompressed data.</param>
	/// <param name="uncompressedSize">The number of bytes to read from <paramref name="uncompressedSize"/>.</param>
	/// <param name="compressedStream">The stream in which to write the compressed data.</param>
	/// <returns>The number of compressed bytes written to <paramref name="compressedStream"/> including the 5 property bytes.</returns>
	public static long CompressLzmaStream(Stream uncompressedStream, long uncompressedSize, Stream compressedStream)
	{
		long basePosition = compressedStream.Position;
		LzmaStream lzmaStream = new LzmaStream(new(), false, compressedStream);
		compressedStream.Write(lzmaStream.Properties);
		CopyToLzma(uncompressedStream, lzmaStream, uncompressedSize);
		lzmaStream.Close();
		return compressedStream.Position - basePosition;
	}

	/// <summary>
	/// Compress some data with LZMA.
	/// </summary>
	/// <param name="uncompressedStream">The source stream with uncompressed data.</param>
	/// <param name="uncompressedSize">The number of bytes to read from <paramref name="uncompressedSize"/>.</param>
	/// <param name="compressedStream">The stream in which to write the compressed data.</param>
	/// <returns>
	/// The number of compressed bytes written to <paramref name="compressedStream"/> including the 5 property bytes
	/// and <see langword="long"/> uncompressed size value.
	/// </returns>
	public static long CompressLzmaSizeStream(Stream uncompressedStream, long uncompressedSize, Stream compressedStream)
	{
		long basePosition = compressedStream.Position;
		LzmaStream lzmaStream = new LzmaStream(new(), false, compressedStream);
		compressedStream.Write(lzmaStream.Properties);
		new BinaryWriter(compressedStream).Write(uncompressedSize);
		CopyToLzma(uncompressedStream, lzmaStream, uncompressedSize);
		lzmaStream.Close();
		return compressedStream.Position - basePosition;
	}

	private static void CopyToLzma(Stream inputStream, LzmaStream lzmaStream, long uncompressedSize)
	{
		byte[] buffer = ArrayPool<byte>.Shared.Rent(1024);
		long totalCopied = 0;
		while (totalCopied < uncompressedSize)
		{
			int read = inputStream.Read(buffer, 0, (int)Math.Min(buffer.Length, uncompressedSize - totalCopied));
			if (read == 0)
			{
				throw new EndOfStreamException();
			}
			lzmaStream.Write(buffer, 0, read);
			totalCopied += read;
		}
	}

	private const int PropertiesSize = 5;
	private const int UncompressedSize = sizeof(long);
}