structured-zstd 0.2.6

//! Encoder fastpath: hot encode functions duplicated per CPU feature set so the
//! whole hot loop stays inside one `#[target_feature]` umbrella and SIMD/BMI2
//! intrinsics inline natively (no ABI barrier).
//!
//! All kernel functions are `unsafe fn`; the explicit inner `unsafe { }` blocks
//! around intrinsic calls are kept for safety documentation (this matches the
//! Rust 2024 recommended style enforced by `unsafe_op_in_unsafe_fn`). The
//! `unused_unsafe` lint sees them as redundant inside an `unsafe fn` body, so
//! we silence it at the module level rather than removing the documentation.
#![allow(unused_unsafe)]
//!
//! # Background
//!
//! In Rust, `#[target_feature(enable = "...")]` creates an ABI boundary: a
//! caller without the same feature set must call the function non-inline. In
//! C, the equivalent intrinsics inline via macros without restriction. That ABI
//! barrier is the dominant structural reason our encoder cannot match the
//! C zstd upstream zstd on per-block latency — every hot-path SIMD call becomes a
//! function call (~100 cycles overhead per BT walk iter, ~32-512 iters per
//! position, thousands of positions per block).
//!
//! # Strategy
//!
//! Each architecture-specific submodule (`neon`, `avx2_bmi2`, `sse2`,
//! `scalar`) holds a duplicate of the hot encode path, with every function in
//! the chain marked with the same `#[target_feature]`. Inside the module
//! everything inlines freely. The single ABI boundary is the dispatcher entry
//! point in this `mod.rs`, called once per encoder invocation.
//!
//! # Variant matrix
//!
//! - `scalar`: portable baseline, no SIMD assumptions. Used on unsupported
//!   targets and as fallback.
//! - `neon` (aarch64 only): NEON is part of the AArch64 baseline ISA but Rust
//!   still flags intrinsics like `vld1q_u8` with `#[target_feature(enable =
//!   "neon")]`, so we still need the umbrella attribute to let them inline.
//! - `sse2` (x86/x86_64): 128-bit SSE2 vector ops, the x86_64 baseline.
//!   Shared by both x86 tiers — they differ only in the optimal parser's
//!   price set, where the SSE4.2 tier can use `_mm_min_epu32`.
//! - `avx2_bmi2` (x86_64): adds AVX2 (32-byte vectors) and BMI2 (`pext`,
//!   `pdep`, `bzhi`) — common on Haswell+ (2013+).
//!
//! # Dispatcher
//!
//! [`select_kernel`] picks the best supported variant once per process via a
//! `OnceLock`. Encoder entry points call through the cached function pointer.
//! The single indirect call is amortized over the entire compression call,
//! and once inside the variant module the call graph is straight-line inlined.
//!
//! # Roadmap inside this module
//!
//! Week 1 (this commit): module scaffold + dispatcher skeleton.
//! Week 2a: match-length / common-prefix-len + `count_match_from_indices`.
//! Week 3a: BT walk (`bt_insert_range`, the collect-matches body) + HC
//!   chain walk.
//! Week 3b: optimal parser DP (`build_optimal_plan_impl` + price helpers).
//! Week 4: entropy encoders (FSE `encode_interleaved`, Huff0 `encode_stream`).
//! Week 5-6: bench vs `perf/pre-intrinsics-refactor-baseline` tag, profile,
//!   finalize.
//!
//! Refactor history and working rules for the multi-week PR #110 effort are
//! captured in the corresponding pull-request description.

pub(crate) mod scalar;

// Every SIMD tier the detection below can pick asks the ceiling first.
#[cfg(any(
    all(
        any(target_arch = "x86", target_arch = "x86_64"),
        any(feature = "kernel-sse", feature = "kernel-avx2")
    ),
    all(
        target_arch = "aarch64",
        target_endian = "little",
        feature = "kernel-neon"
    )
))]
use crate::cpu_kernel::{CpuLevel, cpu_allows};

#[cfg(all(
    target_arch = "aarch64",
    target_endian = "little",
    feature = "kernel-neon"
))]
pub(crate) mod neon;

#[cfg(all(
    any(target_arch = "x86", target_arch = "x86_64"),
    feature = "kernel-sse"
))]
pub(crate) mod sse2;

#[cfg(all(
    any(target_arch = "x86", target_arch = "x86_64"),
    feature = "kernel-avx2"
))]
pub(crate) mod avx2_bmi2;

#[cfg(all(
    target_arch = "wasm32",
    target_feature = "simd128",
    feature = "kernel-simd128"
))]
pub(crate) mod simd128;

/// Runtime-selected variant tag. Picked once per process by [`select_kernel`].
///
/// Each variant corresponds to one of the submodules above and dictates which
/// implementation of the hot encoder path the dispatcher will call into.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub(crate) enum FastpathKernel {
    Scalar,
    #[cfg(all(
        target_arch = "aarch64",
        target_endian = "little",
        feature = "kernel-neon"
    ))]
    Neon,
    /// SSE2-only x86 tier: the 128-bit prefix-compare kernel plus the SSE2
    /// price-set path. Selected on x86 CPUs without SSE4.2, which would
    /// otherwise fall all the way back to [`FastpathKernel::Scalar`].
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-sse"
    ))]
    Sse2,
    /// SSE4.2 x86 tier: same 128-bit kernel as [`FastpathKernel::Sse2`], but
    /// the price set uses the SSE4.1 `min_epu32` compare.
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-sse"
    ))]
    Sse42,
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-avx2"
    ))]
    Avx2Bmi2,
    #[cfg(all(
        target_arch = "wasm32",
        target_feature = "simd128",
        feature = "kernel-simd128"
    ))]
    Simd128,
}

impl FastpathKernel {
    /// Stable lowercase name of the tier, for diagnostics.
    pub(crate) fn name(self) -> &'static str {
        match self {
            FastpathKernel::Scalar => "scalar",
            #[cfg(all(
                target_arch = "aarch64",
                target_endian = "little",
                feature = "kernel-neon"
            ))]
            FastpathKernel::Neon => "neon",
            #[cfg(all(
                any(target_arch = "x86", target_arch = "x86_64"),
                feature = "kernel-sse"
            ))]
            FastpathKernel::Sse2 => "sse2",
            #[cfg(all(
                any(target_arch = "x86", target_arch = "x86_64"),
                feature = "kernel-sse"
            ))]
            FastpathKernel::Sse42 => "sse4.2",
            #[cfg(all(
                any(target_arch = "x86", target_arch = "x86_64"),
                feature = "kernel-avx2"
            ))]
            FastpathKernel::Avx2Bmi2 => "avx2",
            #[cfg(all(
                target_arch = "wasm32",
                target_feature = "simd128",
                feature = "kernel-simd128"
            ))]
            FastpathKernel::Simd128 => "simd128",
        }
    }
}

/// Select the best supported variant for the running CPU. Cached after first
/// call; intended to be invoked once at the entry point of each encoder call
/// so the rest of the call graph can keep working with the resolved kernel
/// value as a const-foldable input.
#[inline]
pub(crate) fn select_kernel() -> FastpathKernel {
    #[cfg(feature = "std")]
    {
        use std::sync::OnceLock;
        static CACHE: OnceLock<FastpathKernel> = OnceLock::new();
        *CACHE.get_or_init(detect_kernel_uncached)
    }
    #[cfg(not(feature = "std"))]
    {
        detect_kernel_uncached()
    }
}

#[inline]
// On wasm32+simd128 the tier is resolved unconditionally to `Simd128` (no
// runtime CPUID), so the trailing `Scalar` fallback is statically unreachable
// there; it stays the reachable fallback on every other target.
#[cfg_attr(
    all(
        target_arch = "wasm32",
        target_feature = "simd128",
        feature = "kernel-simd128"
    ),
    allow(unreachable_code)
)]
fn detect_kernel_uncached() -> FastpathKernel {
    // Each probe covers everything its tier executes, which is more than the
    // vector ops in this module's own submodules: the `Sse42` tier's
    // optimal-parser body calls `priceset::priceset_range_nonabort_sse41`
    // (`_mm_min_epu32`, SSE4.1). A CPU with SSE2 but no SSE4.2 therefore takes
    // the `Sse2` tier, which pairs the same 128-bit prefix-compare kernel with
    // the SSE2 price set, rather than dropping to the scalar kernel.
    #[cfg(all(
        feature = "std",
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-avx2"
    ))]
    {
        if cpu_allows(CpuLevel::Avx2)
            && std::is_x86_feature_detected!("avx2")
            && std::is_x86_feature_detected!("bmi2")
        {
            return FastpathKernel::Avx2Bmi2;
        }
    }
    #[cfg(all(
        feature = "std",
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-sse"
    ))]
    {
        if cpu_allows(CpuLevel::Sse42) && std::is_x86_feature_detected!("sse4.2") {
            return FastpathKernel::Sse42;
        }
        if cpu_allows(CpuLevel::Sse2) && std::is_x86_feature_detected!("sse2") {
            return FastpathKernel::Sse2;
        }
    }
    #[cfg(all(
        feature = "std",
        target_arch = "aarch64",
        target_endian = "little",
        feature = "kernel-neon"
    ))]
    {
        if cpu_allows(CpuLevel::Neon) && std::arch::is_aarch64_feature_detected!("neon") {
            return FastpathKernel::Neon;
        }
    }

    #[cfg(all(
        not(feature = "std"),
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-avx2"
    ))]
    {
        if cpu_allows(CpuLevel::Avx2)
            && cfg!(target_feature = "avx2")
            && cfg!(target_feature = "bmi2")
        {
            return FastpathKernel::Avx2Bmi2;
        }
    }
    #[cfg(all(
        not(feature = "std"),
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-sse"
    ))]
    {
        if cpu_allows(CpuLevel::Sse42) && cfg!(target_feature = "sse4.2") {
            return FastpathKernel::Sse42;
        }
        if cpu_allows(CpuLevel::Sse2) && cfg!(target_feature = "sse2") {
            return FastpathKernel::Sse2;
        }
    }
    #[cfg(all(
        not(feature = "std"),
        target_arch = "aarch64",
        target_endian = "little",
        feature = "kernel-neon"
    ))]
    {
        if cpu_allows(CpuLevel::Neon) && cfg!(target_feature = "neon") {
            return FastpathKernel::Neon;
        }
    }

    // wasm SIMD is a compile-time feature (no runtime detection), so the
    // `+simd128` payload selects the SIMD kernel and the scalar payload never
    // compiles the variant.
    #[cfg(all(
        target_arch = "wasm32",
        target_feature = "simd128",
        feature = "kernel-simd128"
    ))]
    {
        return FastpathKernel::Simd128;
    }

    FastpathKernel::Scalar
}

/// A [`FastpathKernel`] tier as a type. The encoder chooses its tier once, in
/// [`with_encode_tier`], and everything below that choice is compiled per tier:
/// a leaf that runs a different kernel per tier picks it with [`tier_match`],
/// which keeps the one kernel at monomorphisation, so no tier is asked at run
/// time below the choice.
pub(crate) trait EncodeTier {
    /// The tier this type stands for.
    const TIER: FastpathKernel;

    /// Length of the common prefix of `lhs` and `rhs`, at most `max`, with this
    /// tier's compare.
    ///
    /// # Safety
    /// `lhs` / `rhs` must each point to at least `max` initialized bytes, and
    /// the CPU must support the tier, which holds for the tier the encoder
    /// chose.
    unsafe fn common_prefix_len_ptr(lhs: *const u8, rhs: *const u8, max: usize) -> usize;
}

/// One tier type: the struct, its [`FastpathKernel`] and its compare kernel,
/// under the cfg that builds the tier.
macro_rules! encode_tier {
    ($(#[$cfg:meta])* $name:ident, $tier:ident, $prefix:path) => {
        #[doc = concat!("[`FastpathKernel::", stringify!($tier), "`] as a type.")]
        $(#[$cfg])*
        pub(crate) struct $name;
        $(#[$cfg])*
        impl EncodeTier for $name {
            const TIER: FastpathKernel = FastpathKernel::$tier;

            #[inline(always)]
            unsafe fn common_prefix_len_ptr(lhs: *const u8, rhs: *const u8, max: usize) -> usize {
                // SAFETY: forwarded from the caller.
                unsafe { $prefix(lhs, rhs, max) }
            }
        }
    };
}

encode_tier!(ScalarTier, Scalar, scalar::common_prefix_len_ptr);
encode_tier!(
    #[cfg(all(
        target_arch = "aarch64",
        target_endian = "little",
        feature = "kernel-neon"
    ))]
    NeonTier,
    Neon,
    neon::common_prefix_len_ptr
);
// Both SSE tiers compare with the same SSE2 kernel; only the optimal parser's
// price set differs between them.
encode_tier!(
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-sse"
    ))]
    Sse2Tier,
    Sse2,
    sse2::common_prefix_len_ptr
);
encode_tier!(
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-sse"
    ))]
    Sse42Tier,
    Sse42,
    sse2::common_prefix_len_ptr
);
encode_tier!(
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-avx2"
    ))]
    Avx2Bmi2Tier,
    Avx2Bmi2,
    avx2_bmi2::common_prefix_len_ptr
);
encode_tier!(
    #[cfg(all(
        target_arch = "wasm32",
        target_feature = "simd128",
        feature = "kernel-simd128"
    ))]
    Simd128Tier,
    Simd128,
    simd128::common_prefix_len_ptr
);

/// The encoder's one choice of tier: evaluate `$body` with `$k` naming the
/// [`EncodeTier`] type of the tier `$tier` holds.
macro_rules! with_encode_tier {
    ($tier:expr, $k:ident => $body:expr) => {{
        use $crate::encoding::fastpath as tiers;
        match $tier {
            tiers::FastpathKernel::Scalar => {
                type $k = tiers::ScalarTier;
                $body
            }
            #[cfg(all(
                target_arch = "aarch64",
                target_endian = "little",
                feature = "kernel-neon"
            ))]
            tiers::FastpathKernel::Neon => {
                type $k = tiers::NeonTier;
                $body
            }
            #[cfg(all(
                any(target_arch = "x86", target_arch = "x86_64"),
                feature = "kernel-sse"
            ))]
            tiers::FastpathKernel::Sse2 => {
                type $k = tiers::Sse2Tier;
                $body
            }
            #[cfg(all(
                any(target_arch = "x86", target_arch = "x86_64"),
                feature = "kernel-sse"
            ))]
            tiers::FastpathKernel::Sse42 => {
                type $k = tiers::Sse42Tier;
                $body
            }
            #[cfg(all(
                any(target_arch = "x86", target_arch = "x86_64"),
                feature = "kernel-avx2"
            ))]
            tiers::FastpathKernel::Avx2Bmi2 => {
                type $k = tiers::Avx2Bmi2Tier;
                $body
            }
            #[cfg(all(
                target_arch = "wasm32",
                target_feature = "simd128",
                feature = "kernel-simd128"
            ))]
            tiers::FastpathKernel::Simd128 => {
                type $k = tiers::Simd128Tier;
                $body
            }
        }
    }};
}
pub(crate) use with_encode_tier;

/// A `match` on `$k`'s [`EncodeTier::TIER`] that instantiates only the arm of
/// that tier.
///
/// A plain `match` on the constant folds in an optimised build, but an
/// unoptimised one still instantiates every arm's callee for every tier, so
/// each tier's monomorph would carry the kernels of all the others. Each arm
/// here is an `if` on an inline `const`, whose untaken branch the compiler
/// drops before code generation. The arms are checked for exhaustiveness as the
/// `match` they stand for.
macro_rules! tier_match {
    ($k:ty { $( $(#[$attr:meta])* $pat:pat => $arm:expr ),+ $(,)? }) => {{
        const {
            match <$k as $crate::encoding::fastpath::EncodeTier>::TIER {
                $( $(#[$attr])* $pat => (), )+
            }
        };
        'tier: {
            $(
                $(#[$attr])*
                if const {
                    matches!(<$k as $crate::encoding::fastpath::EncodeTier>::TIER, $pat)
                } {
                    break 'tier ($arm);
                }
            )+
            unreachable!()
        }
    }};
}
pub(crate) use tier_match;

/// [`with_encode_tier!`] on this CPU's tier, the one the encoder picks: unit
/// tests drive a backend directly through the kernels it would ship with.
#[cfg(test)]
macro_rules! with_host_tier {
    ($k:ident => $body:expr) => {
        $crate::encoding::fastpath::with_encode_tier!(
            $crate::encoding::fastpath::select_kernel(),
            $k => $body
        )
    };
}
#[cfg(test)]
pub(crate) use with_host_tier;

/// Every tier this CPU may execute, scalar first: each SIMD arm assumes its
/// target feature is present, so a test forcing a tier takes it from here. A
/// test comparing tiers checks the scalar fallback against every SIMD kernel
/// the machine has, including the ones the detection leaves cold.
#[cfg(test)]
pub(crate) fn runnable_tiers() -> alloc::vec::Vec<FastpathKernel> {
    #[allow(unused_mut)]
    let mut tiers = alloc::vec![FastpathKernel::Scalar];
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-sse"
    ))]
    {
        if std::is_x86_feature_detected!("sse2") {
            tiers.push(FastpathKernel::Sse2);
        }
        if std::is_x86_feature_detected!("sse4.2") {
            tiers.push(FastpathKernel::Sse42);
        }
    }
    #[cfg(all(
        any(target_arch = "x86", target_arch = "x86_64"),
        feature = "kernel-avx2"
    ))]
    if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("bmi2") {
        tiers.push(FastpathKernel::Avx2Bmi2);
    }
    #[cfg(all(
        target_arch = "aarch64",
        target_endian = "little",
        feature = "kernel-neon"
    ))]
    tiers.push(FastpathKernel::Neon);
    #[cfg(all(
        target_arch = "wasm32",
        target_feature = "simd128",
        feature = "kernel-simd128"
    ))]
    tiers.push(FastpathKernel::Simd128);
    tiers
}

/// `run` under every [`runnable_tiers`] tier, asserting they all agree, and
/// the answer they agree on.
#[cfg(test)]
pub(crate) fn same_under_every_tier<R: PartialEq + core::fmt::Debug>(
    run: impl Fn(FastpathKernel) -> R,
) -> R {
    let tiers = runnable_tiers();
    let answer = run(tiers[0]);
    for &tier in &tiers[1..] {
        assert_eq!(
            run(tier),
            answer,
            "tier {tier:?} disagrees with {:?}",
            tiers[0]
        );
    }
    answer
}

/// A test's sequence closure as the one sink type tests hand a backend, so
/// each search is compiled once per tier for all of them rather than once per
/// closure and tier.
#[cfg(test)]
pub(crate) fn dyn_sink(
    sink: &mut impl FnMut(crate::encoding::Sequence),
) -> &mut dyn FnMut(crate::encoding::Sequence) {
    sink
}

#[cfg(test)]
mod tests;