BansheeEngine/Source/BansheeUtility/Utility/BsCompression.cpp

//********************************** Banshee Engine (www.banshee3d.com) **************************************************//
//**************** Copyright (c) 2016 Marko Pintera ([email protected]). All rights reserved. **********************//
#include "Utility/BsCompression.h"
#include "FileSystem/BsDataStream.h"

// Third party
#include "snappy.h"
#include "snappy-sinksource.h"
#include "Debug/BsDebug.h"

namespace bs
{
	/** Source accepting a data stream. Used for Snappy compression library. */
	class DataStreamSource : public snappy::Source
	{
	public:
		DataStreamSource(const SPtr<DataStream>& stream)
			: mStream(stream), mReadBuffer(nullptr), mReadBufferContentSize(0), mBufferOffset(0)
		{
			mRemaining = mStream->size() - mStream->tell();

			if (mStream->isFile())
				mReadBuffer = (char*)bs_alloc(2048);
		}

		virtual ~DataStreamSource()
		{
			if (mReadBuffer != nullptr)
				bs_free(mReadBuffer);
		}

		size_t Available() const override
		{
			return mRemaining;
		}

		const char* Peek(size_t* len) override
		{
			if (!mStream->isFile())
			{
				SPtr<MemoryDataStream> memStream = std::static_pointer_cast<MemoryDataStream>(mStream);

				*len = Available();
				return (char*)memStream->getPtr() + mBufferOffset;
			}
			else
			{
				while (mBufferOffset >= mReadBufferContentSize)
				{
					mBufferOffset -= mReadBufferContentSize;
					mReadBufferContentSize = mStream->read(mReadBuffer, 2048);

					if (mReadBufferContentSize == 0)
						break;
				}

				*len = mReadBufferContentSize - mBufferOffset;
				return (char*)(mReadBuffer + mBufferOffset);
			}
		}

		void Skip(size_t n) override
		{
			mBufferOffset += n;
			mRemaining -= n;
		}
	private:
		SPtr<DataStream> mStream;

		// File streams only
		char* mReadBuffer;
		size_t mReadBufferContentSize;
		size_t mRemaining;
		size_t mBufferOffset;
	};

	/** Sink (destination) accepting a data stream. Used for Snappy compression library. */
	class DataStreamSink : public snappy::Sink
	{
		struct BufferPiece
		{
			char* buffer;
			size_t size;
		};

	public:
		DataStreamSink() { }
		virtual ~DataStreamSink()
		{
			for (auto& entry : mBufferPieces)
				bs_free(entry.buffer);
		}

		void Append(const char* data, size_t n) override
		{
			if(mBufferPieces.size() == 0 || mBufferPieces.back().buffer != data)
			{
				BufferPiece piece;
				piece.buffer = (char*)bs_alloc((UINT32)n);
				piece.size = n;

				memcpy(piece.buffer, data, n);
				mBufferPieces.push_back(piece);
			}
			else
			{
				BufferPiece& piece = mBufferPieces.back();
				assert(piece.buffer == data);

				piece.size = n;
			}
		}

		char* GetAppendBuffer(size_t len, char* scratch) override
		{
			BufferPiece piece;
			piece.buffer = (char*)bs_alloc((UINT32)len);
			piece.size = 0;

			mBufferPieces.push_back(piece);
			return piece.buffer;
		}

		char* GetAppendBufferVariable(size_t min_size, size_t desired_size_hint, char* scratch, size_t scratch_size,
			size_t* allocated_size) override
		{
			BufferPiece piece;
			piece.buffer = (char*)bs_alloc((UINT32)desired_size_hint);
			piece.size = 0;

			mBufferPieces.push_back(piece);

			*allocated_size = desired_size_hint;
			return piece.buffer;
		}

		void AppendAndTakeOwnership(char* bytes, size_t n, void(*deleter)(void*, const char*, size_t),
			void *deleter_arg) override
		{
			BufferPiece& piece = mBufferPieces.back();

			if (piece.buffer != bytes)
			{
				memcpy(piece.buffer, bytes, n);
				(*deleter)(deleter_arg, bytes, n);
			}

			piece.size = n;
		}

		SPtr<MemoryDataStream> GetOutput()
		{
			size_t totalSize = 0;
			for (auto& entry : mBufferPieces)
				totalSize += entry.size;

			SPtr<MemoryDataStream> ds = bs_shared_ptr_new<MemoryDataStream>(totalSize);
			for (auto& entry : mBufferPieces)
				ds->write(entry.buffer, entry.size);

			ds->seek(0);
			return ds;
		}

	private:
		Vector<BufferPiece> mBufferPieces;
	};

	SPtr<MemoryDataStream> Compression::compress(SPtr<DataStream>& input)
	{
		DataStreamSource src(input);
		DataStreamSink dst;

		size_t bytesWritten = snappy::Compress(&src, &dst);
		SPtr<MemoryDataStream> output = dst.GetOutput();
		assert(output->size() == bytesWritten);

		return output;
	}

	SPtr<MemoryDataStream> Compression::decompress(SPtr<DataStream>& input)
	{
		DataStreamSource src(input);
		DataStreamSink dst;

		if (!snappy::Uncompress(&src, &dst))
		{
			LOGERR("Decompression failed, corrupt data.");
			return nullptr;
		}

		return dst.GetOutput();
	}
}