1//************************************ bs::framework - Copyright 2018 Marko Pintera **************************************//
2//*********** Licensed under the MIT license. See LICENSE.md for full terms. This notice is not to be removed. ***********//
3#include "Utility/BsCompression.h"
4#include "FileSystem/BsDataStream.h"
5
6// Third party
7#include "snappy.h"
8#include "snappy-sinksource.h"
9#include "Debug/BsDebug.h"
10
11namespace bs
12{
13 /** Source accepting a data stream. Used for Snappy compression library. */
14 class DataStreamSource : public snappy::Source
15 {
16 public:
17 DataStreamSource(const SPtr<DataStream>& stream, std::function<void(float)> reportProgress = nullptr)
18 :mStream(stream), mReportProgress(std::move(reportProgress))
19 {
20 mTotal = mStream->size() - mStream->tell();
21 mRemaining = mTotal;
22
23 if (mStream->isFile())
24 mReadBuffer = (char*)bs_alloc(32768);
25 }
26
27 virtual ~DataStreamSource()
28 {
29 if (mReadBuffer != nullptr)
30 bs_free(mReadBuffer);
31 }
32
33 size_t Available() const override
34 {
35 return mRemaining;
36 }
37
38 const char* Peek(size_t* len) override
39 {
40 if (!mStream->isFile())
41 {
42 SPtr<MemoryDataStream> memStream = std::static_pointer_cast<MemoryDataStream>(mStream);
43
44 *len = Available();
45 return (char*)memStream->getPtr() + mBufferOffset;
46 }
47 else
48 {
49 while (mBufferOffset >= mReadBufferContentSize)
50 {
51 mBufferOffset -= mReadBufferContentSize;
52 mReadBufferContentSize = mStream->read(mReadBuffer, 32768);
53
54 if (mReadBufferContentSize == 0)
55 break;
56 }
57
58 *len = mReadBufferContentSize - mBufferOffset;
59 return (char*)(mReadBuffer + mBufferOffset);
60 }
61 }
62
63 void Skip(size_t n) override
64 {
65 mBufferOffset += n;
66 mRemaining -= n;
67
68 if(mReportProgress)
69 mReportProgress(1.0f - mRemaining / (float)mTotal);
70 }
71 private:
72 SPtr<DataStream> mStream;
73 std::function<void(float)> mReportProgress;
74
75 size_t mRemaining;
76 size_t mTotal;
77 size_t mBufferOffset = 0;
78
79 // File streams only
80 char* mReadBuffer = nullptr;
81 size_t mReadBufferContentSize = 0;
82 };
83
84 /** Sink (destination) accepting a data stream. Used for Snappy compression library. */
85 class DataStreamSink : public snappy::Sink
86 {
87 struct BufferPiece
88 {
89 char* buffer;
90 size_t size;
91 };
92
93 public:
94 DataStreamSink() = default;
95 virtual ~DataStreamSink()
96 {
97 for (auto& entry : mBufferPieces)
98 bs_free(entry.buffer);
99 }
100
101 void Append(const char* data, size_t n) override
102 {
103 if(mBufferPieces.empty() || mBufferPieces.back().buffer != data)
104 {
105 BufferPiece piece;
106 piece.buffer = (char*)bs_alloc((UINT32)n);
107 piece.size = n;
108
109 memcpy(piece.buffer, data, n);
110 mBufferPieces.push_back(piece);
111 }
112 else
113 {
114 BufferPiece& piece = mBufferPieces.back();
115 assert(piece.buffer == data);
116
117 piece.size = n;
118 }
119 }
120
121 char* GetAppendBuffer(size_t len, char* scratch) override
122 {
123 BufferPiece piece;
124 piece.buffer = (char*)bs_alloc((UINT32)len);
125 piece.size = 0;
126
127 mBufferPieces.push_back(piece);
128 return piece.buffer;
129 }
130
131 char* GetAppendBufferVariable(size_t min_size, size_t desired_size_hint, char* scratch, size_t scratch_size,
132 size_t* allocated_size) override
133 {
134 BufferPiece piece;
135 piece.buffer = (char*)bs_alloc((UINT32)desired_size_hint);
136 piece.size = 0;
137
138 mBufferPieces.push_back(piece);
139
140 *allocated_size = desired_size_hint;
141 return piece.buffer;
142 }
143
144 void AppendAndTakeOwnership(char* bytes, size_t n, void(*deleter)(void*, const char*, size_t),
145 void *deleter_arg) override
146 {
147 BufferPiece& piece = mBufferPieces.back();
148
149 if (piece.buffer != bytes)
150 {
151 memcpy(piece.buffer, bytes, n);
152 (*deleter)(deleter_arg, bytes, n);
153 }
154
155 piece.size = n;
156 }
157
158 SPtr<MemoryDataStream> GetOutput()
159 {
160 size_t totalSize = 0;
161 for (auto& entry : mBufferPieces)
162 totalSize += entry.size;
163
164 SPtr<MemoryDataStream> ds = bs_shared_ptr_new<MemoryDataStream>(totalSize);
165 for (auto& entry : mBufferPieces)
166 ds->write(entry.buffer, entry.size);
167
168 ds->seek(0);
169 return ds;
170 }
171
172 private:
173 Vector<BufferPiece> mBufferPieces;
174 };
175
176 SPtr<MemoryDataStream> Compression::compress(SPtr<DataStream>& input, std::function<void(float)> reportProgress)
177 {
178 DataStreamSource src(input, std::move(reportProgress));
179 DataStreamSink dst;
180
181 size_t bytesWritten = snappy::Compress(&src, &dst);
182 SPtr<MemoryDataStream> output = dst.GetOutput();
183 assert(output->size() == bytesWritten);
184
185 return output;
186 }
187
188 SPtr<MemoryDataStream> Compression::decompress(SPtr<DataStream>& input, std::function<void(float)> reportProgress)
189 {
190 DataStreamSource src(input, std::move(reportProgress));
191 DataStreamSink dst;
192
193 if (!snappy::Uncompress(&src, &dst))
194 {
195 LOGERR("Decompression failed, corrupt data.");
196 return nullptr;
197 }
198
199 return dst.GetOutput();
200 }
201}
202