roaring 0.12.0

use crate::bitmap::store::array_store::vector::count_set_lanes;
use fearless_simd::{mask16x8, prelude::*, u16x8};
use fearless_simd_macros::simd;

#[cfg(not(feature = "std"))]
use alloc::vec::Vec;

/// This visitor pattern allows multiple different algorithms to be written over the same data
/// For example: vectorized algorithms can pass a visitor off to a scalar algorithm to finish off
/// a tail that is not a multiple of the vector width.
///
/// Perhaps more importantly: it separates the set algorithms from the operations performed on
/// their results. Future work can utilize the exiting algorithms to trivially implement
/// computing the cardinality of an operation without materializng a new bitmap.
pub trait BinaryOperationVisitor {
    fn visit_vector<S: Simd>(&mut self, value: u16x8<S>, mask: mask16x8<S>);
    fn visit_scalar(&mut self, value: u16);
    fn visit_slice(&mut self, values: &[u16]);
}

/// A simple visitor that stores the computation result to a Vec
/// accessible by calling `into_inner()`
pub struct VecWriter {
    vec: Vec<u16>,
}

impl VecWriter {
    /// `capacity` must be an upper bound on the length of the output.
    /// `visit_scalar` panics rather than grow past it.
    pub fn new(capacity: usize) -> VecWriter {
        // `visit_vector` writes a whole vector before truncating it. The vector algorithms only
        // call it when both inputs hold a whole vector, and so when `capacity` is at least 8.
        // Below that, keep tiny outputs in the smallest allocation.
        let capacity = if capacity < 8 { capacity } else { capacity + 8 };
        let vec = Vec::with_capacity(capacity);
        VecWriter { vec }
    }

    pub fn into_inner(self) -> Vec<u16> {
        // Consider shrinking the vec here.
        // Exactly len could be too aggressive. Len rounded up to next power of 2?
        // Related, but not exact issue: https://github.com/RoaringBitmap/roaring-rs/issues/136
        self.vec
    }
}

impl BinaryOperationVisitor for VecWriter {
    #[simd]
    fn visit_vector<S: Simd>(&mut self, value: u16x8<S>, mask: mask16x8<S>) {
        let result = value.compress(mask);
        let result_size = count_set_lanes(mask);

        // This idiom is better than subslicing result, as it compiles down to an unaligned vector
        // store instr.
        // A more straightforward, but unsafe way would be ptr::write_unaligned and Vec::set_len
        // Writing a vector at once is why the vectorized algorithms do not operate in place
        // first write the entire vector
        self.vec.extend_from_slice(result.as_slice());
        // next truncate the masked out values
        self.vec.truncate(self.vec.len() - (result.len() - result_size));
    }

    fn visit_scalar(&mut self, value: u16) {
        // Lets the compiler drop the grow path of `push` from the merge loops.
        assert!(self.vec.len() < self.vec.capacity());
        self.vec.push(value)
    }

    fn visit_slice(&mut self, values: &[u16]) {
        self.vec.extend_from_slice(values);
    }
}

pub struct CardinalityCounter {
    count: usize,
}

impl CardinalityCounter {
    pub fn new() -> CardinalityCounter {
        CardinalityCounter { count: 0 }
    }

    pub fn into_inner(self) -> u64 {
        self.count as u64
    }
}

impl BinaryOperationVisitor for CardinalityCounter {
    fn visit_vector<S: Simd>(&mut self, _value: u16x8<S>, mask: mask16x8<S>) {
        self.count += count_set_lanes(mask);
    }

    fn visit_scalar(&mut self, _value: u16) {
        self.count += 1;
    }

    fn visit_slice(&mut self, values: &[u16]) {
        self.count += values.len();
    }
}