Deflate.cpp 40 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058
  1. /*
  2. * Copyright (c) 2020, the SerenityOS developers.
  3. * Copyright (c) 2021, Idan Horowitz <idan.horowitz@serenityos.org>
  4. *
  5. * SPDX-License-Identifier: BSD-2-Clause
  6. */
  7. #include <AK/Array.h>
  8. #include <AK/Assertions.h>
  9. #include <AK/BinaryHeap.h>
  10. #include <AK/BinarySearch.h>
  11. #include <AK/BitStream.h>
  12. #include <AK/MemoryStream.h>
  13. #include <string.h>
  14. #include <LibCompress/Deflate.h>
  15. namespace Compress {
  16. static constexpr u8 deflate_special_code_length_copy = 16;
  17. static constexpr u8 deflate_special_code_length_zeros = 17;
  18. static constexpr u8 deflate_special_code_length_long_zeros = 18;
  19. CanonicalCode const& CanonicalCode::fixed_literal_codes()
  20. {
  21. static CanonicalCode code;
  22. static bool initialized = false;
  23. if (initialized)
  24. return code;
  25. code = MUST(CanonicalCode::from_bytes(fixed_literal_bit_lengths));
  26. initialized = true;
  27. return code;
  28. }
  29. CanonicalCode const& CanonicalCode::fixed_distance_codes()
  30. {
  31. static CanonicalCode code;
  32. static bool initialized = false;
  33. if (initialized)
  34. return code;
  35. code = MUST(CanonicalCode::from_bytes(fixed_distance_bit_lengths));
  36. initialized = true;
  37. return code;
  38. }
  39. ErrorOr<CanonicalCode> CanonicalCode::from_bytes(ReadonlyBytes bytes)
  40. {
  41. CanonicalCode code;
  42. auto non_zero_symbols = 0;
  43. auto last_non_zero = -1;
  44. for (size_t i = 0; i < bytes.size(); i++) {
  45. if (bytes[i] != 0) {
  46. non_zero_symbols++;
  47. last_non_zero = i;
  48. }
  49. }
  50. if (non_zero_symbols == 1) { // special case - only 1 symbol
  51. code.m_prefix_table[0] = PrefixTableEntry { static_cast<u16>(last_non_zero), 1u };
  52. code.m_prefix_table[1] = code.m_prefix_table[0];
  53. code.m_max_prefixed_code_length = 1;
  54. if (code.m_bit_codes.size() < static_cast<size_t>(last_non_zero + 1)) {
  55. TRY(code.m_bit_codes.try_resize(last_non_zero + 1));
  56. TRY(code.m_bit_code_lengths.try_resize(last_non_zero + 1));
  57. }
  58. code.m_bit_codes[last_non_zero] = 0;
  59. code.m_bit_code_lengths[last_non_zero] = 1;
  60. return code;
  61. }
  62. struct PrefixCode {
  63. u16 symbol_code { 0 };
  64. u16 symbol_value { 0 };
  65. u16 code_length { 0 };
  66. };
  67. Array<PrefixCode, 1 << CanonicalCode::max_allowed_prefixed_code_length> prefix_codes;
  68. size_t number_of_prefix_codes = 0;
  69. auto next_code = 0;
  70. for (size_t code_length = 1; code_length <= 15; ++code_length) {
  71. next_code <<= 1;
  72. auto start_bit = 1 << code_length;
  73. for (size_t symbol = 0; symbol < bytes.size(); ++symbol) {
  74. if (bytes[symbol] != code_length)
  75. continue;
  76. if (next_code > start_bit)
  77. return Error::from_string_literal("Failed to decode code lengths");
  78. if (code_length <= CanonicalCode::max_allowed_prefixed_code_length) {
  79. if (number_of_prefix_codes >= prefix_codes.size())
  80. return Error::from_string_literal("Invalid canonical Huffman code");
  81. auto& prefix_code = prefix_codes[number_of_prefix_codes++];
  82. prefix_code.symbol_code = next_code;
  83. prefix_code.symbol_value = symbol;
  84. prefix_code.code_length = code_length;
  85. code.m_max_prefixed_code_length = code_length;
  86. } else {
  87. code.m_symbol_codes.append(start_bit | next_code);
  88. code.m_symbol_values.append(symbol);
  89. }
  90. if (code.m_bit_codes.size() < symbol + 1) {
  91. TRY(code.m_bit_codes.try_resize(symbol + 1));
  92. TRY(code.m_bit_code_lengths.try_resize(symbol + 1));
  93. }
  94. code.m_bit_codes[symbol] = fast_reverse16(start_bit | next_code, code_length); // DEFLATE writes huffman encoded symbols as lsb-first
  95. code.m_bit_code_lengths[symbol] = code_length;
  96. next_code++;
  97. }
  98. }
  99. if (next_code != (1 << 15))
  100. return Error::from_string_literal("Failed to decode code lengths");
  101. for (auto [symbol_code, symbol_value, code_length] : prefix_codes) {
  102. if (code_length == 0 || code_length > CanonicalCode::max_allowed_prefixed_code_length)
  103. break;
  104. auto shift = code.m_max_prefixed_code_length - code_length;
  105. symbol_code <<= shift;
  106. for (size_t j = 0; j < (1u << shift); ++j) {
  107. auto index = fast_reverse16(symbol_code + j, code.m_max_prefixed_code_length);
  108. code.m_prefix_table[index] = PrefixTableEntry { symbol_value, code_length };
  109. }
  110. }
  111. return code;
  112. }
  113. ErrorOr<u32> CanonicalCode::read_symbol(LittleEndianInputBitStream& stream) const
  114. {
  115. auto prefix = TRY(stream.peek_bits<size_t>(m_max_prefixed_code_length));
  116. if (auto [symbol_value, code_length] = m_prefix_table[prefix]; code_length != 0) {
  117. stream.discard_previously_peeked_bits(code_length);
  118. return symbol_value;
  119. }
  120. auto code_bits = TRY(stream.read_bits<u16>(m_max_prefixed_code_length));
  121. code_bits = fast_reverse16(code_bits, m_max_prefixed_code_length);
  122. code_bits |= 1 << m_max_prefixed_code_length;
  123. for (size_t i = m_max_prefixed_code_length; i < 16; ++i) {
  124. size_t index;
  125. if (binary_search(m_symbol_codes.span(), code_bits, &index))
  126. return m_symbol_values[index];
  127. code_bits = code_bits << 1 | TRY(stream.read_bit());
  128. }
  129. return Error::from_string_literal("Symbol exceeds maximum symbol number");
  130. }
  131. ErrorOr<void> CanonicalCode::write_symbol(LittleEndianOutputBitStream& stream, u32 symbol) const
  132. {
  133. auto code = symbol < m_bit_codes.size() ? m_bit_codes[symbol] : 0u;
  134. auto length = symbol < m_bit_code_lengths.size() ? m_bit_code_lengths[symbol] : 0u;
  135. TRY(stream.write_bits(code, length));
  136. return {};
  137. }
  138. DeflateDecompressor::CompressedBlock::CompressedBlock(DeflateDecompressor& decompressor, CanonicalCode literal_codes, Optional<CanonicalCode> distance_codes)
  139. : m_decompressor(decompressor)
  140. , m_literal_codes(literal_codes)
  141. , m_distance_codes(distance_codes)
  142. {
  143. }
  144. ErrorOr<bool> DeflateDecompressor::CompressedBlock::try_read_more()
  145. {
  146. if (m_eof == true)
  147. return false;
  148. auto const symbol = TRY(m_literal_codes.read_symbol(*m_decompressor.m_input_stream));
  149. if (symbol >= 286)
  150. return Error::from_string_literal("Invalid deflate literal/length symbol");
  151. if (symbol < 256) {
  152. u8 byte_symbol = symbol;
  153. m_decompressor.m_output_buffer.write({ &byte_symbol, sizeof(byte_symbol) });
  154. return true;
  155. }
  156. if (symbol == 256) {
  157. m_eof = true;
  158. return false;
  159. }
  160. if (!m_distance_codes.has_value())
  161. return Error::from_string_literal("Distance codes have not been initialized");
  162. auto const length = TRY(m_decompressor.decode_length(symbol));
  163. auto const distance_symbol = TRY(m_distance_codes.value().read_symbol(*m_decompressor.m_input_stream));
  164. if (distance_symbol >= 30)
  165. return Error::from_string_literal("Invalid deflate distance symbol");
  166. auto const distance = TRY(m_decompressor.decode_distance(distance_symbol));
  167. auto copied_length = TRY(m_decompressor.m_output_buffer.copy_from_seekback(distance, length));
  168. // TODO: What should we do if the output buffer is full?
  169. VERIFY(copied_length == length);
  170. return true;
  171. }
  172. DeflateDecompressor::UncompressedBlock::UncompressedBlock(DeflateDecompressor& decompressor, size_t length)
  173. : m_decompressor(decompressor)
  174. , m_bytes_remaining(length)
  175. {
  176. }
  177. ErrorOr<bool> DeflateDecompressor::UncompressedBlock::try_read_more()
  178. {
  179. if (m_bytes_remaining == 0)
  180. return false;
  181. if (m_decompressor.m_input_stream->is_eof())
  182. return Error::from_string_literal("Input data ends in the middle of an uncompressed DEFLATE block");
  183. Array<u8, 4096> temporary_buffer;
  184. auto readable_bytes = temporary_buffer.span().trim(min(m_bytes_remaining, m_decompressor.m_output_buffer.empty_space()));
  185. auto read_bytes = TRY(m_decompressor.m_input_stream->read_some(readable_bytes));
  186. auto written_bytes = m_decompressor.m_output_buffer.write(read_bytes);
  187. VERIFY(read_bytes.size() == written_bytes);
  188. m_bytes_remaining -= written_bytes;
  189. return true;
  190. }
  191. ErrorOr<NonnullOwnPtr<DeflateDecompressor>> DeflateDecompressor::construct(MaybeOwned<LittleEndianInputBitStream> stream)
  192. {
  193. auto output_buffer = TRY(CircularBuffer::create_empty(32 * KiB));
  194. return TRY(adopt_nonnull_own_or_enomem(new (nothrow) DeflateDecompressor(move(stream), move(output_buffer))));
  195. }
  196. DeflateDecompressor::DeflateDecompressor(MaybeOwned<LittleEndianInputBitStream> stream, CircularBuffer output_buffer)
  197. : m_input_stream(move(stream))
  198. , m_output_buffer(move(output_buffer))
  199. {
  200. }
  201. DeflateDecompressor::~DeflateDecompressor()
  202. {
  203. if (m_state == State::ReadingCompressedBlock)
  204. m_compressed_block.~CompressedBlock();
  205. if (m_state == State::ReadingUncompressedBlock)
  206. m_uncompressed_block.~UncompressedBlock();
  207. }
  208. ErrorOr<Bytes> DeflateDecompressor::read_some(Bytes bytes)
  209. {
  210. size_t total_read = 0;
  211. while (total_read < bytes.size()) {
  212. auto slice = bytes.slice(total_read);
  213. if (m_state == State::Idle) {
  214. if (m_read_final_block)
  215. break;
  216. m_read_final_block = TRY(m_input_stream->read_bit());
  217. auto const block_type = TRY(m_input_stream->read_bits(2));
  218. if (block_type == 0b00) {
  219. m_input_stream->align_to_byte_boundary();
  220. u16 length = TRY(m_input_stream->read_value<LittleEndian<u16>>());
  221. u16 negated_length = TRY(m_input_stream->read_value<LittleEndian<u16>>());
  222. if ((length ^ 0xffff) != negated_length)
  223. return Error::from_string_literal("Calculated negated length does not equal stored negated length");
  224. m_state = State::ReadingUncompressedBlock;
  225. new (&m_uncompressed_block) UncompressedBlock(*this, length);
  226. continue;
  227. }
  228. if (block_type == 0b01) {
  229. m_state = State::ReadingCompressedBlock;
  230. new (&m_compressed_block) CompressedBlock(*this, CanonicalCode::fixed_literal_codes(), CanonicalCode::fixed_distance_codes());
  231. continue;
  232. }
  233. if (block_type == 0b10) {
  234. CanonicalCode literal_codes;
  235. Optional<CanonicalCode> distance_codes;
  236. TRY(decode_codes(literal_codes, distance_codes));
  237. m_state = State::ReadingCompressedBlock;
  238. new (&m_compressed_block) CompressedBlock(*this, literal_codes, distance_codes);
  239. continue;
  240. }
  241. return Error::from_string_literal("Unhandled block type for Idle state");
  242. }
  243. if (m_state == State::ReadingCompressedBlock) {
  244. auto nread = m_output_buffer.read(slice).size();
  245. while (nread < slice.size() && TRY(m_compressed_block.try_read_more())) {
  246. nread += m_output_buffer.read(slice.slice(nread)).size();
  247. }
  248. total_read += nread;
  249. if (nread == slice.size())
  250. break;
  251. m_compressed_block.~CompressedBlock();
  252. m_state = State::Idle;
  253. continue;
  254. }
  255. if (m_state == State::ReadingUncompressedBlock) {
  256. auto nread = m_output_buffer.read(slice).size();
  257. while (nread < slice.size() && TRY(m_uncompressed_block.try_read_more())) {
  258. nread += m_output_buffer.read(slice.slice(nread)).size();
  259. }
  260. total_read += nread;
  261. if (nread == slice.size())
  262. break;
  263. m_uncompressed_block.~UncompressedBlock();
  264. m_state = State::Idle;
  265. continue;
  266. }
  267. VERIFY_NOT_REACHED();
  268. }
  269. return bytes.slice(0, total_read);
  270. }
  271. bool DeflateDecompressor::is_eof() const { return m_state == State::Idle && m_read_final_block; }
  272. ErrorOr<size_t> DeflateDecompressor::write_some(ReadonlyBytes)
  273. {
  274. return Error::from_errno(EBADF);
  275. }
  276. bool DeflateDecompressor::is_open() const
  277. {
  278. return true;
  279. }
  280. void DeflateDecompressor::close()
  281. {
  282. }
  283. ErrorOr<ByteBuffer> DeflateDecompressor::decompress_all(ReadonlyBytes bytes)
  284. {
  285. FixedMemoryStream memory_stream { bytes };
  286. LittleEndianInputBitStream bit_stream { MaybeOwned<Stream>(memory_stream) };
  287. auto deflate_stream = TRY(DeflateDecompressor::construct(MaybeOwned<LittleEndianInputBitStream>(bit_stream)));
  288. return deflate_stream->read_until_eof(4096);
  289. }
  290. ErrorOr<u32> DeflateDecompressor::decode_length(u32 symbol)
  291. {
  292. if (symbol <= 264)
  293. return symbol - 254;
  294. if (symbol <= 284) {
  295. auto extra_bits = (symbol - 261) / 4;
  296. return (((symbol - 265) % 4 + 4) << extra_bits) + 3 + TRY(m_input_stream->read_bits(extra_bits));
  297. }
  298. if (symbol == 285)
  299. return DeflateDecompressor::max_back_reference_length;
  300. VERIFY_NOT_REACHED();
  301. }
  302. ErrorOr<u32> DeflateDecompressor::decode_distance(u32 symbol)
  303. {
  304. if (symbol <= 3)
  305. return symbol + 1;
  306. if (symbol <= 29) {
  307. auto extra_bits = (symbol / 2) - 1;
  308. return ((symbol % 2 + 2) << extra_bits) + 1 + TRY(m_input_stream->read_bits(extra_bits));
  309. }
  310. VERIFY_NOT_REACHED();
  311. }
  312. ErrorOr<void> DeflateDecompressor::decode_codes(CanonicalCode& literal_code, Optional<CanonicalCode>& distance_code)
  313. {
  314. auto literal_code_count = TRY(m_input_stream->read_bits(5)) + 257;
  315. auto distance_code_count = TRY(m_input_stream->read_bits(5)) + 1;
  316. auto code_length_count = TRY(m_input_stream->read_bits(4)) + 4;
  317. // First we have to extract the code lengths of the code that was used to encode the code lengths of
  318. // the code that was used to encode the block.
  319. u8 code_lengths_code_lengths[19] = { 0 };
  320. for (size_t i = 0; i < code_length_count; ++i) {
  321. code_lengths_code_lengths[code_lengths_code_lengths_order[i]] = TRY(m_input_stream->read_bits(3));
  322. }
  323. // Now we can extract the code that was used to encode the code lengths of the code that was used to
  324. // encode the block.
  325. auto const code_length_code = TRY(CanonicalCode::from_bytes({ code_lengths_code_lengths, sizeof(code_lengths_code_lengths) }));
  326. // Next we extract the code lengths of the code that was used to encode the block.
  327. Vector<u8, 286> code_lengths;
  328. while (code_lengths.size() < literal_code_count + distance_code_count) {
  329. auto symbol = TRY(code_length_code.read_symbol(*m_input_stream));
  330. if (symbol < deflate_special_code_length_copy) {
  331. code_lengths.append(static_cast<u8>(symbol));
  332. } else if (symbol == deflate_special_code_length_copy) {
  333. if (code_lengths.is_empty())
  334. return Error::from_string_literal("Found no codes to copy before a copy block");
  335. auto nrepeat = 3 + TRY(m_input_stream->read_bits(2));
  336. for (size_t j = 0; j < nrepeat; ++j)
  337. code_lengths.append(code_lengths.last());
  338. } else if (symbol == deflate_special_code_length_zeros) {
  339. auto nrepeat = 3 + TRY(m_input_stream->read_bits(3));
  340. for (size_t j = 0; j < nrepeat; ++j)
  341. code_lengths.append(0);
  342. } else {
  343. VERIFY(symbol == deflate_special_code_length_long_zeros);
  344. auto nrepeat = 11 + TRY(m_input_stream->read_bits(7));
  345. for (size_t j = 0; j < nrepeat; ++j)
  346. code_lengths.append(0);
  347. }
  348. }
  349. if (code_lengths.size() != literal_code_count + distance_code_count)
  350. return Error::from_string_literal("Number of code lengths does not match the sum of codes");
  351. // Now we extract the code that was used to encode literals and lengths in the block.
  352. literal_code = TRY(CanonicalCode::from_bytes(code_lengths.span().trim(literal_code_count)));
  353. // Now we extract the code that was used to encode distances in the block.
  354. if (distance_code_count == 1) {
  355. auto length = code_lengths[literal_code_count];
  356. if (length == 0)
  357. return {};
  358. else if (length != 1)
  359. return Error::from_string_literal("Length for a single distance code is longer than 1");
  360. }
  361. distance_code = TRY(CanonicalCode::from_bytes(code_lengths.span().slice(literal_code_count)));
  362. return {};
  363. }
  364. ErrorOr<NonnullOwnPtr<DeflateCompressor>> DeflateCompressor::construct(MaybeOwned<Stream> stream, CompressionLevel compression_level)
  365. {
  366. auto bit_stream = TRY(try_make<LittleEndianOutputBitStream>(move(stream)));
  367. auto deflate_compressor = TRY(adopt_nonnull_own_or_enomem(new (nothrow) DeflateCompressor(move(bit_stream), compression_level)));
  368. return deflate_compressor;
  369. }
  370. DeflateCompressor::DeflateCompressor(NonnullOwnPtr<LittleEndianOutputBitStream> stream, CompressionLevel compression_level)
  371. : m_compression_level(compression_level)
  372. , m_compression_constants(compression_constants[static_cast<int>(m_compression_level)])
  373. , m_output_stream(move(stream))
  374. {
  375. m_symbol_frequencies.fill(0);
  376. m_distance_frequencies.fill(0);
  377. }
  378. DeflateCompressor::~DeflateCompressor()
  379. {
  380. VERIFY(m_finished);
  381. }
  382. ErrorOr<Bytes> DeflateCompressor::read_some(Bytes)
  383. {
  384. return Error::from_errno(EBADF);
  385. }
  386. ErrorOr<size_t> DeflateCompressor::write_some(ReadonlyBytes bytes)
  387. {
  388. VERIFY(!m_finished);
  389. size_t total_written = 0;
  390. while (!bytes.is_empty()) {
  391. auto n_written = bytes.copy_trimmed_to(pending_block().slice(m_pending_block_size));
  392. m_pending_block_size += n_written;
  393. if (m_pending_block_size == block_size)
  394. TRY(flush());
  395. bytes = bytes.slice(n_written);
  396. total_written += n_written;
  397. }
  398. return total_written;
  399. }
  400. bool DeflateCompressor::is_eof() const
  401. {
  402. return true;
  403. }
  404. bool DeflateCompressor::is_open() const
  405. {
  406. return m_output_stream->is_open();
  407. }
  408. void DeflateCompressor::close()
  409. {
  410. }
  411. // Knuth's multiplicative hash on 4 bytes
  412. u16 DeflateCompressor::hash_sequence(u8 const* bytes)
  413. {
  414. constexpr const u32 knuth_constant = 2654435761; // shares no common factors with 2^32
  415. return ((bytes[0] | bytes[1] << 8 | bytes[2] << 16 | bytes[3] << 24) * knuth_constant) >> (32 - hash_bits);
  416. }
  417. size_t DeflateCompressor::compare_match_candidate(size_t start, size_t candidate, size_t previous_match_length, size_t maximum_match_length)
  418. {
  419. VERIFY(previous_match_length < maximum_match_length);
  420. // We firstly check that the match is at least (prev_match_length + 1) long, we check backwards as there's a higher chance the end mismatches
  421. for (ssize_t i = previous_match_length; i >= 0; i--) {
  422. if (m_rolling_window[start + i] != m_rolling_window[candidate + i])
  423. return 0;
  424. }
  425. // Find the actual length
  426. auto match_length = previous_match_length + 1;
  427. while (match_length < maximum_match_length && m_rolling_window[start + match_length] == m_rolling_window[candidate + match_length]) {
  428. match_length++;
  429. }
  430. VERIFY(match_length > previous_match_length);
  431. VERIFY(match_length <= maximum_match_length);
  432. return match_length;
  433. }
  434. size_t DeflateCompressor::find_back_match(size_t start, u16 hash, size_t previous_match_length, size_t maximum_match_length, size_t& match_position)
  435. {
  436. auto max_chain_length = m_compression_constants.max_chain;
  437. if (previous_match_length == 0)
  438. previous_match_length = min_match_length - 1; // we only care about matches that are at least min_match_length long
  439. if (previous_match_length >= maximum_match_length)
  440. return 0; // we can't improve a maximum length match
  441. if (previous_match_length >= m_compression_constants.max_lazy_length)
  442. return 0; // the previous match is already pretty, we shouldn't waste another full search
  443. if (previous_match_length >= m_compression_constants.good_match_length)
  444. max_chain_length /= 4; // we already have a pretty good much, so do a shorter search
  445. auto candidate = m_hash_head[hash];
  446. auto match_found = false;
  447. while (max_chain_length--) {
  448. if (candidate == empty_slot)
  449. break; // no remaining candidates
  450. VERIFY(candidate < start);
  451. if (start - candidate > window_size)
  452. break; // outside the window
  453. auto match_length = compare_match_candidate(start, candidate, previous_match_length, maximum_match_length);
  454. if (match_length != 0) {
  455. match_found = true;
  456. match_position = candidate;
  457. previous_match_length = match_length;
  458. if (match_length == maximum_match_length)
  459. return match_length; // bail if we got the maximum possible length
  460. }
  461. candidate = m_hash_prev[candidate % window_size];
  462. }
  463. if (!match_found)
  464. return 0; // we didn't find any matches
  465. return previous_match_length; // we found matches, but they were at most previous_match_length long
  466. }
  467. ALWAYS_INLINE u8 DeflateCompressor::distance_to_base(u16 distance)
  468. {
  469. return (distance <= 256) ? distance_to_base_lo[distance - 1] : distance_to_base_hi[(distance - 1) >> 7];
  470. }
  471. template<size_t Size>
  472. void DeflateCompressor::generate_huffman_lengths(Array<u8, Size>& lengths, Array<u16, Size> const& frequencies, size_t max_bit_length, u16 frequency_cap)
  473. {
  474. VERIFY((1u << max_bit_length) >= Size);
  475. u16 heap_keys[Size]; // Used for O(n) heap construction
  476. u16 heap_values[Size];
  477. u16 huffman_links[Size * 2] = { 0 };
  478. size_t non_zero_freqs = 0;
  479. for (size_t i = 0; i < Size; i++) {
  480. auto frequency = frequencies[i];
  481. if (frequency == 0)
  482. continue;
  483. if (frequency > frequency_cap) {
  484. frequency = frequency_cap;
  485. }
  486. heap_keys[non_zero_freqs] = frequency; // sort symbols by frequency
  487. heap_values[non_zero_freqs] = Size + non_zero_freqs; // huffman_links "links"
  488. non_zero_freqs++;
  489. }
  490. // special case for only 1 used symbol
  491. if (non_zero_freqs < 2) {
  492. for (size_t i = 0; i < Size; i++)
  493. lengths[i] = (frequencies[i] == 0) ? 0 : 1;
  494. return;
  495. }
  496. BinaryHeap<u16, u16, Size> heap { heap_keys, heap_values, non_zero_freqs };
  497. // build the huffman tree - binary heap is used for efficient frequency comparisons
  498. while (heap.size() > 1) {
  499. u16 lowest_frequency = heap.peek_min_key();
  500. u16 lowest_link = heap.pop_min();
  501. u16 second_lowest_frequency = heap.peek_min_key();
  502. u16 second_lowest_link = heap.pop_min();
  503. u16 new_link = heap.size() + 1;
  504. heap.insert(lowest_frequency + second_lowest_frequency, new_link);
  505. huffman_links[lowest_link] = new_link;
  506. huffman_links[second_lowest_link] = new_link;
  507. }
  508. non_zero_freqs = 0;
  509. for (size_t i = 0; i < Size; i++) {
  510. if (frequencies[i] == 0) {
  511. lengths[i] = 0;
  512. continue;
  513. }
  514. u16 link = huffman_links[Size + non_zero_freqs];
  515. non_zero_freqs++;
  516. size_t bit_length = 1;
  517. while (link != 1) {
  518. bit_length++;
  519. link = huffman_links[link];
  520. }
  521. if (bit_length > max_bit_length) {
  522. VERIFY(frequency_cap != 1);
  523. return generate_huffman_lengths(lengths, frequencies, max_bit_length, frequency_cap / 2);
  524. }
  525. lengths[i] = bit_length;
  526. }
  527. }
  528. void DeflateCompressor::lz77_compress_block()
  529. {
  530. for (auto& slot : m_hash_head) { // initialize chained hash table
  531. slot = empty_slot;
  532. }
  533. auto insert_hash = [&](auto pos, auto hash) {
  534. auto window_pos = pos % window_size;
  535. m_hash_prev[window_pos] = m_hash_head[hash];
  536. m_hash_head[hash] = window_pos;
  537. };
  538. auto emit_literal = [&](auto literal) {
  539. VERIFY(m_pending_symbol_size <= block_size + 1);
  540. auto index = m_pending_symbol_size++;
  541. m_symbol_buffer[index].distance = 0;
  542. m_symbol_buffer[index].literal = literal;
  543. m_symbol_frequencies[literal]++;
  544. };
  545. auto emit_back_reference = [&](auto distance, auto length) {
  546. VERIFY(m_pending_symbol_size <= block_size + 1);
  547. auto index = m_pending_symbol_size++;
  548. m_symbol_buffer[index].distance = distance;
  549. m_symbol_buffer[index].length = length;
  550. m_symbol_frequencies[length_to_symbol[length]]++;
  551. m_distance_frequencies[distance_to_base(distance)]++;
  552. };
  553. size_t previous_match_length = 0;
  554. size_t previous_match_position = 0;
  555. VERIFY(m_compression_constants.great_match_length <= max_match_length);
  556. // our block starts at block_size and is m_pending_block_size in length
  557. auto block_end = block_size + m_pending_block_size;
  558. size_t current_position;
  559. for (current_position = block_size; current_position < block_end - min_match_length + 1; current_position++) {
  560. auto hash = hash_sequence(&m_rolling_window[current_position]);
  561. size_t match_position;
  562. auto match_length = find_back_match(current_position, hash, previous_match_length,
  563. min(m_compression_constants.great_match_length, block_end - current_position), match_position);
  564. insert_hash(current_position, hash);
  565. // if the previous match is as good as the new match, just use it
  566. if (previous_match_length != 0 && previous_match_length >= match_length) {
  567. emit_back_reference((current_position - 1) - previous_match_position, previous_match_length);
  568. // skip all the bytes that are included in this match
  569. for (size_t j = current_position + 1; j < min(current_position - 1 + previous_match_length, block_end - min_match_length + 1); j++) {
  570. insert_hash(j, hash_sequence(&m_rolling_window[j]));
  571. }
  572. current_position = (current_position - 1) + previous_match_length - 1;
  573. previous_match_length = 0;
  574. continue;
  575. }
  576. if (match_length == 0) {
  577. VERIFY(previous_match_length == 0);
  578. emit_literal(m_rolling_window[current_position]);
  579. continue;
  580. }
  581. // if this is a lazy match, and the new match is better than the old one, output previous as literal
  582. if (previous_match_length != 0) {
  583. emit_literal(m_rolling_window[current_position - 1]);
  584. }
  585. previous_match_length = match_length;
  586. previous_match_position = match_position;
  587. }
  588. // clean up leftover lazy match
  589. if (previous_match_length != 0) {
  590. emit_back_reference((current_position - 1) - previous_match_position, previous_match_length);
  591. current_position = (current_position - 1) + previous_match_length;
  592. }
  593. // output remaining literals
  594. while (current_position < block_end) {
  595. emit_literal(m_rolling_window[current_position++]);
  596. }
  597. }
  598. size_t DeflateCompressor::huffman_block_length(Array<u8, max_huffman_literals> const& literal_bit_lengths, Array<u8, max_huffman_distances> const& distance_bit_lengths)
  599. {
  600. size_t length = 0;
  601. for (size_t i = 0; i < 286; i++) {
  602. auto frequency = m_symbol_frequencies[i];
  603. length += literal_bit_lengths[i] * frequency;
  604. if (i >= 257) // back reference length symbols
  605. length += packed_length_symbols[i - 257].extra_bits * frequency;
  606. }
  607. for (size_t i = 0; i < 30; i++) {
  608. auto frequency = m_distance_frequencies[i];
  609. length += distance_bit_lengths[i] * frequency;
  610. length += packed_distances[i].extra_bits * frequency;
  611. }
  612. return length;
  613. }
  614. size_t DeflateCompressor::uncompressed_block_length()
  615. {
  616. auto padding = 8 - ((m_output_stream->bit_offset() + 3) % 8);
  617. // 3 bit block header + align to byte + 2 * 16 bit length fields + block contents
  618. return 3 + padding + (2 * 16) + m_pending_block_size * 8;
  619. }
  620. size_t DeflateCompressor::fixed_block_length()
  621. {
  622. // block header + fixed huffman encoded block contents
  623. return 3 + huffman_block_length(fixed_literal_bit_lengths, fixed_distance_bit_lengths);
  624. }
  625. size_t DeflateCompressor::dynamic_block_length(Array<u8, max_huffman_literals> const& literal_bit_lengths, Array<u8, max_huffman_distances> const& distance_bit_lengths, Array<u8, 19> const& code_lengths_bit_lengths, Array<u16, 19> const& code_lengths_frequencies, size_t code_lengths_count)
  626. {
  627. // block header + literal code count + distance code count + code length count
  628. auto length = 3 + 5 + 5 + 4;
  629. // 3 bits per code_length
  630. length += 3 * code_lengths_count;
  631. for (size_t i = 0; i < code_lengths_frequencies.size(); i++) {
  632. auto frequency = code_lengths_frequencies[i];
  633. length += code_lengths_bit_lengths[i] * frequency;
  634. if (i == deflate_special_code_length_copy) {
  635. length += 2 * frequency;
  636. } else if (i == deflate_special_code_length_zeros) {
  637. length += 3 * frequency;
  638. } else if (i == deflate_special_code_length_long_zeros) {
  639. length += 7 * frequency;
  640. }
  641. }
  642. return length + huffman_block_length(literal_bit_lengths, distance_bit_lengths);
  643. }
  644. ErrorOr<void> DeflateCompressor::write_huffman(CanonicalCode const& literal_code, Optional<CanonicalCode> const& distance_code)
  645. {
  646. auto has_distances = distance_code.has_value();
  647. for (size_t i = 0; i < m_pending_symbol_size; i++) {
  648. if (m_symbol_buffer[i].distance == 0) {
  649. TRY(literal_code.write_symbol(*m_output_stream, m_symbol_buffer[i].literal));
  650. continue;
  651. }
  652. VERIFY(has_distances);
  653. auto symbol = length_to_symbol[m_symbol_buffer[i].length];
  654. TRY(literal_code.write_symbol(*m_output_stream, symbol));
  655. // Emit extra bits if needed
  656. TRY(m_output_stream->write_bits<u16>(m_symbol_buffer[i].length - packed_length_symbols[symbol - 257].base_length, packed_length_symbols[symbol - 257].extra_bits));
  657. auto base_distance = distance_to_base(m_symbol_buffer[i].distance);
  658. TRY(distance_code.value().write_symbol(*m_output_stream, base_distance));
  659. // Emit extra bits if needed
  660. TRY(m_output_stream->write_bits<u16>(m_symbol_buffer[i].distance - packed_distances[base_distance].base_distance, packed_distances[base_distance].extra_bits));
  661. }
  662. return {};
  663. }
  664. size_t DeflateCompressor::encode_huffman_lengths(Array<u8, max_huffman_literals + max_huffman_distances> const& lengths, size_t lengths_count, Array<code_length_symbol, max_huffman_literals + max_huffman_distances>& encoded_lengths)
  665. {
  666. size_t encoded_count = 0;
  667. size_t i = 0;
  668. while (i < lengths_count) {
  669. if (lengths[i] == 0) {
  670. auto zero_count = 0;
  671. for (size_t j = i; j < min(lengths_count, i + 138) && lengths[j] == 0; j++)
  672. zero_count++;
  673. if (zero_count < 3) { // below minimum repeated zero count
  674. encoded_lengths[encoded_count++].symbol = 0;
  675. i++;
  676. continue;
  677. }
  678. if (zero_count <= 10) {
  679. encoded_lengths[encoded_count].symbol = deflate_special_code_length_zeros;
  680. encoded_lengths[encoded_count++].count = zero_count;
  681. } else {
  682. encoded_lengths[encoded_count].symbol = deflate_special_code_length_long_zeros;
  683. encoded_lengths[encoded_count++].count = zero_count;
  684. }
  685. i += zero_count;
  686. continue;
  687. }
  688. encoded_lengths[encoded_count++].symbol = lengths[i++];
  689. auto copy_count = 0;
  690. for (size_t j = i; j < min(lengths_count, i + 6) && lengths[j] == lengths[i - 1]; j++)
  691. copy_count++;
  692. if (copy_count >= 3) {
  693. encoded_lengths[encoded_count].symbol = deflate_special_code_length_copy;
  694. encoded_lengths[encoded_count++].count = copy_count;
  695. i += copy_count;
  696. continue;
  697. }
  698. }
  699. return encoded_count;
  700. }
  701. size_t DeflateCompressor::encode_block_lengths(Array<u8, max_huffman_literals> const& literal_bit_lengths, Array<u8, max_huffman_distances> const& distance_bit_lengths, Array<code_length_symbol, max_huffman_literals + max_huffman_distances>& encoded_lengths, size_t& literal_code_count, size_t& distance_code_count)
  702. {
  703. literal_code_count = max_huffman_literals;
  704. distance_code_count = max_huffman_distances;
  705. VERIFY(literal_bit_lengths[256] != 0); // Make sure at least the EndOfBlock marker is present
  706. while (literal_bit_lengths[literal_code_count - 1] == 0)
  707. literal_code_count--;
  708. // Drop trailing zero lengths, keeping at least one
  709. while (distance_bit_lengths[distance_code_count - 1] == 0 && distance_code_count > 1)
  710. distance_code_count--;
  711. Array<u8, max_huffman_literals + max_huffman_distances> all_lengths {};
  712. size_t lengths_count = 0;
  713. for (size_t i = 0; i < literal_code_count; i++) {
  714. all_lengths[lengths_count++] = literal_bit_lengths[i];
  715. }
  716. for (size_t i = 0; i < distance_code_count; i++) {
  717. all_lengths[lengths_count++] = distance_bit_lengths[i];
  718. }
  719. return encode_huffman_lengths(all_lengths, lengths_count, encoded_lengths);
  720. }
  721. ErrorOr<void> DeflateCompressor::write_dynamic_huffman(CanonicalCode const& literal_code, size_t literal_code_count, Optional<CanonicalCode> const& distance_code, size_t distance_code_count, Array<u8, 19> const& code_lengths_bit_lengths, size_t code_length_count, Array<code_length_symbol, max_huffman_literals + max_huffman_distances> const& encoded_lengths, size_t encoded_lengths_count)
  722. {
  723. TRY(m_output_stream->write_bits(literal_code_count - 257, 5));
  724. TRY(m_output_stream->write_bits(distance_code_count - 1, 5));
  725. TRY(m_output_stream->write_bits(code_length_count - 4, 4));
  726. for (size_t i = 0; i < code_length_count; i++) {
  727. TRY(m_output_stream->write_bits(code_lengths_bit_lengths[code_lengths_code_lengths_order[i]], 3));
  728. }
  729. auto code_lengths_code = MUST(CanonicalCode::from_bytes(code_lengths_bit_lengths));
  730. for (size_t i = 0; i < encoded_lengths_count; i++) {
  731. auto encoded_length = encoded_lengths[i];
  732. TRY(code_lengths_code.write_symbol(*m_output_stream, encoded_length.symbol));
  733. if (encoded_length.symbol == deflate_special_code_length_copy) {
  734. TRY(m_output_stream->write_bits<u8>(encoded_length.count - 3, 2));
  735. } else if (encoded_length.symbol == deflate_special_code_length_zeros) {
  736. TRY(m_output_stream->write_bits<u8>(encoded_length.count - 3, 3));
  737. } else if (encoded_length.symbol == deflate_special_code_length_long_zeros) {
  738. TRY(m_output_stream->write_bits<u8>(encoded_length.count - 11, 7));
  739. }
  740. }
  741. TRY(write_huffman(literal_code, distance_code));
  742. return {};
  743. }
  744. ErrorOr<void> DeflateCompressor::flush()
  745. {
  746. TRY(m_output_stream->write_bits(m_finished, 1));
  747. // if this is just an empty block to signify the end of the deflate stream use the smallest block possible (10 bits total)
  748. if (m_pending_block_size == 0) {
  749. VERIFY(m_finished); // we shouldn't be writing empty blocks unless this is the final one
  750. TRY(m_output_stream->write_bits(0b01u, 2)); // fixed huffman codes
  751. TRY(m_output_stream->write_bits(0b0000000u, 7)); // end of block symbol
  752. TRY(m_output_stream->align_to_byte_boundary());
  753. return {};
  754. }
  755. auto write_uncompressed = [&]() -> ErrorOr<void> {
  756. TRY(m_output_stream->write_bits(0b00u, 2)); // no compression
  757. TRY(m_output_stream->align_to_byte_boundary());
  758. TRY(m_output_stream->write_value<LittleEndian<u16>>(m_pending_block_size));
  759. TRY(m_output_stream->write_value<LittleEndian<u16>>(~m_pending_block_size));
  760. TRY(m_output_stream->write_until_depleted(pending_block().slice(0, m_pending_block_size)));
  761. return {};
  762. };
  763. if (m_compression_level == CompressionLevel::STORE) { // disabled compression fast path
  764. TRY(write_uncompressed());
  765. m_pending_block_size = 0;
  766. return {};
  767. }
  768. // The following implementation of lz77 compression and huffman encoding is based on the reference implementation by Hans Wennborg https://www.hanshq.net/zip.html
  769. // this reads from the pending block and writes to m_symbol_buffer
  770. lz77_compress_block();
  771. // insert EndOfBlock marker to the symbol buffer
  772. m_symbol_buffer[m_pending_symbol_size].distance = 0;
  773. m_symbol_buffer[m_pending_symbol_size++].literal = 256;
  774. m_symbol_frequencies[256]++;
  775. // generate optimal dynamic huffman code lengths
  776. Array<u8, max_huffman_literals> dynamic_literal_bit_lengths {};
  777. Array<u8, max_huffman_distances> dynamic_distance_bit_lengths {};
  778. generate_huffman_lengths(dynamic_literal_bit_lengths, m_symbol_frequencies, 15); // deflate data huffman can use up to 15 bits per symbol
  779. generate_huffman_lengths(dynamic_distance_bit_lengths, m_distance_frequencies, 15);
  780. // encode literal and distance lengths together in deflate format
  781. Array<code_length_symbol, max_huffman_literals + max_huffman_distances> encoded_lengths {};
  782. size_t literal_code_count;
  783. size_t distance_code_count;
  784. auto encoded_lengths_count = encode_block_lengths(dynamic_literal_bit_lengths, dynamic_distance_bit_lengths, encoded_lengths, literal_code_count, distance_code_count);
  785. // count code length frequencies
  786. Array<u16, 19> code_lengths_frequencies { 0 };
  787. for (size_t i = 0; i < encoded_lengths_count; i++) {
  788. code_lengths_frequencies[encoded_lengths[i].symbol]++;
  789. }
  790. // generate optimal huffman code lengths code lengths
  791. Array<u8, 19> code_lengths_bit_lengths {};
  792. generate_huffman_lengths(code_lengths_bit_lengths, code_lengths_frequencies, 7); // deflate code length huffman can use up to 7 bits per symbol
  793. // calculate actual code length code lengths count (without trailing zeros)
  794. auto code_lengths_count = code_lengths_bit_lengths.size();
  795. while (code_lengths_bit_lengths[code_lengths_code_lengths_order[code_lengths_count - 1]] == 0)
  796. code_lengths_count--;
  797. auto uncompressed_size = uncompressed_block_length();
  798. auto fixed_huffman_size = fixed_block_length();
  799. auto dynamic_huffman_size = dynamic_block_length(dynamic_literal_bit_lengths, dynamic_distance_bit_lengths, code_lengths_bit_lengths, code_lengths_frequencies, code_lengths_count);
  800. // If the compression somehow didn't reduce the size enough, just write out the block uncompressed as it allows for much faster decompression
  801. if (uncompressed_size <= min(fixed_huffman_size, dynamic_huffman_size)) {
  802. TRY(write_uncompressed());
  803. } else if (fixed_huffman_size <= dynamic_huffman_size) {
  804. // If the fixed and dynamic huffman codes come out the same size, prefer the fixed version, as it takes less time to decode fixed huffman codes.
  805. TRY(m_output_stream->write_bits(0b01u, 2));
  806. TRY(write_huffman(CanonicalCode::fixed_literal_codes(), CanonicalCode::fixed_distance_codes()));
  807. } else {
  808. // dynamic huffman codes
  809. TRY(m_output_stream->write_bits(0b10u, 2));
  810. auto literal_code = MUST(CanonicalCode::from_bytes(dynamic_literal_bit_lengths));
  811. auto distance_code_or_error = CanonicalCode::from_bytes(dynamic_distance_bit_lengths);
  812. Optional<CanonicalCode> distance_code;
  813. if (!distance_code_or_error.is_error())
  814. distance_code = distance_code_or_error.release_value();
  815. TRY(write_dynamic_huffman(literal_code, literal_code_count, distance_code, distance_code_count, code_lengths_bit_lengths, code_lengths_count, encoded_lengths, encoded_lengths_count));
  816. }
  817. if (m_finished)
  818. TRY(m_output_stream->align_to_byte_boundary());
  819. // reset all block specific members
  820. m_pending_block_size = 0;
  821. m_pending_symbol_size = 0;
  822. m_symbol_frequencies.fill(0);
  823. m_distance_frequencies.fill(0);
  824. // On the final block this copy will potentially produce an invalid search window, but since its the final block we dont care
  825. pending_block().copy_trimmed_to({ m_rolling_window, block_size });
  826. return {};
  827. }
  828. ErrorOr<void> DeflateCompressor::final_flush()
  829. {
  830. VERIFY(!m_finished);
  831. m_finished = true;
  832. TRY(flush());
  833. TRY(m_output_stream->flush_buffer_to_stream());
  834. return {};
  835. }
  836. ErrorOr<ByteBuffer> DeflateCompressor::compress_all(ReadonlyBytes bytes, CompressionLevel compression_level)
  837. {
  838. auto output_stream = TRY(try_make<AllocatingMemoryStream>());
  839. auto deflate_stream = TRY(DeflateCompressor::construct(MaybeOwned<Stream>(*output_stream), compression_level));
  840. TRY(deflate_stream->write_until_depleted(bytes));
  841. TRY(deflate_stream->final_flush());
  842. auto buffer = TRY(ByteBuffer::create_uninitialized(output_stream->used_buffer_size()));
  843. TRY(output_stream->read_until_filled(buffer));
  844. return buffer;
  845. }
  846. }