puffin_stream.cc 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483
  1. // Copyright 2017 The Chromium OS Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style license that can be
  3. // found in the LICENSE file.
  4. #include "puffin/src/puffin_stream.h"
  5. #include <algorithm>
  6. #include <memory>
  7. #include <string>
  8. #include <utility>
  9. #include <vector>
  10. #include "puffin/src/bit_reader.h"
  11. #include "puffin/src/bit_writer.h"
  12. #include "puffin/src/include/puffin/common.h"
  13. #include "puffin/src/include/puffin/huffer.h"
  14. #include "puffin/src/include/puffin/puffer.h"
  15. #include "puffin/src/include/puffin/stream.h"
  16. #include "puffin/src/logging.h"
  17. #include "puffin/src/puff_reader.h"
  18. #include "puffin/src/puff_writer.h"
  19. using std::shared_ptr;
  20. using std::unique_ptr;
  21. using std::vector;
  22. namespace puffin {
  23. namespace {
  24. bool CheckArgsIntegrity(uint64_t puff_size,
  25. const vector<BitExtent>& deflates,
  26. const vector<ByteExtent>& puffs) {
  27. TEST_AND_RETURN_FALSE(puffs.size() == deflates.size());
  28. // Check if the |puff_size| is actually greater than the last byte of the last
  29. // puff in |puffs|.
  30. if (!puffs.empty()) {
  31. TEST_AND_RETURN_FALSE(puff_size >=
  32. puffs.back().offset + puffs.back().length);
  33. }
  34. // Check to make sure |puffs| and |deflates| are sorted and non-overlapping.
  35. auto is_overlap = [](const auto& a, const auto& b) {
  36. return (a.offset + a.length) > b.offset;
  37. };
  38. TEST_AND_RETURN_FALSE(deflates.end() == std::adjacent_find(deflates.begin(),
  39. deflates.end(),
  40. is_overlap));
  41. TEST_AND_RETURN_FALSE(puffs.end() == std::adjacent_find(puffs.begin(),
  42. puffs.end(),
  43. is_overlap));
  44. return true;
  45. }
  46. } // namespace
  47. UniqueStreamPtr PuffinStream::CreateForPuff(UniqueStreamPtr stream,
  48. shared_ptr<Puffer> puffer,
  49. uint64_t puff_size,
  50. const vector<BitExtent>& deflates,
  51. const vector<ByteExtent>& puffs,
  52. size_t max_cache_size) {
  53. TEST_AND_RETURN_VALUE(CheckArgsIntegrity(puff_size, deflates, puffs),
  54. nullptr);
  55. TEST_AND_RETURN_VALUE(stream->Seek(0), nullptr);
  56. UniqueStreamPtr puffin_stream(new PuffinStream(std::move(stream), puffer,
  57. nullptr, puff_size, deflates,
  58. puffs, max_cache_size));
  59. TEST_AND_RETURN_VALUE(puffin_stream->Seek(0), nullptr);
  60. return puffin_stream;
  61. }
  62. UniqueStreamPtr PuffinStream::CreateForHuff(UniqueStreamPtr stream,
  63. shared_ptr<Huffer> huffer,
  64. uint64_t puff_size,
  65. const vector<BitExtent>& deflates,
  66. const vector<ByteExtent>& puffs) {
  67. TEST_AND_RETURN_VALUE(CheckArgsIntegrity(puff_size, deflates, puffs),
  68. nullptr);
  69. TEST_AND_RETURN_VALUE(stream->Seek(0), nullptr);
  70. UniqueStreamPtr puffin_stream(new PuffinStream(
  71. std::move(stream), nullptr, huffer, puff_size, deflates, puffs, 0));
  72. TEST_AND_RETURN_VALUE(puffin_stream->Seek(0), nullptr);
  73. return puffin_stream;
  74. }
  75. PuffinStream::PuffinStream(UniqueStreamPtr stream,
  76. shared_ptr<Puffer> puffer,
  77. shared_ptr<Huffer> huffer,
  78. uint64_t puff_size,
  79. const vector<BitExtent>& deflates,
  80. const vector<ByteExtent>& puffs,
  81. size_t max_cache_size)
  82. : stream_(std::move(stream)),
  83. puffer_(puffer),
  84. huffer_(huffer),
  85. puff_stream_size_(puff_size),
  86. deflates_(deflates),
  87. puffs_(puffs),
  88. is_for_puff_(puffer_ ? true : false),
  89. max_cache_size_(max_cache_size) {
  90. // Building upper bounds for faster seek.
  91. upper_bounds_.reserve(puffs.size());
  92. for (const auto& puff : puffs) {
  93. upper_bounds_.emplace_back(puff.offset + puff.length);
  94. }
  95. upper_bounds_.emplace_back(puff_stream_size_ + 1);
  96. // We can pass the size of the deflate stream too, but it is not necessary
  97. // yet. We cannot get the size of stream from itself, because we might be
  98. // writing into it and its size is not defined yet.
  99. uint64_t deflate_stream_size = puff_stream_size_;
  100. if (!puffs.empty()) {
  101. deflate_stream_size =
  102. ((deflates.back().offset + deflates.back().length) / 8) +
  103. puff_stream_size_ - (puffs.back().offset + puffs.back().length);
  104. }
  105. deflates_.emplace_back(deflate_stream_size * 8, 0);
  106. puffs_.emplace_back(puff_stream_size_, 0);
  107. // Look for the largest puff and deflate extents and get proper size buffers.
  108. uint64_t max_puff_length = 0;
  109. for (const auto& puff : puffs) {
  110. max_puff_length = std::max(max_puff_length, puff.length);
  111. }
  112. puff_buffer_ = std::make_shared<Buffer>(max_puff_length + 1);
  113. if (max_cache_size_ < max_puff_length) {
  114. max_cache_size_ = 0; // It means we are not caching puffs.
  115. }
  116. uint64_t max_deflate_length = 0;
  117. for (const auto& deflate : deflates) {
  118. max_deflate_length = std::max(max_deflate_length, deflate.length * 8);
  119. }
  120. deflate_buffer_ = std::make_unique<Buffer>(max_deflate_length + 2);
  121. }
  122. bool PuffinStream::GetSize(uint64_t* size) {
  123. *size = puff_stream_size_;
  124. return true;
  125. }
  126. bool PuffinStream::GetOffset(uint64_t* offset) {
  127. *offset = puff_pos_ + skip_bytes_;
  128. return true;
  129. }
  130. bool PuffinStream::Seek(uint64_t offset) {
  131. TEST_AND_RETURN_FALSE(!closed_);
  132. if (!is_for_puff_) {
  133. // For huffing we should not seek, only seek to zero is accepted.
  134. TEST_AND_RETURN_FALSE(offset == 0);
  135. }
  136. TEST_AND_RETURN_FALSE(offset <= puff_stream_size_);
  137. // We are searching first available puff which either includes the |offset| or
  138. // it is the next available puff after |offset|.
  139. auto next_puff_iter =
  140. std::upper_bound(upper_bounds_.begin(), upper_bounds_.end(), offset);
  141. TEST_AND_RETURN_FALSE(next_puff_iter != upper_bounds_.end());
  142. auto next_puff_idx = std::distance(upper_bounds_.begin(), next_puff_iter);
  143. cur_puff_ = std::next(puffs_.begin(), next_puff_idx);
  144. cur_deflate_ = std::next(deflates_.begin(), next_puff_idx);
  145. if (offset < cur_puff_->offset) {
  146. // between two puffs.
  147. puff_pos_ = offset;
  148. auto back_track_bytes = cur_puff_->offset - puff_pos_;
  149. deflate_bit_pos_ = ((cur_deflate_->offset + 7) / 8 - back_track_bytes) * 8;
  150. if (cur_puff_ != puffs_.begin()) {
  151. auto prev_deflate = std::prev(cur_deflate_);
  152. if (deflate_bit_pos_ < prev_deflate->offset + prev_deflate->length) {
  153. deflate_bit_pos_ = prev_deflate->offset + prev_deflate->length;
  154. }
  155. }
  156. } else {
  157. // Inside a puff.
  158. puff_pos_ = cur_puff_->offset;
  159. deflate_bit_pos_ = cur_deflate_->offset;
  160. }
  161. skip_bytes_ = offset - puff_pos_;
  162. if (!is_for_puff_ && offset == 0) {
  163. TEST_AND_RETURN_FALSE(stream_->Seek(0));
  164. TEST_AND_RETURN_FALSE(SetExtraByte());
  165. }
  166. return true;
  167. }
  168. bool PuffinStream::Close() {
  169. closed_ = true;
  170. return stream_->Close();
  171. }
  172. bool PuffinStream::Read(void* buffer, size_t count) {
  173. TEST_AND_RETURN_FALSE(!closed_);
  174. TEST_AND_RETURN_FALSE(is_for_puff_);
  175. if (cur_puff_ == puffs_.end()) {
  176. TEST_AND_RETURN_FALSE(count == 0);
  177. }
  178. auto bytes = static_cast<uint8_t*>(buffer);
  179. uint64_t length = count;
  180. uint64_t bytes_read = 0;
  181. while (bytes_read < length) {
  182. if (puff_pos_ < cur_puff_->offset) {
  183. // Reading between two deflates. We also read bytes that have at least one
  184. // bit of a deflate bit stream. The byte which has both deflate and raw
  185. // data will be shifted or masked off the deflate bits and the remaining
  186. // value will be saved in the puff stream as an byte integer.
  187. uint64_t start_byte = (deflate_bit_pos_ / 8);
  188. uint64_t end_byte = (cur_deflate_->offset + 7) / 8;
  189. auto bytes_to_read = std::min(length - bytes_read, end_byte - start_byte);
  190. TEST_AND_RETURN_FALSE(bytes_to_read >= 1);
  191. TEST_AND_RETURN_FALSE(stream_->Seek(start_byte));
  192. TEST_AND_RETURN_FALSE(stream_->Read(bytes + bytes_read, bytes_to_read));
  193. // If true, we read the first byte of the curret deflate. So we have to
  194. // mask out the deflate bits (which are most significant bits.)
  195. if ((start_byte + bytes_to_read) * 8 > cur_deflate_->offset) {
  196. bytes[bytes_read + bytes_to_read - 1] &=
  197. (1 << (cur_deflate_->offset & 7)) - 1;
  198. }
  199. // If true, we read the last byte of the previous deflate and we have to
  200. // shift it out. The least significat bits belongs to the deflate
  201. // stream. The order of these last two conditions are important because a
  202. // byte can contain a finishing deflate and a starting deflate with some
  203. // bits between them so we have to modify correctly. Keep in mind that in
  204. // this situation both are modifying the same byte.
  205. if (start_byte * 8 < deflate_bit_pos_) {
  206. bytes[bytes_read] >>= deflate_bit_pos_ & 7;
  207. }
  208. // Pass |deflate_bit_pos_| for all the read bytes.
  209. deflate_bit_pos_ -= deflate_bit_pos_ & 7;
  210. deflate_bit_pos_ += bytes_to_read * 8;
  211. if (deflate_bit_pos_ > cur_deflate_->offset) {
  212. // In case it reads into the first byte of the current deflate.
  213. deflate_bit_pos_ = cur_deflate_->offset;
  214. }
  215. bytes_read += bytes_to_read;
  216. puff_pos_ += bytes_to_read;
  217. TEST_AND_RETURN_FALSE(puff_pos_ <= cur_puff_->offset);
  218. } else {
  219. // Reading the deflate itself. We read all bytes including the first and
  220. // last byte (which may partially include a deflate bit). Here we keep the
  221. // |puff_pos_| point to the first byte of the puffed stream and
  222. // |skip_bytes_| shows how many bytes in the puff we have copied till now.
  223. auto start_byte = (cur_deflate_->offset / 8);
  224. auto end_byte = (cur_deflate_->offset + cur_deflate_->length + 7) / 8;
  225. auto bytes_to_read = end_byte - start_byte;
  226. // Puff directly to buffer if it has space.
  227. bool puff_directly_into_buffer =
  228. max_cache_size_ == 0 && (skip_bytes_ == 0) &&
  229. (length - bytes_read >= cur_puff_->length);
  230. auto cur_puff_idx = std::distance(puffs_.begin(), cur_puff_);
  231. if (max_cache_size_ == 0 ||
  232. !GetPuffCache(cur_puff_idx, cur_puff_->length, &puff_buffer_)) {
  233. // Did not find the puff buffer in cache. We have to build it.
  234. deflate_buffer_->resize(bytes_to_read);
  235. TEST_AND_RETURN_FALSE(stream_->Seek(start_byte));
  236. TEST_AND_RETURN_FALSE(
  237. stream_->Read(deflate_buffer_->data(), bytes_to_read));
  238. BufferBitReader bit_reader(deflate_buffer_->data(), bytes_to_read);
  239. BufferPuffWriter puff_writer(puff_directly_into_buffer
  240. ? bytes + bytes_read
  241. : puff_buffer_->data(),
  242. cur_puff_->length);
  243. // Drop the first unused bits.
  244. size_t extra_bits_len = cur_deflate_->offset & 7;
  245. TEST_AND_RETURN_FALSE(bit_reader.CacheBits(extra_bits_len));
  246. bit_reader.DropBits(extra_bits_len);
  247. TEST_AND_RETURN_FALSE(
  248. puffer_->PuffDeflate(&bit_reader, &puff_writer, nullptr));
  249. TEST_AND_RETURN_FALSE(bytes_to_read == bit_reader.Offset());
  250. TEST_AND_RETURN_FALSE(cur_puff_->length == puff_writer.Size());
  251. } else {
  252. // Just seek to proper location.
  253. TEST_AND_RETURN_FALSE(stream_->Seek(start_byte + bytes_to_read));
  254. }
  255. // Copy from puff buffer to output if needed.
  256. auto bytes_to_copy =
  257. std::min(length - bytes_read, cur_puff_->length - skip_bytes_);
  258. if (!puff_directly_into_buffer) {
  259. memcpy(bytes + bytes_read, puff_buffer_->data() + skip_bytes_,
  260. bytes_to_copy);
  261. }
  262. skip_bytes_ += bytes_to_copy;
  263. bytes_read += bytes_to_copy;
  264. // Move to next puff.
  265. if (puff_pos_ + skip_bytes_ == cur_puff_->offset + cur_puff_->length) {
  266. puff_pos_ += skip_bytes_;
  267. skip_bytes_ = 0;
  268. deflate_bit_pos_ = cur_deflate_->offset + cur_deflate_->length;
  269. cur_puff_++;
  270. cur_deflate_++;
  271. if (cur_puff_ == puffs_.end()) {
  272. break;
  273. }
  274. }
  275. }
  276. }
  277. TEST_AND_RETURN_FALSE(bytes_read == length);
  278. return true;
  279. }
  280. bool PuffinStream::Write(const void* buffer, size_t count) {
  281. TEST_AND_RETURN_FALSE(!closed_);
  282. TEST_AND_RETURN_FALSE(!is_for_puff_);
  283. auto bytes = static_cast<const uint8_t*>(buffer);
  284. uint64_t length = count;
  285. uint64_t bytes_wrote = 0;
  286. while (bytes_wrote < length) {
  287. if (deflate_bit_pos_ < (cur_deflate_->offset & ~7ull)) {
  288. // Between two puffs or before the first puff. We know that we are
  289. // starting from the byte boundary because we have already processed the
  290. // non-deflate bits of the last byte of the last deflate. Here we don't
  291. // process any byte that has deflate bit.
  292. TEST_AND_RETURN_FALSE((deflate_bit_pos_ & 7) == 0);
  293. auto copy_len =
  294. std::min((cur_deflate_->offset / 8) - (deflate_bit_pos_ / 8),
  295. length - bytes_wrote);
  296. TEST_AND_RETURN_FALSE(stream_->Write(bytes + bytes_wrote, copy_len));
  297. bytes_wrote += copy_len;
  298. puff_pos_ += copy_len;
  299. deflate_bit_pos_ += copy_len * 8;
  300. } else {
  301. // We are in a puff. We have to buffer incoming bytes until we reach the
  302. // size of the current puff so we can huff :). If the last bit of the
  303. // current deflate does not end in a byte boundary, then we have to read
  304. // one more byte to fill up the last byte of the deflate stream before
  305. // doing anything else.
  306. // |deflate_bit_pos_| now should be in the same byte as
  307. // |cur_deflate->offset|.
  308. if (deflate_bit_pos_ < cur_deflate_->offset) {
  309. last_byte_ |= bytes[bytes_wrote++] << (deflate_bit_pos_ & 7);
  310. skip_bytes_ = 0;
  311. deflate_bit_pos_ = cur_deflate_->offset;
  312. puff_pos_++;
  313. TEST_AND_RETURN_FALSE(puff_pos_ == cur_puff_->offset);
  314. }
  315. auto copy_len = std::min(length - bytes_wrote,
  316. cur_puff_->length + extra_byte_ - skip_bytes_);
  317. TEST_AND_RETURN_FALSE(puff_buffer_->size() >= skip_bytes_ + copy_len);
  318. memcpy(puff_buffer_->data() + skip_bytes_, bytes + bytes_wrote, copy_len);
  319. skip_bytes_ += copy_len;
  320. bytes_wrote += copy_len;
  321. if (skip_bytes_ == cur_puff_->length + extra_byte_) {
  322. // |puff_buffer_| is full, now huff into the |deflate_buffer_|.
  323. auto start_byte = cur_deflate_->offset / 8;
  324. auto end_byte = (cur_deflate_->offset + cur_deflate_->length + 7) / 8;
  325. auto bytes_to_write = end_byte - start_byte;
  326. deflate_buffer_->resize(bytes_to_write);
  327. BufferBitWriter bit_writer(deflate_buffer_->data(), bytes_to_write);
  328. BufferPuffReader puff_reader(puff_buffer_->data(), cur_puff_->length);
  329. // Write last byte if it has any.
  330. TEST_AND_RETURN_FALSE(
  331. bit_writer.WriteBits(cur_deflate_->offset & 7, last_byte_));
  332. last_byte_ = 0;
  333. TEST_AND_RETURN_FALSE(huffer_->HuffDeflate(&puff_reader, &bit_writer));
  334. TEST_AND_RETURN_FALSE(bit_writer.Size() == bytes_to_write);
  335. TEST_AND_RETURN_FALSE(puff_reader.BytesLeft() == 0);
  336. deflate_bit_pos_ = cur_deflate_->offset + cur_deflate_->length;
  337. if (extra_byte_ == 1) {
  338. deflate_buffer_->data()[bytes_to_write - 1] |=
  339. puff_buffer_->data()[cur_puff_->length] << (deflate_bit_pos_ & 7);
  340. deflate_bit_pos_ = (deflate_bit_pos_ + 7) & ~7ull;
  341. } else if ((deflate_bit_pos_ & 7) != 0) {
  342. // This happens if current and next deflate finish and end on the same
  343. // byte, then we cannot write into output until we have huffed the
  344. // next puff buffer, so untill then we cache it into |last_byte_| and
  345. // we won't write it out.
  346. last_byte_ = deflate_buffer_->data()[bytes_to_write - 1];
  347. bytes_to_write--;
  348. }
  349. // Write |deflate_buffer_| into output.
  350. TEST_AND_RETURN_FALSE(
  351. stream_->Write(deflate_buffer_->data(), bytes_to_write));
  352. // Move to the next deflate/puff.
  353. puff_pos_ += skip_bytes_;
  354. skip_bytes_ = 0;
  355. cur_puff_++;
  356. cur_deflate_++;
  357. if (cur_puff_ == puffs_.end()) {
  358. break;
  359. }
  360. // Find if need an extra byte to cache at the end.
  361. TEST_AND_RETURN_FALSE(SetExtraByte());
  362. }
  363. }
  364. }
  365. TEST_AND_RETURN_FALSE(bytes_wrote == length);
  366. return true;
  367. }
  368. bool PuffinStream::SetExtraByte() {
  369. TEST_AND_RETURN_FALSE(cur_deflate_ != deflates_.end());
  370. if ((cur_deflate_ + 1) == deflates_.end()) {
  371. extra_byte_ = 0;
  372. return true;
  373. }
  374. uint64_t end_bit = cur_deflate_->offset + cur_deflate_->length;
  375. if ((end_bit & 7) && ((end_bit + 7) & ~7ull) <= (cur_deflate_ + 1)->offset) {
  376. extra_byte_ = 1;
  377. } else {
  378. extra_byte_ = 0;
  379. }
  380. return true;
  381. }
  382. bool PuffinStream::GetPuffCache(int puff_id,
  383. uint64_t puff_size,
  384. shared_ptr<Buffer>* buffer) {
  385. bool found = false;
  386. // Search for it.
  387. std::pair<int, shared_ptr<Buffer>> cache;
  388. // TODO(*): Find a faster way of doing this? Maybe change the data structure
  389. // that supports faster search.
  390. for (auto iter = caches_.begin(); iter != caches_.end(); ++iter) {
  391. if (iter->first == puff_id) {
  392. cache = std::move(*iter);
  393. found = true;
  394. // Remove it so later we can add it to the begining of the list.
  395. caches_.erase(iter);
  396. break;
  397. }
  398. }
  399. // If not found, either create one or get one from the list.
  400. if (!found) {
  401. // If |caches_| were full, remove last ones in the list (least used), until
  402. // we have enough space for the new cache.
  403. while (!caches_.empty() && cur_cache_size_ + puff_size > max_cache_size_) {
  404. cache = std::move(caches_.back());
  405. caches_.pop_back(); // Remove it from the list.
  406. cur_cache_size_ -= cache.second->capacity();
  407. }
  408. // If we have not populated the cache yet, create one.
  409. if (!cache.second) {
  410. cache.second = std::make_shared<Buffer>(puff_size);
  411. }
  412. cache.second->resize(puff_size);
  413. constexpr uint64_t kMaxSizeDifference = 20 * 1024;
  414. if (puff_size + kMaxSizeDifference < cache.second->capacity()) {
  415. cache.second->shrink_to_fit();
  416. }
  417. cur_cache_size_ += cache.second->capacity();
  418. cache.first = puff_id;
  419. }
  420. *buffer = cache.second;
  421. // By now we have either removed a cache or created new one. Now we have to
  422. // insert it in the front of the list so it becomes the most recently used
  423. // one.
  424. caches_.push_front(std::move(cache));
  425. return found;
  426. }
  427. } // namespace puffin