authorgravatar for 4678790+dweiller@users.noreply.github.comDominic <4678790+dweiller@users.noreply.github.com> 2023-01-22 13:32:16+11:00
committergravatar for 4678790+dweiller@users.noreply.github.comDominic <4678790+dweiller@users.noreply.github.com> 2023-02-20 09:09:05+11:00
log18091723d5afa8001e0fd71274dc4b74d601d0e1
treee004ebd9da208b8bce207397b86962dea1f8eab4
parent61cb5143872ec2f3ae9c2942a03e18968bf27761

std.compress.zstandard: cleanup decodeBlock

1 files changed, 78 insertions(+), 70 deletions(-)

lib/std/compress/zstandard/decompress.zig+78-70
......@@ -65,12 +65,14 @@ const DecodeState = struct {
6565 match_fse_buffer: []Table.Fse,
6666 literal_fse_buffer: []Table.Fse,
6767
68 literal_written_count: usize,
68 fse_tables_undefined: bool,
6969
7070 literal_stream_reader: ReverseBitReader(ReversedByteReader.Reader),
7171 literal_stream_bytes: ReversedByteReader,
7272 literal_stream_index: usize,
73 huffman_tree: Literals.HuffmanTree,
73 huffman_tree: ?Literals.HuffmanTree,
74
75 literal_written_count: usize,
7476
7577 fn StateData(comptime max_accuracy_log: comptime_int) type {
7678 return struct {
......@@ -129,7 +131,6 @@ const DecodeState = struct {
129131 src: []const u8,
130132 comptime choice: DataType,
131133 mode: Sequences.Header.Mode,
132 first_compressed_block: bool,
133134 ) !usize {
134135 const field_name = @tagName(choice);
135136 switch (mode) {
......@@ -162,7 +163,7 @@ const DecodeState = struct {
162163 dumpFseTable(field_name, @field(self, field_name).table.fse);
163164 return counting_reader.bytes_read;
164165 },
165 .repeat => return if (first_compressed_block) error.RepeatModeFirst else 0,
166 .repeat => return if (self.fse_tables_undefined) error.RepeatModeFirst else 0,
166167 }
167168 }
168169
......@@ -275,7 +276,7 @@ const DecodeState = struct {
275276 },
276277 .compressed, .treeless => {
277278 // const written_bytes_per_stream = (literals.header.regenerated_size + 3) / 4;
278 const huffman_tree = self.huffman_tree;
279 const huffman_tree = self.huffman_tree orelse unreachable;
279280 const max_bit_count = huffman_tree.max_bit_count;
280281 const starting_bit_count = Literals.HuffmanTree.weightToBitCount(
281282 huffman_tree.nodes[huffman_tree.symbol_count_minus_one].weight,
......@@ -399,14 +400,14 @@ pub fn decodeFrameBlocks(dest: []u8, src: []const u8, consumed_count: *usize, ha
399400 .match_fse_buffer = &match_fse_data,
400401 .offset_fse_buffer = &offset_fse_data,
401402
403 .fse_tables_undefined = true,
404
402405 .literal_written_count = 0,
403406 .literal_stream_reader = undefined,
404407 .literal_stream_bytes = undefined,
405408 .literal_stream_index = undefined,
406 .huffman_tree = undefined,
409 .huffman_tree = null,
407410 };
408 var first_compressed_block = true;
409 var first_compressed_literals = true;
410411 var written_count: usize = 0;
411412 while (true) : ({
412413 block_header = decodeBlockHeader(src[bytes_read..][0..3]);
......@@ -417,8 +418,6 @@ pub fn decodeFrameBlocks(dest: []u8, src: []const u8, consumed_count: *usize, ha
417418 src[bytes_read..],
418419 block_header,
419420 &decode_state,
420 &first_compressed_block,
421 &first_compressed_literals,
422421 &bytes_read,
423422 written_count,
424423 );
......@@ -430,13 +429,77 @@ pub fn decodeFrameBlocks(dest: []u8, src: []const u8, consumed_count: *usize, ha
430429 return written_count;
431430}
432431
432fn decodeRawBlock(dest: []u8, src: []const u8, block_size: u21, consumed_count: *usize) usize {
433 log.debug("writing raw block - size {d}", .{block_size});
434 const data = src[0..block_size];
435 std.mem.copy(u8, dest, data);
436 consumed_count.* += block_size;
437 return block_size;
438}
439
440fn decodeRleBlock(dest: []u8, src: []const u8, block_size: u21, consumed_count: *usize) usize {
441 log.debug("writing rle block - '{x}'x{d}", .{ src[0], block_size });
442 var write_pos: usize = 0;
443 while (write_pos < block_size) : (write_pos += 1) {
444 dest[write_pos] = src[0];
445 }
446 consumed_count.* += 1;
447 return block_size;
448}
449
450fn prepareDecodeState(
451 decode_state: *DecodeState,
452 src: []const u8,
453 literals: Literals,
454 sequences_header: Sequences.Header,
455) !usize {
456 if (literals.huffman_tree) |tree| {
457 decode_state.huffman_tree = tree;
458 } else if (literals.header.block_type == .treeless and decode_state.huffman_tree == null) {
459 return error.TreelessLiteralsFirst;
460 }
461
462 switch (literals.header.block_type) {
463 .raw, .rle => {},
464 .compressed, .treeless => {
465 decode_state.literal_stream_index = 0;
466 switch (literals.streams) {
467 .one => |slice| try decode_state.initLiteralStream(slice),
468 .four => |streams| try decode_state.initLiteralStream(streams[0]),
469 }
470 },
471 }
472
473 if (sequences_header.sequence_count > 0) {
474 var bytes_read = try decode_state.updateFseTable(
475 src,
476 .literal,
477 sequences_header.literal_lengths,
478 );
479
480 bytes_read += try decode_state.updateFseTable(
481 src[bytes_read..],
482 .offset,
483 sequences_header.offsets,
484 );
485
486 bytes_read += try decode_state.updateFseTable(
487 src[bytes_read..],
488 .match,
489 sequences_header.match_lengths,
490 );
491 decode_state.fse_tables_undefined = false;
492
493 return bytes_read;
494 }
495 return 0;
496}
497
433498pub fn decodeBlock(
434499 dest: []u8,
435500 src: []const u8,
436501 block_header: frame.ZStandard.Block.Header,
437502 decode_state: *DecodeState,
438 first_compressed_block: *bool,
439 first_compressed_literals: *bool,
440503 consumed_count: *usize,
441504 written_count: usize,
442505) !usize {
......@@ -445,69 +508,14 @@ pub fn decodeBlock(
445508 if (block_maximum_size < block_size) return error.BlockSizeOverMaximum;
446509 // TODO: we probably want to enable safety for release-fast and release-small (or insert custom checks)
447510 switch (block_header.block_type) {
448 .raw => {
449 log.debug("writing raw block - size {d}", .{block_size});
450 const data = src[0..block_size];
451 std.mem.copy(u8, dest[written_count..], data);
452 consumed_count.* += block_size;
453 return block_size;
454 },
455 .rle => {
456 log.debug("writing rle block - '{x}'x{d}", .{ src[0], block_size });
457 var write_pos: usize = written_count;
458 while (write_pos < block_size + written_count) : (write_pos += 1) {
459 dest[write_pos] = src[0];
460 }
461 consumed_count.* += 1;
462 return block_size;
463 },
511 .raw => return decodeRawBlock(dest[written_count..], src, block_size, consumed_count),
512 .rle => return decodeRleBlock(dest[written_count..], src, block_size, consumed_count),
464513 .compressed => {
465514 var bytes_read: usize = 0;
466515 const literals = try decodeLiteralsSection(src, &bytes_read);
467516 const sequences_header = try decodeSequencesHeader(src[bytes_read..], &bytes_read);
468517
469 if (first_compressed_literals.* and literals.header.block_type == .treeless)
470 return error.TreelessLiteralsFirst;
471
472 if (literals.huffman_tree) |tree| {
473 decode_state.huffman_tree = tree;
474 first_compressed_literals.* = false;
475 }
476
477 switch (literals.header.block_type) {
478 .raw, .rle => {},
479 .compressed, .treeless => {
480 decode_state.literal_stream_index = 0;
481 switch (literals.streams) {
482 .one => |slice| try decode_state.initLiteralStream(slice),
483 .four => |streams| try decode_state.initLiteralStream(streams[0]),
484 }
485 },
486 }
487
488 if (sequences_header.sequence_count > 0) {
489 bytes_read += try decode_state.updateFseTable(
490 src[bytes_read..],
491 .literal,
492 sequences_header.literal_lengths,
493 first_compressed_block.*,
494 );
495
496 bytes_read += try decode_state.updateFseTable(
497 src[bytes_read..],
498 .offset,
499 sequences_header.offsets,
500 first_compressed_block.*,
501 );
502
503 bytes_read += try decode_state.updateFseTable(
504 src[bytes_read..],
505 .match,
506 sequences_header.match_lengths,
507 first_compressed_block.*,
508 );
509 first_compressed_block.* = false;
510 }
518 bytes_read += try prepareDecodeState(decode_state, src[bytes_read..], literals, sequences_header);
511519
512520 var bytes_written: usize = 0;
513521 if (sequences_header.sequence_count > 0) {