lzip_parallel/

deflate_scan.rs

//! DEFLATE full-flush boundary scanner.
//!
//! A Z_FULL_FLUSH point resets the LZ77 back-reference window to empty,
//! making all bytes that follow independently decompressible.  In the
//! compressed byte stream this always produces the 5-byte sequence:
//!
//!   00          BFINAL=0, BTYPE=00 (stored), 5 padding zero-bits
//!   00 00       LEN  = 0x0000 (zero-length stored block)
//!   FF FF       NLEN = 0xFFFF (one's complement)
//!
//! We scan for the 4-byte LEN+NLEN pattern `00 00 FF FF`.  The next
//! independently-decompressible segment starts at `match_pos + 4`.
//!
//! Mirrors lbzip2-rs block_scan.rs: parallel candidate-offset forward-scan,
//! one task per physical core, no look-behind verification (Q1 resolved).

use rayon::prelude::*;

/// A full-flush boundary.
///
/// `start` is the byte offset where the next independently-decompressible
/// segment begins (the byte immediately after the `FF FF` NLEN field).
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FlushBoundary {
    pub start: usize,
}

/// Scan forward from `from` for the next `00 00 FF FF` pattern.
///
/// Returns a `FlushBoundary` whose `start` is the first byte of the
/// segment that follows the flush marker, or `None` if not found.
pub fn find_next_flush(buf: &[u8], from: usize) -> Option<FlushBoundary> {
    if buf.len() < 4 || from + 4 > buf.len() {
        return None;
    }

    #[cfg(target_arch = "x86_64")]
    {
        if std::arch::is_x86_feature_detected!("avx512bw") {
            return unsafe { find_next_flush_avx512(buf, from) };
        }
        if std::arch::is_x86_feature_detected!("avx2") {
            return unsafe { find_next_flush_avx2(buf, from) };
        }
    }

    find_next_flush_scalar(buf, from)
}

/// Scalar fallback for flush scan.
fn find_next_flush_scalar(buf: &[u8], from: usize) -> Option<FlushBoundary> {
    if buf.len() < 4 {
        return None;
    }
    let end = buf.len() - 3;
    let mut i = from;
    while i < end {
        // Fast skip: the pattern requires buf[i+2]==0xFF.
        if buf[i + 2] != 0xFF {
            i += 1;
            continue;
        }
        if buf[i] == 0x00 && buf[i + 1] == 0x00 && buf[i + 3] == 0xFF {
            return Some(FlushBoundary { start: i + 4 });
        }
        i += 1;
    }
    None
}

/// AVX-512BW flush scan: searches 64 bytes at a time.
///
/// Same algorithm as AVX2 but with 64-byte vectors (zmm registers).
/// On Zen 4+ and Intel Skylake-X+, processes 2× the data per iteration.
///
/// # Safety
/// Requires AVX-512BW. Caller ensures `buf.len() >= 4`.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512bw")]
unsafe fn find_next_flush_avx512(buf: &[u8], from: usize) -> Option<FlushBoundary> {
    use core::arch::x86_64::*;

    unsafe {
        let zero_vec = _mm512_setzero_si512();
        let ff_vec = _mm512_set1_epi8(-1i8);

        let ptr = buf.as_ptr();
        let len = buf.len();
        let mut i = from;

        // Process 64 bytes at a time; need 3 extra for pattern overlap
        while i + 67 <= len {
            let chunk = _mm512_loadu_si512(ptr.add(i) as *const __m512i);

            // Compare each byte to 0x00 and 0xFF — returns 64-bit masks
            let zero_mask = _mm512_cmpeq_epi8_mask(chunk, zero_vec);
            let ff_mask = _mm512_cmpeq_epi8_mask(chunk, ff_vec);

            // Pattern 00 00 FF FF: shifted AND of masks
            let candidates = (zero_mask & (zero_mask >> 1)) & ((ff_mask >> 2) & (ff_mask >> 3));

            if candidates != 0 {
                let bit_pos = candidates.trailing_zeros() as usize;
                let pos = i + bit_pos;
                if pos + 4 <= len
                    && buf[pos] == 0x00
                    && buf[pos + 1] == 0x00
                    && buf[pos + 2] == 0xFF
                    && buf[pos + 3] == 0xFF
                {
                    return Some(FlushBoundary { start: pos + 4 });
                }
                i += bit_pos + 1;
                continue;
            }

            i += 64;
        }

        // Fall back to AVX2/scalar for the tail
        find_next_flush_scalar(buf, i)
    }
}

/// AVX2 flush scan: searches 32 bytes at a time using SIMD comparison.
///
/// Strategy: compare all bytes against 0x00, then against 0xFF.
/// A valid pattern `00 00 FF FF` means: zero[i] && zero[i+1] && ff[i+2] && ff[i+3].
/// We use shifted AND of the bitmasks to find candidates.
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn find_next_flush_avx2(buf: &[u8], from: usize) -> Option<FlushBoundary> {
    use core::arch::x86_64::*;

    unsafe {
        let zero_vec = _mm256_setzero_si256();
        let ff_vec = _mm256_set1_epi8(-1); // 0xFF as i8

        let ptr = buf.as_ptr();
        let len = buf.len();

        let mut i = from;

        // Process 32 bytes at a time
        // We need to check positions i..i+35 (32 bytes + 3 for pattern overlap)
        while i + 35 <= len {
            let chunk = _mm256_loadu_si256(ptr.add(i) as *const __m256i);

            // Compare each byte to 0x00 and 0xFF
            let eq_zero = _mm256_cmpeq_epi8(chunk, zero_vec);
            let eq_ff = _mm256_cmpeq_epi8(chunk, ff_vec);

            let zero_mask = _mm256_movemask_epi8(eq_zero) as u32;
            let ff_mask = _mm256_movemask_epi8(eq_ff) as u32;

            // Pattern 00 00 FF FF at position j means:
            //   zero_mask bit j set, zero_mask bit j+1 set,
            //   ff_mask bit j+2 set, ff_mask bit j+3 set.
            let candidates = (zero_mask & (zero_mask >> 1)) & ((ff_mask >> 2) & (ff_mask >> 3));

            if candidates != 0 {
                let bit_pos = candidates.trailing_zeros() as usize;
                let pos = i + bit_pos;
                if pos + 4 <= len
                    && buf[pos] == 0x00
                    && buf[pos + 1] == 0x00
                    && buf[pos + 2] == 0xFF
                    && buf[pos + 3] == 0xFF
                {
                    return Some(FlushBoundary { start: pos + 4 });
                }
                i += bit_pos + 1;
                continue;
            }

            i += 32;
        }

        // Scalar tail for remaining bytes
        find_next_flush_scalar(buf, i)
    }
}

/// Find all full-flush boundaries in `buf`.
pub fn find_all_flushes(buf: &[u8]) -> Vec<FlushBoundary> {
    let mut result = Vec::new();
    let mut pos = 0;
    while let Some(b) = find_next_flush(buf, pos) {
        result.push(b);
        pos = b.start;
    }
    result
}

/// Parallel split-point search: N−1 rayon tasks each find the nearest flush
/// boundary at or after their nominal byte offset.
///
/// `n_splits` is typically the number of physical cores.
/// Returns up to `n_splits − 1` boundaries, sorted and deduplicated.
pub fn split_boundaries_parallel(buf: &[u8], n_splits: usize) -> Vec<FlushBoundary> {
    if n_splits <= 1 || buf.len() < 8 {
        return Vec::new();
    }

    let mut found: Vec<Option<FlushBoundary>> = crate::thread_pool().install(|| {
        (1..n_splits)
            .into_par_iter()
            .map(|i| {
                let nominal = buf.len() * i / n_splits;
                find_next_flush(buf, nominal)
            })
            .collect()
    });

    let mut result: Vec<FlushBoundary> = found.drain(..).flatten().collect();
    result.sort_by_key(|b| b.start);
    result.dedup_by_key(|b| b.start);
    result
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_flush_marker() -> Vec<u8> {
        vec![0x00, 0x00, 0x00, 0xFF, 0xFF]
    }

    #[test]
    fn find_flush_at_start() {
        let mut buf = make_flush_marker();
        buf.extend_from_slice(&[0xAA, 0xBB]);
        // Pattern 00 00 FF FF starts at offset 1 (the LEN bytes)
        let b = find_next_flush(&buf, 0).unwrap();
        assert_eq!(b.start, 5); // after the 5-byte sequence
    }

    #[test]
    fn find_flush_mid_buffer() {
        let mut buf = vec![0xDE, 0xAD, 0xBE, 0xEF];
        buf.extend_from_slice(&make_flush_marker());
        buf.extend_from_slice(&[0x01, 0x02]);
        let b = find_next_flush(&buf, 0).unwrap();
        assert_eq!(b.start, 9);
    }

    #[test]
    fn no_flush_in_random_data() {
        // High-byte data is unlikely to match
        let buf: Vec<u8> = (0..64).map(|i| (i * 37 + 1) as u8).collect();
        // Just check it doesn't panic
        let _ = find_next_flush(&buf, 0);
    }

    #[test]
    fn find_all_flushes_two_markers() {
        let mut buf = Vec::new();
        buf.extend_from_slice(&make_flush_marker()); // boundary at 5
        buf.extend_from_slice(&[0xAA; 10]);
        buf.extend_from_slice(&make_flush_marker()); // boundary at 20
        let all = find_all_flushes(&buf);
        assert_eq!(all.len(), 2);
        assert_eq!(all[0].start, 5);
        assert_eq!(all[1].start, 20);
    }
}