zstd_complete/encoding/

workspace.rs

//! Reusable slice-to-slice compression workspaces.

use alloc::vec::Vec;
use core::{fmt, marker::PhantomData, mem::MaybeUninit};

use super::{levels::c_port, CompressionLevel};
use crate::workspace::{Arena, ArenaError, ArenaSize, ReusableVec};

use c_port::{
    WorkspaceCctxParameters as CctxParameters, WorkspaceDFastMatchState as DFastMatchState,
    WorkspaceFastMatchState as FastMatchState, WorkspaceFrameBlockState as FrameBlockState,
    WorkspaceGreedyMatchState as GreedyMatchState, WorkspaceLazyBlockStrategy as LazyBlockStrategy,
    WorkspaceLdmWorkspace as LdmWorkspace, WorkspaceOptBlockState as OptBlockState,
    WorkspaceOptFrameStrategy as OptFrameStrategy, WorkspaceParamSwitch as ParamSwitch,
    WorkspaceStrategy as Strategy,
};

/// Errors reported while creating or using an encoder workspace.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum EncoderWorkspaceError {
    Capacity(crate::workspace::CapacityError),
    /// A prior encode did not complete. Recreate the context before reuse.
    ContextPoisoned,
    /// Requested strategy would enter an owned-only entropy path.
    UnsupportedWorkspaceConfiguration,
    InsufficientWorkspace { required: usize, provided: usize },
    InputTooLarge { maximum: usize, provided: usize },
    OutputTooSmall { required: usize, provided: usize },
    SizeOverflow,
}

impl fmt::Display for EncoderWorkspaceError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Capacity(error) => write!(formatter, "encoder storage admission failed: {error:?}"),
            Self::UnsupportedWorkspaceConfiguration => formatter.write_str("encoder configuration is not supported by caller-backed storage"),
            Self::ContextPoisoned => formatter.write_str("encoder context must be recreated after an incomplete encode"),
            Self::InsufficientWorkspace { required, provided } => write!(
                formatter,
                "encoder workspace needs {required} bytes but only {provided} were provided"
            ),
            Self::InputTooLarge { maximum, provided } => write!(
                formatter,
                "encoder workspace accepts at most {maximum} input bytes but received {provided}"
            ),
            Self::OutputTooSmall { required, provided } => write!(
                formatter,
                "encoded output needs {required} bytes but only {provided} were provided"
            ),
            Self::SizeOverflow => formatter.write_str("encoder workspace size overflowed"),
        }
    }
}

// Boxing a strategy would add a heap allocation to caller-backed contexts;
// the enum deliberately keeps the selected state inline in the arena owner.
#[allow(clippy::large_enum_variant)]
enum MatchWorkspace {
    Uncompressed,
    Fast(FastMatchState),
    DFast(DFastMatchState),
    Greedy {
        state: GreedyMatchState,
        depth: LazyBlockStrategy,
    },
    Opt {
        state: OptBlockState,
        strategy: OptFrameStrategy,
    },
}

impl From<crate::workspace::CapacityError> for EncoderWorkspaceError {
    fn from(error: crate::workspace::CapacityError) -> Self { Self::Capacity(error) }
}

struct EncoderCore {
    // Set before any codec mutation and cleared only after complete success.
    // This also survives caught unwinds: taken scratch cannot silently become
    // an owned default buffer on a later sensitive operation.
    poisoned: bool,
    level: i32,
    maximum_input: usize,
    base_cctx: CctxParameters,
    frame: FrameBlockState,
    matcher: MatchWorkspace,
    ldm: Option<LdmWorkspace>,
}

impl EncoderCore {
    fn add_size(
        size: &mut ArenaSize,
        level: i32,
        maximum_input: usize,
    ) -> Result<CctxParameters, EncoderWorkspaceError> {
        let cctx = CctxParameters::for_level(level, maximum_input as u64, 0);
        FrameBlockState::add_workspace_size(size).map_err(map_arena)?;
        let block_size = crate::common::MAX_BLOCK_SIZE as usize;
        if level == 0 {
            return Ok(cctx);
        }
        match cctx.compression.strategy {
            Strategy::Fast => {
                FastMatchState::add_workspace_size(size, cctx.compression, block_size)
                    .map_err(map_arena)?;
            }
            Strategy::DFast => {
                DFastMatchState::add_workspace_size(size, cctx.compression, block_size)
                    .map_err(map_arena)?;
            }
            Strategy::Greedy | Strategy::Lazy | Strategy::Lazy2 | Strategy::BtLazy2 => {
                GreedyMatchState::add_workspace_size(size, cctx.compression, block_size)
                    .map_err(map_arena)?;
            }
            Strategy::BtOpt | Strategy::BtUltra | Strategy::BtUltra2 => {
                OptBlockState::add_workspace_size(size, cctx.compression, block_size)
                    .map_err(map_arena)?;
            }
        }
        if cctx.ldm.enable_ldm == ParamSwitch::Enable {
            LdmWorkspace::add_workspace_size(size, cctx.ldm, maximum_input).map_err(map_arena)?;
        }
        Ok(cctx)
    }

    fn new_in(
        arena: &mut Arena<'_>,
        level: i32,
        maximum_input: usize,
        cctx: CctxParameters,
    ) -> Result<Self, EncoderWorkspaceError> {
        let block_size = crate::common::MAX_BLOCK_SIZE as usize;
        let frame = FrameBlockState::new_in(arena, cctx.compression, cctx.max_block_size)
            .map_err(map_arena)?;
        let matcher = if level == 0 {
            MatchWorkspace::Uncompressed
        } else {
            match cctx.compression.strategy {
                Strategy::Fast => MatchWorkspace::Fast(
                    FastMatchState::new_in(arena, cctx.compression, block_size)
                        .map_err(map_arena)?,
                ),
                Strategy::DFast => MatchWorkspace::DFast(
                    DFastMatchState::new_in(arena, cctx.compression, block_size)
                        .map_err(map_arena)?,
                ),
                strategy @ (Strategy::Greedy
                | Strategy::Lazy
                | Strategy::Lazy2
                | Strategy::BtLazy2) => MatchWorkspace::Greedy {
                    state: GreedyMatchState::new_in(arena, cctx.compression, block_size)
                        .map_err(map_arena)?,
                    depth: match strategy {
                        Strategy::Greedy => LazyBlockStrategy::Greedy,
                        Strategy::Lazy => LazyBlockStrategy::Lazy,
                        Strategy::Lazy2 => LazyBlockStrategy::Lazy2,
                        Strategy::BtLazy2 => LazyBlockStrategy::BtLazy2,
                        _ => unreachable!(),
                    },
                },
                strategy @ (Strategy::BtOpt | Strategy::BtUltra | Strategy::BtUltra2) => {
                    MatchWorkspace::Opt {
                        state: OptBlockState::new_in(arena, cctx.compression, block_size)
                            .map_err(map_arena)?,
                        strategy: match strategy {
                            Strategy::BtOpt => OptFrameStrategy::BtOpt,
                            Strategy::BtUltra => OptFrameStrategy::BtUltra,
                            Strategy::BtUltra2 => OptFrameStrategy::BtUltra2,
                            _ => unreachable!(),
                        },
                    }
                }
            }
        };
        let ldm = if cctx.ldm.enable_ldm == ParamSwitch::Enable {
            Some(LdmWorkspace::new_in(arena, cctx.ldm, maximum_input).map_err(map_arena)?)
        } else {
            None
        };
        Ok(Self {
            poisoned: false,
            level,
            maximum_input,
            base_cctx: cctx,
            frame,
            matcher,
            ldm,
        })
    }

    fn encode_into<'output>(
        &mut self,
        input: &[u8],
        output: &'output mut [u8],
    ) -> Result<&'output [u8], EncoderWorkspaceError> {
        if self.poisoned { return Err(EncoderWorkspaceError::ContextPoisoned); }
        if input.len() > self.maximum_input {
            return Err(EncoderWorkspaceError::InputTooLarge {
                maximum: self.maximum_input,
                provided: input.len(),
            });
        }
        let required = checked_output_size(input.len())?;
        if output.len() < required {
            return Err(EncoderWorkspaceError::OutputTooSmall {
                required,
                provided: output.len(),
            });
        }
        self.poisoned = true;
        let cctx = CctxParameters::from_compression_parameters(
            self.level,
            self.base_cctx.compression,
            input.len() as u64,
        );
        // Target mode uses owned clone-based alternatives. Static encoding
        // currently never requests it; refuse a future accidental routing change.
        if cctx.target_c_block_size != 0
            || !FrameBlockState::supports_workspace_entropy(cctx.compression)
        {
            return Err(EncoderWorkspaceError::UnsupportedWorkspaceConfiguration);
        }
        // SAFETY: the guarded facade never owns caller storage, including on
        // early exit. Every growth is admitted before writes by the fixed-capacity
        // facade; caller storage cannot be reallocated even if sizing is wrong.
        let reusable = unsafe { ReusableVec::from_static_parts(output.as_mut_ptr(), output.len()) };
        let (mut output_vec, lease) = reusable.lease_vec();
        let ldm = &mut self.ldm;
        match &mut self.matcher {
            MatchWorkspace::Uncompressed => encode_uncompressed(input, &mut output_vec),
            MatchWorkspace::Fast(state) => {
                c_port::workspace_encode_fast(input, cctx, &mut self.frame, state, &mut output_vec)
            }
            MatchWorkspace::DFast(state) => {
                c_port::workspace_encode_dfast(input, cctx, &mut self.frame, state, &mut output_vec)
            }
            MatchWorkspace::Greedy { state, depth } => c_port::workspace_encode_greedy(
                input,
                cctx,
                *depth,
                &mut self.frame,
                state,
                &mut output_vec,
            ),
            MatchWorkspace::Opt { state, strategy } => c_port::workspace_encode_opt(
                input,
                cctx,
                *strategy,
                &mut self.frame,
                state,
                ldm.as_mut(),
                &mut output_vec,
            ),
        }?;
        let length = output_vec.len();
        let reusable = ReusableVec::recover_vec(output_vec, lease);
        drop(reusable);
        self.poisoned = false;
        Ok(&output[..length])
    }
}

fn encode_uncompressed(input: &[u8], output: &mut impl crate::workspace::Growable<u8>) -> Result<(), crate::workspace::CapacityError> {
    use crate::{
        blocks::block::BlockType,
        encoding::{block_header::BlockHeader, frame_header::FrameHeader},
    };
    output.clear();
    FrameHeader {
        frame_content_size: Some(input.len() as u64),
        single_segment: true,
        content_checksum: false,
        dictionary_id: None,
        window_size: None,
    }
    .serialize(output)?;
    if input.is_empty() {
        BlockHeader {
            last_block: true,
            block_type: BlockType::Raw,
            block_size: 0,
        }
        .serialize(output)?;
        return Ok(());
    }
    let block_size = crate::common::MAX_BLOCK_SIZE as usize;
    for (index, block) in input.chunks(block_size).enumerate() {
        BlockHeader {
            last_block: (index + 1) * block_size >= input.len(),
            block_type: BlockType::Raw,
            block_size: block.len() as u32,
        }
        .serialize(output)?;
        output.extend_from_slice(block)?;
    }
    Ok(())
}

/// An owned compression workspace reusable without further allocation.
pub struct EncoderWorkspace {
    core: EncoderCore,
    storage: Vec<MaybeUninit<u8>>,
}

impl EncoderWorkspace {
    pub fn new(
        level: CompressionLevel,
        maximum_input: usize,
    ) -> Result<Self, EncoderWorkspaceError> {
        let required = Self::required_size(level, maximum_input)?;
        let mut storage = Vec::with_capacity(required);
        storage.resize_with(required, MaybeUninit::uninit);
        let mut arena = Arena::new(&mut storage);
        let mut size = ArenaSize::new();
        let cctx = EncoderCore::add_size(&mut size, level.get(), maximum_input)?;
        let core = EncoderCore::new_in(&mut arena, level.get(), maximum_input, cctx)?;
        Ok(Self { core, storage })
    }

    pub fn required_size(
        level: CompressionLevel,
        maximum_input: usize,
    ) -> Result<usize, EncoderWorkspaceError> {
        let mut size = ArenaSize::new();
        EncoderCore::add_size(&mut size, level.get(), maximum_input)?;
        Ok(size.finish())
    }

    pub fn required_output_size(maximum_input: usize) -> Result<usize, EncoderWorkspaceError> {
        checked_output_size(maximum_input)
    }

    pub fn encode_into<'output>(
        &mut self,
        input: &[u8],
        output: &'output mut [u8],
    ) -> Result<&'output [u8], EncoderWorkspaceError> {
        self.core.encode_into(input, output)
    }

    pub fn workspace_bytes(&self) -> usize {
        self.storage.len()
    }

    pub const fn maximum_input(&self) -> usize {
        self.core.maximum_input
    }
}

#[cfg(feature = "std")]
impl std::error::Error for EncoderWorkspaceError {}

/// A compression context whose working state lives entirely in caller storage.
pub struct StaticEncoderWorkspace<'storage> {
    core: EncoderCore,
    workspace_bytes: usize,
    marker: PhantomData<&'storage mut [MaybeUninit<u8>]>,
}

impl<'storage> StaticEncoderWorkspace<'storage> {
    pub fn new(
        storage: &'storage mut [u8],
        level: CompressionLevel,
        maximum_input: usize,
    ) -> Result<Self, EncoderWorkspaceError> {
        Self::new_uninit(initialized_bytes_as_uninit(storage), level, maximum_input)
    }

    /// Constructs a workspace without initializing caller-provided bytes.
    pub fn new_uninit(
        storage: &'storage mut [MaybeUninit<u8>],
        level: CompressionLevel,
        maximum_input: usize,
    ) -> Result<Self, EncoderWorkspaceError> {
        let required = EncoderWorkspace::required_size(level, maximum_input)?;
        if storage.len() < required {
            return Err(EncoderWorkspaceError::InsufficientWorkspace {
                required,
                provided: storage.len(),
            });
        }
        let mut size = ArenaSize::new();
        let cctx = EncoderCore::add_size(&mut size, level.get(), maximum_input)?;
        let mut arena = Arena::new(storage);
        let core = EncoderCore::new_in(&mut arena, level.get(), maximum_input, cctx)?;
        Ok(Self {
            core,
            workspace_bytes: storage.len(),
            marker: PhantomData,
        })
    }

    pub fn encode_into<'output>(
        &mut self,
        input: &[u8],
        output: &'output mut [u8],
    ) -> Result<&'output [u8], EncoderWorkspaceError> {
        self.core.encode_into(input, output)
    }

    pub const fn maximum_input(&self) -> usize {
        self.core.maximum_input
    }

    pub const fn workspace_bytes(&self) -> usize {
        self.workspace_bytes
    }
}

fn initialized_bytes_as_uninit(storage: &mut [u8]) -> &mut [MaybeUninit<u8>] {
    // SAFETY: `MaybeUninit<u8>` has the same layout as `u8`, and treating
    // initialized bytes as possibly uninitialized only weakens the invariant.
    unsafe { core::slice::from_raw_parts_mut(storage.as_mut_ptr().cast(), storage.len()) }
}

fn map_arena(error: ArenaError) -> EncoderWorkspaceError {
    match error {
        ArenaError::Capacity(error) => EncoderWorkspaceError::Capacity(error),
        ArenaError::InsufficientStorage { required, provided } => {
            EncoderWorkspaceError::InsufficientWorkspace { required, provided }
        }
        ArenaError::CapacityExceeded { .. } | ArenaError::SizeOverflow => {
            EncoderWorkspaceError::SizeOverflow
        }
    }
}

fn checked_output_size(input_size: usize) -> Result<usize, EncoderWorkspaceError> {
    let bound = c_port::workspace_compress_bound(input_size);
    if bound == 0 {
        Err(EncoderWorkspaceError::SizeOverflow)
    } else {
        Ok(bound)
    }
}