writerslogic_warm/core/

intersection.rs

// Copyright 2024-2026 WritersLogic Contributors
// SPDX-License-Identifier: Apache-2.0

//! Sparse intersection primitives.
//!
//! Galloping intersection for skewed sizes, sorted merge for balanced.

/// Count common elements in two sorted u32 slices.
///
/// Adaptively picks the best algorithm:
/// - When one slice is much smaller (|a| * 8 < |b|), uses galloping search
///   which is O(|small| * log(|large|)) and beats O(|a| + |b|) merge.
/// - Otherwise uses the standard sorted merge at O(|a| + |b|).
#[inline]
pub fn sparse_intersection_count(a: &[u32], b: &[u32]) -> usize {
    if a.is_empty() || b.is_empty() {
        return 0;
    }

    // Ensure a is the smaller slice for the skew check.
    let (small, large) = if a.len() <= b.len() { (a, b) } else { (b, a) };
    // Skewed sizes gallop before any block kernel: a block kernel walks the
    // large side linearly.
    if small.len() * 8 < large.len() {
        return galloping_intersection_count(small, large);
    }

    #[cfg(target_arch = "x86_64")]
    if let Some(count) = crate::core::simd_math::simd_intersection_count(small, large) {
        return count;
    }

    merge_intersection_count(small, large)
}

/// Sorted merge intersection count: O(|a| + |b|).
#[inline]
fn merge_intersection_count(a: &[u32], b: &[u32]) -> usize {
    let mut count = 0;
    let mut i = 0;
    let mut j = 0;
    while i < a.len() && j < b.len() {
        match a[i].cmp(&b[j]) {
            std::cmp::Ordering::Less => i += 1,
            std::cmp::Ordering::Greater => j += 1,
            std::cmp::Ordering::Equal => {
                count += 1;
                i += 1;
                j += 1;
            }
        }
    }
    count
}

/// Galloping (exponential search) intersection: O(|small| * log(|large|)).
///
/// For each element in `small`, does an exponential search in `large` to find
/// the matching position. Maintains a cursor into `large` so total work on
/// the large array is bounded by O(|small| * log(|large|/|small|)).
#[inline]
fn galloping_intersection_count(small: &[u32], large: &[u32]) -> usize {
    let mut count = 0;
    let mut lo = 0; // cursor into large; advances monotonically
    for &val in small {
        // Skip elements in large that are less than val using exponential search.
        // First, find a bound by doubling the step.
        let mut step = 1;
        while lo + step < large.len() && large[lo + step] < val {
            step *= 2;
        }
        // Binary search within [lo, min(lo+step, len))
        let hi = (lo + step).min(large.len());
        // Find first position >= val
        lo = binary_search_left(&large[lo..hi], val) + lo;
        if lo < large.len() && large[lo] == val {
            count += 1;
            lo += 1; // move past this match
        }
    }
    count
}

/// Binary search for the leftmost position where `slice[pos] >= target`.
#[inline]
fn binary_search_left(slice: &[u32], target: u32) -> usize {
    let mut lo = 0;
    let mut hi = slice.len();
    while lo < hi {
        let mid = lo + (hi - lo) / 2;
        if slice[mid] < target {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    lo
}

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

    fn naive(a: &[u32], b: &[u32]) -> usize {
        let set: std::collections::HashSet<u32> = b.iter().copied().collect();
        a.iter().filter(|x| set.contains(x)).count()
    }

    struct Lcg(u64);
    impl Lcg {
        fn next(&mut self) -> u64 {
            self.0 = self
                .0
                .wrapping_mul(6364136223846793005)
                .wrapping_add(1442695040888963407);
            self.0 >> 33
        }
    }

    /// Sorted, deduplicated sample of `n` values below `universe` (`universe >= n`).
    fn sample(rng: &mut Lcg, n: usize, universe: u64) -> Vec<u32> {
        let mut set = std::collections::BTreeSet::new();
        while set.len() < n {
            set.insert((rng.next() % universe) as u32);
        }
        set.into_iter().collect()
    }

    fn check_all(a: &[u32], b: &[u32]) {
        let want = naive(a, b);
        assert_eq!(sparse_intersection_count(a, b), want, "dispatch a,b");
        assert_eq!(sparse_intersection_count(b, a), want, "dispatch b,a");
        assert_eq!(merge_intersection_count(a, b), want, "merge");
        let (small, large) = if a.len() <= b.len() { (a, b) } else { (b, a) };
        assert_eq!(galloping_intersection_count(small, large), want, "gallop");
        #[cfg(target_arch = "x86_64")]
        if let Some(c) = crate::core::simd_math::simd_intersection_count(a, b) {
            assert_eq!(c, want, "avx2");
        }
    }

    #[test]
    fn all_kernels_match_naive_count() {
        for seed in 0..300u64 {
            let mut rng = Lcg(seed);
            let na = (rng.next() % 2001) as usize;
            let nb = (rng.next() % 2001) as usize;
            // A universe near max(na, nb) forces heavy overlap; a wide one forces little.
            let universe = (na.max(nb) as u64 + 1) * (1 + rng.next() % 8);
            check_all(
                &sample(&mut rng, na, universe),
                &sample(&mut rng, nb, universe),
            );
        }
    }

    #[test]
    fn kernels_match_on_edge_shapes_and_tails() {
        let mut rng = Lcg(7);
        for na in 0..=17usize {
            for nb in (0..=17usize).chain([64, 65, 600]) {
                for universe in [40u64, 2000, 1 << 20] {
                    let universe = universe.max(na.max(nb) as u64);
                    check_all(
                        &sample(&mut rng, na, universe),
                        &sample(&mut rng, nb, universe),
                    );
                }
            }
        }
        let ident = sample(&mut rng, 777, 5000);
        check_all(&ident, &ident);
        let evens: Vec<u32> = (0..1000).map(|x| x * 2).collect();
        let odds: Vec<u32> = (0..1000).map(|x| x * 2 + 1).collect();
        check_all(&evens, &odds);
        check_all(&[], &odds);
        check_all(&[], &[]);
    }
}