structured-zstd 0.2.6

//! Gathers a block's literal runs into the literals buffer after matching.
//!
//! The matcher reports each sequence's literal length, and the runs sit in the
//! block's source in order: run `i` starts where run `i - 1` and its match
//! ended. Copying them in one pass after the match loop keeps the copy out of
//! that loop, and lets the pass run under the CPU tier chosen once for the
//! block, so no copy below the dispatch asks which kernel to use.

use crate::decoding::simd_copy;
use crate::encoding::fastpath::EncodeTier;
use crate::encoding::workspace::RegionVec;

use super::RawSequence;

/// Run length at or above which a run is handed to `memcpy`: the copy is
/// bandwidth-bound there and ERMS `rep movsb` beats a vector loop. Below it the
/// inline kernels win on the call and startup cost they avoid.
const LITERAL_INLINE_COPY_MAX: usize = 2048;

/// Append the literal runs of `sequences`, then `tail` trailing literal bytes,
/// from `block` to `dst`, with the copies of the tier `K`.
///
/// # Panics
/// If the sequences and the tail reach past `block`, which a matcher that
/// reports its sequences correctly never produces.
pub(super) fn gather_literals<K: EncodeTier>(
    block: &[u8],
    sequences: &[RawSequence],
    tail: usize,
    dst: &mut RegionVec<u8>,
) {
    // The literals of a block are at most the block, and the buffer is sized
    // for the largest block, so this holds on every internal path. It is still
    // checked, once, because the loop below writes through a raw pointer.
    if dst.capacity() - dst.len() < block.len() {
        crate::encoding::workspace::capacity_exhausted();
    }
    crate::encoding::fastpath::tier_match!(K {
        #[cfg(all(
            any(target_arch = "x86", target_arch = "x86_64"),
            feature = "kernel-avx2"
        ))]
        // SAFETY: the tier is only selected after confirming AVX2 (and BMI2).
        crate::encoding::fastpath::FastpathKernel::Avx2Bmi2 => unsafe {
            gather_avx2(block, sequences, tail, dst)
        },
        #[cfg(all(
            any(target_arch = "x86", target_arch = "x86_64"),
            feature = "kernel-sse"
        ))]
        // SAFETY: both tiers are only selected after confirming SSE2.
        crate::encoding::fastpath::FastpathKernel::Sse2
        | crate::encoding::fastpath::FastpathKernel::Sse42 => unsafe {
            gather_sse2(block, sequences, tail, dst)
        },
        #[cfg(all(
            target_arch = "aarch64",
            target_endian = "little",
            feature = "kernel-neon"
        ))]
        crate::encoding::fastpath::FastpathKernel::Neon => gather_neon(block, sequences, tail, dst),
        #[cfg(all(
            target_arch = "wasm32",
            target_feature = "simd128",
            feature = "kernel-simd128"
        ))]
        crate::encoding::fastpath::FastpathKernel::Simd128 => {
            gather_simd128(block, sequences, tail, dst)
        },
        _ => gather_scalar(block, sequences, tail, dst),
    })
}

/// Bytes of the block that must remain past a run for it to be copied in whole
/// chunks: the widest chunk rounds a run up by at most 31 bytes, read from the
/// block and written into the literals buffer, which is at least as long.
const WILDCOPY_SLACK: usize = 32;

/// One pass over the runs. A run with [`WILDCOPY_SLACK`] bytes of block after
/// it is copied in whole `$chunk`-byte chunks by `$wild`, overshooting into
/// bytes the next run overwrites (upstream zstd `ZSTD_storeSeq` +
/// `ZSTD_wildcopy`); the runs near the block's end take the exact ladder, with
/// `$medium` for `33..2048` bytes. A macro so each tier's loop is compiled
/// inside its own `#[target_feature]` function and its kernels inline into it.
macro_rules! gather_body {
    ($block:expr, $sequences:expr, $tail:expr, $dst:expr, $medium:path, $wild:path, $chunk:expr) => {{
        let block: &[u8] = $block;
        let dst: &mut RegionVec<u8> = $dst;
        let src = block.as_ptr();
        let src_len = block.len();
        let start = dst.len();
        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;
                // 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` 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);
                    // 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.
                        // The first chunk is a constant-length copy, a single
                        // load and store, so an empty run costs less to copy
                        // than to test for.
                        if len < LITERAL_INLINE_COPY_MAX {
                            $wild(s, d, $chunk);
                            if len > $chunk {
                                $wild(
                                    s.add($chunk),
                                    d.add($chunk),
                                    (len - $chunk).next_multiple_of($chunk),
                                );
                            }
                        } else {
                            core::ptr::copy_nonoverlapping(s, d, len);
                        }
                    } else {
                        // `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);
                            }
                        } else if len < LITERAL_INLINE_COPY_MAX {
                            $medium(s, d, len);
                        } else {
                            core::ptr::copy_nonoverlapping(s, d, len);
                        }
                    }
                }
                written += len;
            }};
        }
        for seq in $sequences {
            let ll = seq.ll as usize;
            copy_run!(ll);
            pos += ll;
            // The next run is read at `pos`, so the match is held to the block
            // as the literals are; `pos <= src_len` holds on every iteration.
            let ml = seq.ml as usize;
            assert!(
                ml <= src_len - pos,
                "sequences reach past the block they describe"
            );
            pos += ml;
        }
        copy_run!($tail);
        // Once per block, in every build: a custom matcher that stops short
        // would leave the block's end out of the frame.
        assert_eq!(pos + $tail, src_len, "sequences must cover the block");
        // SAFETY: `[start, start + written)` was written above, within capacity.
        unsafe { dst.set_len(start + written) };
    }};
}

#[cfg(all(
    any(target_arch = "x86", target_arch = "x86_64"),
    feature = "kernel-avx2"
))]
#[target_feature(enable = "avx2")]
unsafe fn gather_avx2(
    block: &[u8],
    sequences: &[RawSequence],
    tail: usize,
    dst: &mut RegionVec<u8>,
) {
    gather_body!(
        block,
        sequences,
        tail,
        dst,
        simd_copy::copy_exact_avx2,
        simd_copy::copy_avx2,
        32
    )
}

#[cfg(all(
    any(target_arch = "x86", target_arch = "x86_64"),
    feature = "kernel-sse"
))]
#[target_feature(enable = "sse2")]
unsafe fn gather_sse2(
    block: &[u8],
    sequences: &[RawSequence],
    tail: usize,
    dst: &mut RegionVec<u8>,
) {
    gather_body!(
        block,
        sequences,
        tail,
        dst,
        simd_copy::copy_exact_sse2,
        simd_copy::copy_sse2,
        16
    )
}

#[cfg(all(
    target_arch = "aarch64",
    target_endian = "little",
    feature = "kernel-neon"
))]
fn gather_neon(block: &[u8], sequences: &[RawSequence], tail: usize, dst: &mut RegionVec<u8>) {
    #[cfg(target_feature = "neon")]
    gather_body!(
        block,
        sequences,
        tail,
        dst,
        simd_copy::copy_exact_neon,
        simd_copy::copy_neon,
        16
    );
    #[cfg(not(target_feature = "neon"))]
    gather_body!(
        block,
        sequences,
        tail,
        dst,
        simd_copy::copy_exact_u64,
        simd_copy::copy_scalar,
        core::mem::size_of::<usize>()
    );
}

/// wasm SIMD is a compile-time feature, so the tier needs no
/// `#[target_feature]` wrapper; it takes its own loop all the same, with the
/// 16-byte `v128` kernels.
#[cfg(all(
    target_arch = "wasm32",
    target_feature = "simd128",
    feature = "kernel-simd128"
))]
fn gather_simd128(block: &[u8], sequences: &[RawSequence], tail: usize, dst: &mut RegionVec<u8>) {
    gather_body!(
        block,
        sequences,
        tail,
        dst,
        simd_copy::copy_exact_simd128,
        simd_copy::copy_simd128,
        16
    )
}

fn gather_scalar(block: &[u8], sequences: &[RawSequence], tail: usize, dst: &mut RegionVec<u8>) {
    gather_body!(
        block,
        sequences,
        tail,
        dst,
        simd_copy::copy_exact_u64,
        simd_copy::copy_scalar,
        core::mem::size_of::<usize>()
    )
}