diff --git a/zstd/src/bit_io/bit_writer.rs b/zstd/src/bit_io/bit_writer.rs index 9a198e82..1004093a 100644 --- a/zstd/src/bit_io/bit_writer.rs +++ b/zstd/src/bit_io/bit_writer.rs @@ -372,29 +372,25 @@ impl>> BitBurst<'_, V> { /// /// # Safety /// - /// The accumulator must have room: `bits_in_partial + num_bits <= 64`, and - /// `bits` must be clean above `num_bits` (`num_bits == 64 || bits >> - /// num_bits == 0`). Neither violation touches memory, but either corrupts - /// the stream with nothing to detect it, so the caller proves both from - /// its bit budget. Debug builds assert them. + /// The accumulator must keep a bit free: `bits_in_partial + num_bits < 64` + /// (upstream's `assert(nbBits + bitPos < 64)`), and `bits` must be clean + /// above `num_bits` (`bits >> num_bits == 0`). Neither violation touches + /// memory, but either corrupts the stream with nothing to detect it, so the + /// caller proves both from its bit budget. Debug builds assert them. #[inline(always)] pub(crate) unsafe fn add(&mut self, bits: u64, num_bits: usize) { debug_assert!( - self.bits_in_partial + num_bits <= 64, + self.bits_in_partial + num_bits < 64, "{num_bits} bits do not fit beside the {} held", self.bits_in_partial, ); debug_assert!( - num_bits == 64 || bits >> num_bits == 0, + bits >> num_bits == 0, "value has dirty high bits beyond num_bits={num_bits}", ); - // Masked so a full accumulator cannot make this a shift by 64, which is - // undefined for `u64`. x86 and AArch64 shifts mask the count - // themselves, so the mask costs nothing, where a `num_bits == 0` early - // return was a branch on every add. The accumulator is full only when - // `num_bits == 0`, and the second precondition then makes `bits == 0`, - // so the OR adds nothing whatever the count. - self.partial |= bits << (self.bits_in_partial & 63); + // `bits_in_partial < 64` by the first precondition, so the shift is in + // range. + self.partial |= bits << self.bits_in_partial; self.bits_in_partial += num_bits; } @@ -423,13 +419,9 @@ impl>> BitBurst<'_, V> { .write_unaligned(self.partial.to_le()); self.cursor = self.cursor.add(nb_bytes); } - // A full accumulator (64 bits) leaves nothing: shifting a `u64` by 64 - // is undefined, so it is zeroed instead. - self.partial = if nb_bytes == 8 { - 0 - } else { - self.partial >> (nb_bytes * 8) - }; + // `add` keeps the accumulator under 64 bits, so at most 7 bytes leave + // and the shift is in range (upstream `BIT_flushBitsFast`, unguarded). + self.partial >>= nb_bytes * 8; self.bits_in_partial &= 7; } diff --git a/zstd/src/bit_io/bit_writer/tests.rs b/zstd/src/bit_io/bit_writer/tests.rs index 1472cc73..c9698326 100644 --- a/zstd/src/bit_io/bit_writer/tests.rs +++ b/zstd/src/bit_io/bit_writer/tests.rs @@ -160,28 +160,41 @@ fn add_multiple_aligned() { assert_eq!(vec![0xFF, 0xF0, 0x0F, 0x00], bw.dump()); } -/// A burst flushing a full accumulator writes all eight bytes and holds -/// nothing after, so the next bits land on a clean byte rather than on bits -/// of the flushed value. +/// A burst flushing the fullest accumulator it allows (63 bits) writes the +/// seven whole bytes, keeps the last 7 bits, and lands the next bits right +/// after them rather than on bits of the flushed value. #[test] -fn a_burst_flushes_a_full_accumulator_clean() { +fn a_burst_flushes_63_bits_and_keeps_the_partial_byte() { let mut bw = BitWriter::new(); let mut burst = bw.burst(16); - // SAFETY: 64 bits into an empty accumulator, then 8 into an empty one, + // SAFETY: 63 bits into an empty accumulator, then 1 beside the 7 kept, // within the 16 bytes reserved. unsafe { - burst.add(0x0123_4567_89AB_CDEF, 64); + burst.add(0x7123_4567_89AB_CDEF, 63); burst.flush(); - burst.add(0xA5, 8); + burst.add(0x1, 1); burst.flush(); } burst.finish(); assert_eq!( bw.dump(), - vec![0xEF, 0xCD, 0xAB, 0x89, 0x67, 0x45, 0x23, 0x01, 0xA5] + vec![0xEF, 0xCD, 0xAB, 0x89, 0x67, 0x45, 0x23, 0xF1] ); } +/// An add that would fill all 64 bits breaks the burst's contract (the flush +/// shifts by the whole bytes held, which must stay under 8); debug builds +/// catch it at the add. +#[test] +#[cfg(debug_assertions)] +#[should_panic(expected = "do not fit")] +fn a_burst_refuses_a_full_accumulator() { + let mut bw = BitWriter::new(); + let mut burst = bw.burst(16); + // SAFETY: the debug assertion fires before anything is written. + unsafe { burst.add(0x0123_4567_89AB_CDEF, 64) }; +} + // #[test] // fn catches_more_than_in_buf() { // todo!(); diff --git a/zstd/src/encoding/blocks/compressed.rs b/zstd/src/encoding/blocks/compressed.rs index 3824222f..df37655b 100644 --- a/zstd/src/encoding/blocks/compressed.rs +++ b/zstd/src/encoding/blocks/compressed.rs @@ -2000,10 +2000,10 @@ fn fill_and_count( ); // SAFETY: the codes are at most 35 (LL), 52 (ML) and 31 (OF), under the // histogram's `SEQUENCE_CODE_SLOTS`. The LL and ML bounds hold in release - // because `encode_literal_length` and `encode_match_len` index - // `LL_EXTRA_BITS` (36 entries) and `ML_EXTRA_BITS` (53) with a bounds - // check, so a larger code panics there before reaching this write; the - // OF code is `ilog2` of a `u32`. Those lookups must stay checked. + // because `encode_literal_length` and `encode_match_len` assert the + // length below 2^17 before deriving a code from its high bit, so a larger + // code panics there before reaching this write; the OF code is `ilog2` + // of a `u32`. Those asserts must stay. unsafe { *ll_counts.get_unchecked_mut(usize::from(ll_code)) += 1; *ml_counts.get_unchecked_mut(usize::from(ml_code)) += 1; @@ -2889,6 +2889,18 @@ macro_rules! encode_sequences_body { // body is what leaves the accumulator with under a byte for the next // round. if sequences.len() > 1 { + // Streams 0 / 1 / 2 are LL / ML / OF; built only for a block that + // runs the loop. A local on purpose, not scratch kept on the + // compressor: addressed off the stack pointer it needs no base + // register, and the loop keeps two of the three states in + // registers. Held behind a pointer the same loop spilled one state + // more, +0.4..0.6% instructions at level 1 against the per-block + // clearing this costs. + let mut view = crate::fse::fse_encoder::SequenceStateTables::new(); + view.set(0, ll_table); + view.set(1, ml_table); + view.set(2, of_table); + let view = &view; // At most 63 bits a sequence leave as whole bytes, under 8 bytes; // 12 a sequence and 64 over covers that and the up to 7 bytes the // burst drains from the init-state writes above as it starts, @@ -2932,8 +2944,8 @@ macro_rules! encode_sequences_body { // derivation assumed. #[cfg(debug_assertions)] let mut state_diff_bits = 0usize; - if let Some(table) = of_table { - let next = table.next_state(of_code, of_state); + { + let next = view.next_state::<2>(of_code, of_state); let diff = crate::fse::fse_encoder::transition_bits(of_state, next.num_bits); unsafe { burst.add(diff as u64, next.num_bits as usize); @@ -2944,8 +2956,8 @@ macro_rules! encode_sequences_body { } of_state = next.index; } - if let Some(table) = ml_table { - let next = table.next_state(ml_code, ml_state); + { + let next = view.next_state::<1>(ml_code, ml_state); let diff = crate::fse::fse_encoder::transition_bits(ml_state, next.num_bits); unsafe { burst.add(diff as u64, next.num_bits as usize); @@ -2956,8 +2968,8 @@ macro_rules! encode_sequences_body { } ml_state = next.index; } - if let Some(table) = ll_table { - let next = table.next_state(ll_code, ll_state); + { + let next = view.next_state::<0>(ll_code, ll_state); let diff = crate::fse::fse_encoder::transition_bits(ll_state, next.num_bits); unsafe { burst.add(diff as u64, next.num_bits as usize); @@ -3002,10 +3014,10 @@ macro_rules! encode_sequences_body { // the leftover, the diffs and all three extra fields have // to fit. debug_assert!( - extra_bits_total >= 31 || 7 + state_diff_bits + extra_bits_total <= 64, + extra_bits_total >= 31 || 7 + state_diff_bits + extra_bits_total < 64, "no flush at {extra_bits_total} extra bits, but 7 leftover + \ - {state_diff_bits} diff bits + {extra_bits_total} would overflow the \ - 64-bit accumulator", + {state_diff_bits} diff bits + {extra_bits_total} would fill the \ + accumulator, which keeps a bit free", ); } if extra_bits_total >= 31 { @@ -3046,6 +3058,7 @@ macro_rules! encode_sequences_body { } burst.finish(); } + // An RLE stream has no state to write. if let Some(table) = ml_table { writer.write_bits(ml_state as u64, table.table_size.ilog2() as usize); } @@ -3178,6 +3191,20 @@ const ML_EXTRA_BITS: [u8; 53] = [ 1, 1, 1, 1, 2, 2, 3, 3, 4, 4, 5, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, ]; +/// A short length's code in the low byte and its extra-bit width in the high +/// byte, so one load answers both. +const fn fuse_code_bits(codes: &[u8; N], bits: &[u8]) -> [u16; N] { + let mut fused = [0u16; N]; + let mut i = 0; + while i < N { + fused[i] = codes[i] as u16 | (bits[codes[i] as usize] as u16) << 8; + i += 1; + } + fused +} +const LL_CODE_BITS: [u16; 64] = fuse_code_bits(&LL_CODE, &LL_EXTRA_BITS); +const ML_CODE_BITS: [u16; 128] = fuse_code_bits(&ML_CODE, &ML_EXTRA_BITS); + /// Split a literal-length into its FSE symbol, the extra bits to append, and /// how many of them (upstream `ZSTD_LLcode` plus the `LL_bits` lookup its /// caller does). @@ -3186,15 +3213,22 @@ const ML_EXTRA_BITS: [u8; 53] = [ /// compiled to a chain of comparisons; upstream reaches the same answer with a /// table index. Every code's baseline is a multiple of its own extra-bit width, /// so masking off those bits is the same subtraction the ranges spelled out. +/// From 64 up a code carries exactly `ilog2(len)` extra bits, so that branch +/// needs no table at all. #[inline] pub(crate) fn encode_literal_length(len: u32) -> (u8, u32, usize) { - debug_assert!(len < 131_072, "literal length {len} out of encodable range"); - let code = if len < 64 { - LL_CODE[len as usize] + let (code, bits) = if len < 64 { + let fused = LL_CODE_BITS[len as usize]; + (fused as u8, (fused >> 8) as usize) } else { - (len.ilog2() + LL_DELTA_CODE) as u8 + // RFC 8878 3.1.1.3.2.1.1: literal-length code 35 is the last, its + // baseline 65536 with 16 extra bits. + // No value in the message: formatting it keeps `len` on the stack + // across the whole calling loop. + assert!(len < 131_072, "literal length out of encodable range"); + let log = len.ilog2(); + ((log + LL_DELTA_CODE) as u8, log as usize) }; - let bits = LL_EXTRA_BITS[code as usize] as usize; (code, len & ((1u32 << bits) - 1), bits) } @@ -3205,17 +3239,18 @@ pub(crate) fn encode_literal_length(len: u32) -> (u8, u32, usize) { /// Table-driven for the same reason as [`encode_literal_length`]. #[inline] pub(crate) fn encode_match_len(len: u32) -> (u8, u32, usize) { - debug_assert!( - (3..131_075).contains(&len), - "match length {len} out of encodable range", - ); + debug_assert!(len >= 3, "match length {len} out of encodable range"); let base = len - 3; - let code = if base < 128 { - ML_CODE[base as usize] + let (code, bits) = if base < 128 { + let fused = ML_CODE_BITS[base as usize]; + (fused as u8, (fused >> 8) as usize) } else { - (base.ilog2() + ML_DELTA_CODE) as u8 + // RFC 8878 3.1.1.3.2.1.1: match-length code 52 is the last, its + // baseline 65539 with 16 extra bits. + assert!(base < 131_072, "match length out of encodable range"); + let log = base.ilog2(); + ((log + ML_DELTA_CODE) as u8, log as usize) }; - let bits = ML_EXTRA_BITS[code as usize] as usize; (code, base & ((1u32 << bits) - 1), bits) } diff --git a/zstd/src/encoding/blocks/compressed/literal_runs.rs b/zstd/src/encoding/blocks/compressed/literal_runs.rs index b36073c9..7bf63194 100644 --- a/zstd/src/encoding/blocks/compressed/literal_runs.rs +++ b/zstd/src/encoding/blocks/compressed/literal_runs.rs @@ -92,21 +92,29 @@ macro_rules! gather_body { let out = unsafe { dst.as_mut_ptr().add(start) }; let mut written = 0usize; let mut pos = 0usize; + // A run's length is a `u32`, but the tail is a matcher's unchecked sum; + // bounding it by the block keeps the single-compare room test below + // from wrapping. Once per block. + assert!( + $tail <= src_len, + "sequences reach past the block they describe" + ); macro_rules! copy_run { ($len:expr) => {{ let len = $len; - // `pos <= src_len` holds on every iteration. - let room = src_len - pos; // SAFETY (both arms): `written <= pos`, and the capacity check in // `gather_literals` gives the buffer `src_len` bytes past // `start`. The chunked arm reads and writes at most - // `max($chunk, len + $chunk - 1) <= len + WILDCOPY_SLACK <= room` - // bytes from `pos` / `written` (`$chunk <= WILDCOPY_SLACK`); the - // exact arm, `len <= room` bytes. + // `max($chunk, len + $chunk - 1) <= len + WILDCOPY_SLACK` bytes + // from `pos` / `written` (`$chunk <= WILDCOPY_SLACK`), inside the + // block by its guard; the exact arm, `len <= src_len - pos`. unsafe { let s = src.add(pos); let d = out.add(written); - if room >= WILDCOPY_SLACK && len <= room - WILDCOPY_SLACK { + // One compare in 64 bits, where no operand can overflow + // (a run is a `u32`, `pos` and the tail are bounded by the + // block), rather than a room test and a length test. + if pos as u64 + len as u64 + WILDCOPY_SLACK as u64 <= src_len as u64 { // Upstream's shape (`ZSTD_wildcopy`): one chunk // unconditionally, which is the whole run for most of // them, and a loop only for the rest of a longer one. @@ -126,7 +134,11 @@ macro_rules! gather_body { core::ptr::copy_nonoverlapping(s, d, len); } } else { - assert!(len <= room, "sequences reach past the block they describe"); + // `pos <= src_len` holds on every iteration. + assert!( + len <= src_len - pos, + "sequences reach past the block they describe" + ); if len <= 32 { if len != 0 { simd_copy::copy_exact_small(s, d, len); diff --git a/zstd/src/encoding/blocks/compressed/tests.rs b/zstd/src/encoding/blocks/compressed/tests.rs index 29a2bce0..5d07ad14 100644 --- a/zstd/src/encoding/blocks/compressed/tests.rs +++ b/zstd/src/encoding/blocks/compressed/tests.rs @@ -1246,6 +1246,21 @@ fn sequences_past_the_block_are_refused() { ); } +/// A trailing literal count past the block is refused before any copy, however +/// large: a custom matcher reports it unchecked, and near `usize::MAX` the +/// chunked copy's room test would wrap and let a copy of that length through. +#[test] +#[should_panic(expected = "sequences reach past the block they describe")] +fn a_tail_past_the_block_is_refused_before_the_copy() { + let (_ws, mut literals) = literals_buffer(64); + super::literal_runs::gather_literals::( + &[0x5A; 64], + &[], + usize::MAX - 16, + &mut literals, + ); +} + /// Every literal-gather tier this CPU can run, not only the one it selects, so /// a narrower tier's path is exercised on a wider machine. fn runnable_fastpath_kernels() -> Vec { diff --git a/zstd/src/encoding/frame_compressor.rs b/zstd/src/encoding/frame_compressor.rs index c65159e8..bb1fe944 100644 --- a/zstd/src/encoding/frame_compressor.rs +++ b/zstd/src/encoding/frame_compressor.rs @@ -889,7 +889,9 @@ fn presplit_record_fingerprint( /// 2-byte hashing for `hashLog >= 9`; this helper exists so the borders /// heuristic doesn't pay for that wider hash on its 512-byte windows. fn presplit_record_byte_histogram(fp: &mut PreSplitFingerprint, src: &[u8]) { - fp.events.fill(0); + // Only the byte slots: a `hash_log == 8` fingerprint is recorded, compared + // and read in those alone. + fp.events[..256].fill(0); for &b in src { fp.events[b as usize] += 1; } diff --git a/zstd/src/fse/fse_encoder.rs b/zstd/src/fse/fse_encoder.rs index f3f35b79..a039aad4 100644 --- a/zstd/src/fse/fse_encoder.rs +++ b/zstd/src/fse/fse_encoder.rs @@ -315,6 +315,75 @@ impl FSETable { } } +/// Symbols a sequence-code stream can carry: the widest alphabet, match +/// lengths, has 53 codes (RFC 8878 3.1.1.3.2.1.1), under 64. +const SEQUENCE_CODE_SYMBOLS: usize = 64; + +/// The parts of the three sequence streams' tables the sequence loop reads — +/// transforms, next states, sizes — side by side at fixed offsets, so the +/// loop reaches all three from one base and the registers three table +/// pointers held stay free for the stream states, which are on the loop's +/// dependency chain. A stream in RLE mode is a one-state table whose +/// transitions write nothing, as upstream's `FSE_buildCTable_rle`. +pub(crate) struct SequenceStateTables { + symbol_tt: [[SymbolTT; SEQUENCE_CODE_SYMBOLS]; 3], + next: [[u16; MAX_FSE_TABLE_SIZE]; 3], + table_size: [usize; 3], +} + +impl SequenceStateTables { + pub(crate) fn new() -> Self { + Self { + symbol_tt: [[SymbolTT::default(); SEQUENCE_CODE_SYMBOLS]; 3], + next: [[0; MAX_FSE_TABLE_SIZE]; 3], + table_size: [0; 3], + } + } + + /// Take stream `stream`'s parts from `table`, or an RLE stream's + /// one-state table for `None`. + pub(crate) fn set(&mut self, stream: usize, table: Option<&FSETable>) { + match table { + Some(table) => { + self.symbol_tt[stream].copy_from_slice(&table.symbol_tt[..SEQUENCE_CODE_SYMBOLS]); + let size = table.table_size; + self.next[stream][..size].copy_from_slice(&table.state_table_flat[..size]); + self.table_size[stream] = size; + } + None => { + // One state: `nb_bits = (1 + 0) >> 16 = 0`, slot `1 - 1 = 0`. + self.symbol_tt[stream] = [SymbolTT { + delta_nb_bits: 0, + delta_find_state: -1, + }; SEQUENCE_CODE_SYMBOLS]; + self.next[stream][0] = 0; + self.table_size[stream] = 1; + } + } + } + + /// Stream `S`'s transition on `symbol` out of state `idx`, as + /// [`FSETable::next_state`]. + #[inline(always)] + pub(crate) fn next_state(&self, symbol: u8, idx: usize) -> State { + let tt = self.symbol_tt[S][symbol as usize % SEQUENCE_CODE_SYMBOLS]; + let size = self.table_size[S]; + let value = (size + idx) as u32; + let nb_bits = ((value + tt.delta_nb_bits) >> 16) as usize; + let slot = ((value >> nb_bits) as i32 + tt.delta_find_state) as usize; + debug_assert!( + slot < size, + "next-state slot {slot} outside a {size}-entry table" + ); + State { + num_bits: nb_bits as u8, + // SAFETY: `slot < table_size <= MAX_FSE_TABLE_SIZE` (asserted + // above, held by the build that filled this stream). + index: unsafe { *self.next[S].get_unchecked(slot) } as usize, + } + } +} + /// Write the NCount description of the normalized distribution of `len` /// symbols, `prob(i)` the probability of symbol `i`, at `acc_log` (RFC 8878 /// 4.1.1), padded to a byte. The probabilities are read where they live, so a diff --git a/zstd/src/huff0/huf_cstream.rs b/zstd/src/huff0/huf_cstream.rs index 66d21614..c9912b7f 100644 --- a/zstd/src/huff0/huf_cstream.rs +++ b/zstd/src/huff0/huf_cstream.rs @@ -135,11 +135,12 @@ macro_rules! huf_flush { ($bc:ident, $bp:ident, $cursor:ident, $overflow:ident, $out_base:ident, $end_ptr:ident, $fast:expr) => {{ let nb_bits = ($bp & 0xFF) as usize; let nb_bytes = nb_bits >> 3; - let chunk = if nb_bits == 0 { - 0 - } else { - $bc >> (HUF_BITS_IN_CONTAINER - nb_bits) - }; + // Upstream's `bitContainer >> (64 - nbBits)`, one masked shift: with no + // bits pending the shift wraps to 0 and the store writes the container + // as is, but the cursor does not move, so the next flush overwrites + // those bytes and none of them is committed. A zero test here sat on + // the store's data path. + let chunk = $bc.wrapping_shr((HUF_BITS_IN_CONTAINER - nb_bits) as u32); $bp &= 7; let bytes = chunk.to_le_bytes(); unsafe { @@ -301,14 +302,12 @@ impl<'a> HufCStream<'a> { pub(crate) fn flush_bits(&mut self) { let nb_bits = (self.bit_pos[0] & 0xFF) as usize; let nb_bytes = nb_bits >> 3; - // Top `nb_bits` of the container become the next bytes. - // Upstream zstd uses `bitContainer >> (HUF_BITS_IN_CONTAINER - nb_bits)`. - // Guard the shift: `nb_bits == 0` would shift by 64 (UB in Rust). - let bit_container = if nb_bits == 0 { - 0 - } else { - self.bit_container[0] >> (HUF_BITS_IN_CONTAINER - nb_bits) - }; + // Top `nb_bits` of the container become the next bytes, as upstream's + // `bitContainer >> (HUF_BITS_IN_CONTAINER - nb_bits)`. With no bits + // pending the masked shift leaves the container whole; the cursor does + // not move, so those bytes are never committed (see `huf_flush!`). + let bit_container = + self.bit_container[0].wrapping_shr((HUF_BITS_IN_CONTAINER - nb_bits) as u32); // Mask `bit_pos` to keep the leftover < 8 bits in the low 3 bits. self.bit_pos[0] &= 7; // 8-byte LE write at `cursor`. Bytes at [cursor+nb_bytes..cursor+8] diff --git a/zstd/src/huff0/huff0_encoder.rs b/zstd/src/huff0/huff0_encoder.rs index 913598bd..8f76242c 100644 --- a/zstd/src/huff0/huff0_encoder.rs +++ b/zstd/src/huff0/huff0_encoder.rs @@ -1057,11 +1057,15 @@ impl HuffmanTable { /// probes, per split candidate; most are measured and thrown away, so /// handing the last one back turns two allocations per build into none. pub(crate) fn build_from_weights_reusing(weights: &[usize], spare: Option) -> Self { - let weight_sum = weights - .iter() - .copied() - .filter(|&weight| weight > 0) - .map(|weight| 1 << (weight - 1)) + // Symbols per weight. Everything below needs only these thirteen + // counts: the weight sum that gives the table log, and the symbols + // per code length that give each length's first code. + let mut per_weight = [0u16; 13]; + for &weight in weights { + per_weight[weight] += 1; + } + let weight_sum = (1..per_weight.len()) + .map(|weight| usize::from(per_weight[weight]) << (weight - 1)) .sum::(); if !weight_sum.is_power_of_two() { panic!("This is an internal error"); @@ -1069,17 +1073,19 @@ impl HuffmanTable { let table_log = highest_bit_set(weight_sum) - 1; let mut table = match spare { Some(mut spare) => { - // Both buffers are overwritten in full below, so the recycled - // contents do not survive; only their allocations do. - // Exact reservations: an amortised `resize` would double a - // buffer grown one symbol at a time past the 256 entries any - // table holds, beyond what `max_heap_size` budgets. + // Both buffers are overwritten below, so the recycled contents + // do not survive; only their allocations do. Exact + // reservation: an amortised one would double a buffer grown + // one symbol at a time past the 256 entries any table holds, + // beyond what `max_heap_size` budgets. + // The fill below writes the first `weights.len()` packed codes, + // absent symbols as zero; past them only the previous table's + // symbols can be non-zero. + let stale = spare.codes.len().max(weights.len()); + debug_assert_eq!(spare.packed_codes.len(), 256); + spare.packed_codes[weights.len()..stale].fill(0); spare.codes.clear(); spare.codes.reserve_exact(weights.len()); - spare.codes.resize(weights.len(), (0, 0)); - spare.packed_codes.clear(); - spare.packed_codes.reserve_exact(256); - spare.packed_codes.resize(256, 0u64); spare.table_log = table_log as u32; #[cfg(feature = "std")] { @@ -1091,37 +1097,35 @@ impl HuffmanTable { spare } None => HuffmanTable { - codes: alloc::vec![(0, 0); weights.len()], + codes: Vec::with_capacity(weights.len()), packed_codes: alloc::vec![0u64; 256], table_log: table_log as u32, #[cfg(feature = "std")] cached_encoded_weight_description: WeightDescription::default(), }, }; - let mut nb_per_rank = [0u16; 13]; - for &weight in weights { - if weight > 0 { - let nb_bits = table_log + 1 - weight; - nb_per_rank[nb_bits] += 1; - } - } + // A code of `nb_bits` bits has weight `table_log + 1 - nb_bits`. let mut val_per_rank = [0u16; 13]; let mut min = 0u16; for nb_bits in (1..=table_log).rev() { val_per_rank[nb_bits] = min; - min = min.wrapping_add(nb_per_rank[nb_bits]) >> 1; - } - for (symbol, &weight) in weights.iter().enumerate() { - if weight == 0 { - continue; - } - let nb_bits = table_log + 1 - weight; - let value = val_per_rank[nb_bits]; - val_per_rank[nb_bits] += 1; - table.codes[symbol] = (value as u32, nb_bits as u8); - table.packed_codes[symbol] = - super::huf_cstream::pack_huf_celt(value as u32, nb_bits as u8); + min = min.wrapping_add(per_weight[table_log + 1 - nb_bits]) >> 1; } + let packed_codes = &mut table.packed_codes; + table + .codes + .extend(weights.iter().enumerate().map(|(symbol, &weight)| { + if weight == 0 { + packed_codes[symbol] = 0; + return (0, 0); + } + let nb_bits = table_log + 1 - weight; + let value = val_per_rank[nb_bits]; + val_per_rank[nb_bits] += 1; + packed_codes[symbol] = + super::huf_cstream::pack_huf_celt(value as u32, nb_bits as u8); + (value as u32, nb_bits as u8) + })); table } @@ -1564,14 +1568,15 @@ fn build_huffman_leaf_depths_into(counts: &[usize], nodes: &mut Vec) { node_nb += 1; } - for leaf_idx in 0..leaf_count { - let mut depth = 0u8; - let mut parent = nodes[leaf_idx].parent; - while parent != NO_PARENT { - depth += 1; - parent = nodes[parent as usize].parent; - } - nodes[leaf_idx].nb_bits = depth; + // Depths top-down in one pass, as upstream's `HUF_buildTree`: every node's + // parent was built after it, so it sits at a higher index and is already + // done when the descending walk reaches the child. Walking each leaf up to + // the root instead was a chain of dependent loads per leaf. + nodes[node_root].nb_bits = 0; + for n in (0..node_root).rev() { + let parent = nodes[n].parent as usize; + debug_assert!(parent > n && parent <= node_root); + nodes[n].nb_bits = nodes[parent].nb_bits + 1; } // The leaves keep their (count desc, symbol asc) order: the tree build