structured-zstd 0.2.6

use alloc::boxed::Box;
#[cfg(not(target_has_atomic = "ptr"))]
use alloc::rc::Rc;
#[cfg(target_has_atomic = "ptr")]
use alloc::sync::Arc;
use alloc::vec::Vec;
use core::convert::TryInto;

use crate::decoding::errors::DictionaryDecodeError;
use crate::decoding::scratch::FSEScratch;
use crate::decoding::scratch::HuffmanScratch;

/// Zstandard includes support for "raw content" dictionaries, that store bytes optionally used
/// during sequence execution.
///
/// <https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#dictionary-format>
#[derive(Clone)]
pub struct Dictionary {
    /// A 4 byte value used by decoders to check if they can use
    /// the correct dictionary.
    ///
    /// Zero means unidentified: a raw-content dictionary has no header to
    /// carry an ID, and the frames built from it record none, so it can only
    /// be supplied explicitly and never resolved from a frame header.
    /// Registration by ID
    /// ([`FrameDecoder::add_dict`](crate::decoding::FrameDecoder::add_dict))
    /// therefore still requires a non-zero one.
    pub id: u32,
    /// A dictionary can contain an entropy table, either FSE or
    /// Huffman. Boxed for the same reason as `DecoderScratch::fse`: the
    /// ~12 KiB of inline tables must not ride the stack through
    /// construction.
    pub fse: Box<FSEScratch>,
    /// A dictionary can contain an entropy table, either FSE or
    /// Huffman.
    pub huf: HuffmanScratch,
    /// The content of a dictionary acts as a "past" in front of data
    /// to compress or decompress,
    /// so it can be referenced in sequence commands.
    /// As long as the amount of data decoded from this frame is less than or
    /// equal to Window_Size, sequence commands may specify offsets longer than
    /// the total length of decoded output so far to reference back to the
    /// dictionary, even parts of the dictionary with offsets larger than Window_Size.
    /// After the total output has surpassed Window_Size however,
    /// this is no longer allowed and the dictionary is no longer accessible
    pub dict_content: Vec<u8>,
    /// The 3 most recent offsets are stored so that they can be used
    /// during sequence execution, see
    /// <https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#repeat-offsets>
    /// for more.
    pub offset_hist: [u32; 3],
}

/// What a [`DictionaryHandle`] shares: the parsed dictionary and, for one
/// parsed from serialized bytes, the bytes in front of its content (magic
/// number, ID, tables, repeat offsets). A v0.5 or v0.6 frame takes a
/// dictionary of this layout whole as history, header too.
pub(crate) struct HeldDictionary {
    pub(crate) dict: Dictionary,
    #[cfg(feature = "legacy")]
    pub(crate) serialized_header: Vec<u8>,
}

/// A parsed dictionary held by however many users need it at once.
///
/// Both sides prime frame after frame from one dictionary, so what they hold is
/// shared rather than copied: `Arc` where atomics exist, `Rc` where they do not.
#[cfg(target_has_atomic = "ptr")]
pub(crate) type SharedDictionary = Arc<HeldDictionary>;
#[cfg(not(target_has_atomic = "ptr"))]
pub(crate) type SharedDictionary = Rc<HeldDictionary>;

/// Shared pre-parsed dictionary handle for repeated decoding.
///
/// Uses `Arc` on targets with atomics and falls back to `Rc` otherwise.
#[derive(Clone)]
pub struct DictionaryHandle {
    inner: SharedDictionary,
}

/// This 4 byte (little endian) magic number refers to the start of a dictionary
pub const MAGIC_NUM: [u8; 4] = [0x37, 0xA4, 0x30, 0xEC];

/// Most bytes [`Dictionary::decode_dict`] accepts in front of a dictionary's
/// content: magic number and ID, a Huffman description (its size byte and at
/// most 127 more), the three FSE descriptions at their largest, and the
/// repeat offsets.
#[cfg(feature = "legacy")]
const MAX_SERIALIZED_HEADER_LEN: usize = 8
    + 128
    + max_fse_description_len(
        crate::blocks::sequence_section::MAX_OFFSET_CODE,
        crate::decoding::sequence_section_decoder::OF_MAX_LOG,
    )
    + max_fse_description_len(
        crate::blocks::sequence_section::MAX_MATCH_LENGTH_CODE,
        crate::decoding::sequence_section_decoder::ML_MAX_LOG,
    )
    + max_fse_description_len(
        crate::blocks::sequence_section::MAX_LITERAL_LENGTH_CODE,
        crate::decoding::sequence_section_decoder::LL_MAX_LOG,
    )
    + 12;

/// Most bytes an FSE description of symbols `0..=max_symbol` at accuracy log
/// `max_log` takes: the 4-bit log, then per symbol a count of at most
/// `max_log + 1` bits and, amortized, at most 2 bits of zero-run flags.
#[cfg(feature = "legacy")]
const fn max_fse_description_len(max_symbol: u8, max_log: u8) -> usize {
    (4 + (max_symbol as usize + 1) * (max_log as usize + 3)).div_ceil(8)
}

impl Dictionary {
    /// Heap bytes owned by this dictionary: the content plus the parsed
    /// entropy tables' heap. `fse` is a `Box`, so the fixed-size FSE decode
    /// arrays are heap bytes rather than inline in `size_of::<Dictionary>()`;
    /// charge that `size_of` here so a caller budgeting against the reported
    /// figure sees them.
    pub fn heap_bytes(&self) -> usize {
        self.dict_content.capacity()
            + core::mem::size_of::<FSEScratch>()
            + self.fse.heap_bytes()
            + self.huf.heap_bytes()
    }

    /// Most [`heap_bytes`](Self::heap_bytes) a dictionary holding
    /// `content_len` bytes of content can report, so a caller can budget for
    /// one before parsing it.
    ///
    /// # Examples
    ///
    /// ```
    /// use structured_zstd::decoding::Dictionary;
    ///
    /// let dict = Dictionary::from_raw_content(0, vec![7u8; 1024]).unwrap();
    /// assert!(dict.heap_bytes() <= Dictionary::max_heap_bytes(1024));
    /// ```
    pub const fn max_heap_bytes(content_len: usize) -> usize {
        content_len
            + core::mem::size_of::<FSEScratch>()
            + FSEScratch::max_heap_bytes()
            + HuffmanScratch::max_heap_bytes()
    }

    /// Build a dictionary from raw content bytes (without entropy table sections).
    ///
    /// This is primarily intended for dictionaries produced by the `dict-builder`
    /// module, which currently emits raw-content dictionaries.
    ///
    /// An `id` of 0 means the dictionary is unidentified, which is what any
    /// plain file used as a dictionary is: frames built with it record no
    /// dictionary ID, so it can only ever be supplied explicitly, never
    /// resolved from a frame header. Registration by ID
    /// ([`FrameDecoder::add_dict`](crate::decoding::FrameDecoder::add_dict))
    /// still requires a non-zero one, since the ID is the key it is stored
    /// under.
    pub fn from_raw_content(
        id: u32,
        dict_content: Vec<u8>,
    ) -> Result<Dictionary, DictionaryDecodeError> {
        if dict_content.is_empty() {
            return Err(DictionaryDecodeError::DictionaryTooSmall { got: 0, need: 1 });
        }

        Ok(Dictionary {
            id,
            fse: Box::new(FSEScratch::new()),
            huf: HuffmanScratch::new(),
            dict_content,
            offset_hist: [1, 4, 8],
        })
    }

    /// Parses the dictionary from `raw`, initializes its tables,
    /// and returns a fully constructed [`Dictionary`] whose `id` can be
    /// checked against the frame's `dict_id`.
    pub fn decode_dict(raw: &[u8]) -> Result<Dictionary, DictionaryDecodeError> {
        Self::decode_dict_inner(raw, true)
    }

    /// Whether `raw` is a serialized dictionary rather than raw content: it
    /// starts with [`MAGIC_NUM`] and has room for the ID after it. A shorter
    /// blob is raw content whatever it starts with, as upstream zstd reads it
    /// (zstd_ddict.c:97, zstd_compress.c:5206).
    ///
    /// ```
    /// use structured_zstd::decoding::{Dictionary, DICTIONARY_MAGIC};
    /// assert!(!Dictionary::is_serialized(&DICTIONARY_MAGIC));
    /// assert!(Dictionary::is_serialized(&[0x37, 0xA4, 0x30, 0xEC, 7, 0, 0, 0]));
    /// ```
    pub fn is_serialized(raw: &[u8]) -> bool {
        raw.len() >= 8 && raw.starts_with(&MAGIC_NUM)
    }

    /// Loads whichever kind of dictionary `raw` holds, the way `zstd -D` does:
    /// a [serialized](Self::is_serialized) dictionary with entropy tables and
    /// an ID, and anything else is taken as raw content,
    /// which is why any file can be handed to `-D`. A raw-content dictionary
    /// has no ID, so it must be supplied explicitly on both sides — see
    /// [`Self::from_raw_content`].
    pub fn from_serialized_or_raw_content(raw: &[u8]) -> Result<Dictionary, DictionaryDecodeError> {
        if Self::is_serialized(raw) {
            Self::decode_dict(raw)
        } else if raw.is_empty() {
            // A zero-sized buffer is a dictionary with nothing in it rather
            // than a malformed one: `ZSTD_createDDict(NULL, 0)` builds a
            // usable `DDict` referencing no content, and
            // `ZSTD_CCtx_loadDictionary` with an empty buffer is how a caller
            // says "no dictionary". [`Self::from_raw_content`] still refuses
            // it, because naming raw content and handing over none is the
            // caller asking for a dictionary that cannot exist.
            Ok(Dictionary {
                id: 0,
                fse: Box::new(FSEScratch::new()),
                huf: HuffmanScratch::new(),
                dict_content: Vec::new(),
                offset_hist: [1, 4, 8],
            })
        } else {
            Self::from_raw_content(0, raw.to_vec())
        }
    }

    /// Parse a dictionary for ENCODER use: builds the entropy
    /// probabilities/weights needed by `to_encoder_table` but skips the
    /// decode-only work the encoder never reads — the FSE *decoding*
    /// tables + their `enrich_*` post-passes, and the HUF decode lookup
    /// table (`packed_decode`). Produces a [`Dictionary`] whose FSE
    /// `symbol_probabilities` / `accuracy_log` and HUF `bits` /
    /// `max_num_bits` match `decode_dict` exactly, so the encoder entropy
    /// tables — and thus the emitted frame — are byte-identical; only the
    /// wasted decode-table builds are dropped. Offset history + content
    /// are parsed the same way.
    /// Crate-internal: the returned [`Dictionary`] deliberately has no
    /// decode lookup tables (`packed_decode` / FSE `decode`), so it is
    /// NOT safe to feed into a [`FrameDecoder`](crate::decoding::FrameDecoder)
    /// — Huffman decode would index an empty `packed_decode`. The only caller
    /// is `EncoderDictionary::from_bytes`, which wraps the result in the
    /// encoder-only `EncoderDictionary` type (no decode path), so this
    /// incomplete dictionary can never escape to the decode side. Keeping
    /// this `pub(crate)` keeps it off the public `Dictionary` API entirely.
    pub(crate) fn decode_dict_for_encoding(
        raw: &[u8],
    ) -> Result<Dictionary, DictionaryDecodeError> {
        Self::decode_dict_inner(raw, false)
    }

    /// Shared dictionary parser. `build_decode_tables` selects whether the
    /// FSE/HUF tables get their full decoding tables (FSE decode table +
    /// `enrich_*`, HUF `packed_decode`; decoder path) or only the
    /// probability/weight parse (encoder path — see
    /// [`Self::decode_dict_for_encoding`]).
    fn decode_dict_inner(
        raw: &[u8],
        build_decode_tables: bool,
    ) -> Result<Dictionary, DictionaryDecodeError> {
        const MIN_MAGIC_AND_ID_LEN: usize = 8;
        const OFFSET_HISTORY_LEN: usize = 12;

        if raw.len() < MIN_MAGIC_AND_ID_LEN {
            return Err(DictionaryDecodeError::DictionaryTooSmall {
                got: raw.len(),
                need: MIN_MAGIC_AND_ID_LEN,
            });
        }

        let mut new_dict = Dictionary {
            id: 0,
            fse: Box::new(FSEScratch::new()),
            huf: HuffmanScratch::new(),
            dict_content: Vec::new(),
            offset_hist: [1, 4, 8],
        };

        let magic_num: [u8; 4] = raw[..4].try_into().expect("optimized away");
        if magic_num != MAGIC_NUM {
            return Err(DictionaryDecodeError::BadMagicNum { got: magic_num });
        }

        let dict_id = raw[4..8].try_into().expect("optimized away");
        let dict_id = u32::from_le_bytes(dict_id);
        if dict_id == 0 {
            return Err(DictionaryDecodeError::ZeroDictionaryId);
        }
        new_dict.id = dict_id;

        let raw_tables = &raw[8..];

        let huf_size = if build_decode_tables {
            new_dict.huf.table.build_decoder(raw_tables)?
        } else {
            new_dict.huf.table.build_weights_only(raw_tables)?
        };
        let raw_tables = &raw_tables[huf_size as usize..];

        let of_size = if build_decode_tables {
            let n = new_dict.fse.offsets.build_decoder(
                raw_tables,
                crate::decoding::sequence_section_decoder::OF_MAX_LOG,
            )?;
            new_dict.fse.offsets.enrich_for_offsets();
            // Compute the pipeline-gate long-offset share ONCE here, while the
            // dictionary handle is built, so the per-decode `init_from_dict`
            // path can COPY it instead of re-walking the offsets table on every
            // `decode_*_with_dict_handle` call (the dict is immutable, so the
            // share never changes after this).
            new_dict.fse.offsets_long_share =
                crate::decoding::sequence_section_decoder::compute_offsets_long_share(
                    &new_dict.fse.offsets,
                );
            n
        } else {
            new_dict.fse.offsets.read_table_probabilities(
                raw_tables,
                crate::decoding::sequence_section_decoder::OF_MAX_LOG,
            )?
        };
        let raw_tables = &raw_tables[of_size..];

        let ml_size = if build_decode_tables {
            let n = new_dict.fse.match_lengths.build_decoder(
                raw_tables,
                crate::decoding::sequence_section_decoder::ML_MAX_LOG,
            )?;
            new_dict
                .fse
                .match_lengths
                .enrich_with_packed_seq_meta(&crate::decoding::sequence_section_decoder::ML_META);
            n
        } else {
            new_dict.fse.match_lengths.read_table_probabilities(
                raw_tables,
                crate::decoding::sequence_section_decoder::ML_MAX_LOG,
            )?
        };
        let raw_tables = &raw_tables[ml_size..];

        let ll_size = if build_decode_tables {
            let n = new_dict.fse.literal_lengths.build_decoder(
                raw_tables,
                crate::decoding::sequence_section_decoder::LL_MAX_LOG,
            )?;
            new_dict
                .fse
                .literal_lengths
                .enrich_with_packed_seq_meta(&crate::decoding::sequence_section_decoder::LL_META);
            n
        } else {
            new_dict.fse.literal_lengths.read_table_probabilities(
                raw_tables,
                crate::decoding::sequence_section_decoder::LL_MAX_LOG,
            )?
        };
        let raw_tables = &raw_tables[ll_size..];

        if raw_tables.len() < OFFSET_HISTORY_LEN {
            return Err(DictionaryDecodeError::DictionaryTooSmall {
                got: raw_tables.len(),
                need: OFFSET_HISTORY_LEN,
            });
        }

        let offset1 = raw_tables[0..4].try_into().expect("optimized away");
        let offset1 = u32::from_le_bytes(offset1);

        let offset2 = raw_tables[4..8].try_into().expect("optimized away");
        let offset2 = u32::from_le_bytes(offset2);

        let offset3 = raw_tables[8..12].try_into().expect("optimized away");
        let offset3 = u32::from_le_bytes(offset3);

        if offset1 == 0 {
            return Err(DictionaryDecodeError::ZeroRepeatOffsetInDictionary { index: 0 });
        }
        if offset2 == 0 {
            return Err(DictionaryDecodeError::ZeroRepeatOffsetInDictionary { index: 1 });
        }
        if offset3 == 0 {
            return Err(DictionaryDecodeError::ZeroRepeatOffsetInDictionary { index: 2 });
        }

        new_dict.offset_hist[0] = offset1;
        new_dict.offset_hist[1] = offset2;
        new_dict.offset_hist[2] = offset3;

        let raw_content = &raw_tables[12..];
        new_dict.dict_content.extend(raw_content);

        Ok(new_dict)
    }

    /// Convert this parsed dictionary into a reusable shared handle.
    pub fn into_handle(self) -> DictionaryHandle {
        DictionaryHandle::from_dictionary(self)
    }
}

impl DictionaryHandle {
    /// Wrap an already-parsed dictionary in a shared handle.
    pub fn from_dictionary(dict: Dictionary) -> Self {
        Self {
            inner: SharedDictionary::new(HeldDictionary {
                dict,
                #[cfg(feature = "legacy")]
                serialized_header: Vec::new(),
            }),
        }
    }

    /// The bytes of the serialized dictionary in front of its content, empty
    /// for one built from raw content or from a parsed [`Dictionary`].
    #[cfg(feature = "legacy")]
    pub(crate) fn serialized_header(&self) -> &[u8] {
        &self.inner.serialized_header
    }

    /// Heap bytes the handle keeps alive: the shared allocation the
    /// dictionary lives in and the dictionary's own content and tables.
    /// Clones share all of it.
    ///
    /// # Examples
    ///
    /// ```
    /// use structured_zstd::decoding::{Dictionary, DictionaryHandle};
    ///
    /// let handle = DictionaryHandle::from_dictionary(
    ///     Dictionary::from_raw_content(0, vec![7u8; 1024]).unwrap(),
    /// );
    /// assert!(handle.heap_size() > handle.as_dict().heap_bytes());
    /// assert!(handle.heap_size() <= DictionaryHandle::max_heap_size(1024));
    /// ```
    pub fn heap_size(&self) -> usize {
        let bytes = Self::shared_allocation_size() + self.inner.dict.heap_bytes();
        #[cfg(feature = "legacy")]
        let bytes = bytes + self.inner.serialized_header.capacity();
        bytes
    }

    /// Most [`Self::heap_size`] a handle over a dictionary of `content_len`
    /// bytes can report, serialized header included, for budgeting before
    /// the dictionary is parsed.
    pub fn max_heap_size(content_len: usize) -> usize {
        let bytes = Self::shared_allocation_size() + Dictionary::max_heap_bytes(content_len);
        #[cfg(feature = "legacy")]
        let bytes = bytes + MAX_SERIALIZED_HEADER_LEN;
        bytes
    }

    /// Heap bytes that `handles` alone keep alive: each distinct dictionary
    /// among them once, and only when no handle outside them shares it. For an
    /// owner holding several handles, such as a decoding context that keeps a
    /// dictionary attached and its decoder keeps the last frame's.
    ///
    /// # Examples
    ///
    /// ```
    /// use structured_zstd::decoding::{Dictionary, DictionaryHandle};
    ///
    /// let handle = DictionaryHandle::from_dictionary(
    ///     Dictionary::from_raw_content(0, vec![7u8; 64]).unwrap(),
    /// );
    /// let copy = handle.clone();
    /// // Both handles are in the set: counted once.
    /// assert_eq!(
    ///     DictionaryHandle::exclusive_heap_size([&handle, &copy]),
    ///     handle.heap_size()
    /// );
    /// // `copy` is held elsewhere: nothing is this set's alone.
    /// assert_eq!(DictionaryHandle::exclusive_heap_size([&handle]), 0);
    /// ```
    pub fn exclusive_heap_size<'a, I>(handles: I) -> usize
    where
        I: IntoIterator<Item = &'a DictionaryHandle>,
        I::IntoIter: Clone,
    {
        let handles = handles.into_iter();
        let mut total = 0;
        for (position, handle) in handles.clone().enumerate() {
            let same =
                |other: &&DictionaryHandle| SharedDictionary::ptr_eq(&handle.inner, &other.inner);
            // The first handle to an allocation speaks for all of them.
            if handles.clone().take(position).any(|other| same(&other)) {
                continue;
            }
            // Holders, not references: one handle listed twice is one count.
            let held_here = handles
                .clone()
                .enumerate()
                .filter(|(at, other)| {
                    same(other)
                        && !handles
                            .clone()
                            .take(*at)
                            .any(|earlier| core::ptr::eq(earlier, *other))
                })
                .count();
            if SharedDictionary::strong_count(&handle.inner) == held_here {
                total += handle.heap_size();
            }
        }
        total
    }

    /// Bytes of the shared allocation itself: `Arc` and `Rc` both place two
    /// word-sized counts ahead of the value, padded to its alignment.
    fn shared_allocation_size() -> usize {
        core::alloc::Layout::new::<[usize; 2]>()
            .extend(core::alloc::Layout::new::<HeldDictionary>())
            .expect("a fixed-size layout extends")
            .0
            .pad_to_align()
            .size()
    }

    /// Parse a serialized dictionary and return a reusable shared handle.
    ///
    /// With the `legacy` feature a dictionary in the v0.5 or v0.6 layout is
    /// taken too, whole as content with no ID: frames of its format parse
    /// its tables from it, as libzstd's `ZSTD_decompress_usingDict` does.
    pub fn decode_dict(raw: &[u8]) -> Result<Self, DictionaryDecodeError> {
        #[cfg(feature = "legacy")]
        if crate::decoding::legacy::is_older_dictionary(raw) {
            return Dictionary::from_raw_content(0, raw.to_vec()).map(Self::from_dictionary);
        }
        let dict = Dictionary::decode_dict(raw)?;
        // The content is the tail of the serialized bytes.
        #[cfg(feature = "legacy")]
        let serialized_header = raw[..raw.len() - dict.dict_content.len()].to_vec();
        #[cfg(feature = "legacy")]
        debug_assert!(serialized_header.capacity() <= MAX_SERIALIZED_HEADER_LEN);
        Ok(Self {
            inner: SharedDictionary::new(HeldDictionary {
                dict,
                #[cfg(feature = "legacy")]
                serialized_header,
            }),
        })
    }

    /// Load whichever kind of dictionary `raw` holds, as `ZSTD_createDDict`
    /// does: a [serialized](Dictionary::is_serialized) dictionary, or raw
    /// content. See
    /// [`Dictionary::from_serialized_or_raw_content`].
    pub fn from_serialized_or_raw_content(raw: &[u8]) -> Result<Self, DictionaryDecodeError> {
        if Dictionary::is_serialized(raw) {
            Self::decode_dict(raw)
        } else {
            Dictionary::from_serialized_or_raw_content(raw).map(Self::from_dictionary)
        }
    }

    pub fn id(&self) -> u32 {
        self.inner.dict.id
    }

    pub fn as_dict(&self) -> &Dictionary {
        &self.inner.dict
    }
}

impl AsRef<Dictionary> for DictionaryHandle {
    fn as_ref(&self) -> &Dictionary {
        self.as_dict()
    }
}

impl From<Dictionary> for DictionaryHandle {
    fn from(dict: Dictionary) -> Self {
        DictionaryHandle::from_dictionary(dict)
    }
}

#[cfg(test)]
mod tests;