async-ecs/async-ecs/src/misc/bit_producer.rs

use hibitset::BitSetLike;
use parallel_iterator::{Folder, Producer};

use crate::misc::bit_average;

use super::BitIter;

pub struct BitProducer<'a, T> {
    pub(crate) iter: BitIter<&'a T>,
    pub(crate) splits: u8,
}

impl<'a, T> BitProducer<'a, T>
where
    T: BitSetLike,
{
    pub fn new(iter: BitIter<&'a T>, splits: u8) -> Self {
        Self { iter, splits }
    }

    fn handle_level(&mut self, level: usize) -> Option<Self> {
        if self.iter.masks[level] == 0 {
            None
        } else {
            let level_prefix = self.iter.prefix.get(level).cloned().unwrap_or(0);
            let first_bit = self.iter.masks[level].trailing_zeros();

            bit_average(self.iter.masks[level])
                .map(|average_bit| {
                    let mask = (1 << average_bit) - 1;
                    let mut other = BitProducer::new(
                        BitIter::new(self.iter.set, [0; LAYERS], [0; LAYERS - 1]),
                        self.splits,
                    );

                    other.iter.masks[level] = self.iter.masks[level] & !mask;
                    other.iter.prefix[level - 1] = (level_prefix | average_bit as u32) << BITS;
                    other.iter.prefix[level..].copy_from_slice(&self.iter.prefix[level..]);

                    self.iter.masks[level] &= mask;
                    self.iter.prefix[level - 1] = (level_prefix | first_bit) << BITS;

                    other
                })
                .or_else(|| {
                    let idx = level_prefix as usize | first_bit as usize;

                    self.iter.prefix[level - 1] = (idx as u32) << BITS;
                    self.iter.masks[level] = 0;
                    self.iter.masks[level - 1] = self.iter.set.get_from_layer(level - 1, idx);

                    None
                })
        }
    }
}

impl<'a, T> Producer for BitProducer<'a, T>
where
    T: BitSetLike + Send + Sync,
{
    type Item = u32;

    fn split(mut self) -> (Self, Option<Self>) {
        let other = {
            let top_layer = LAYERS - 1;
            let mut h = self.handle_level(top_layer);

            for i in 1..self.splits {
                h = h.or_else(|| self.handle_level(top_layer - i as usize));
            }

            h
        };

        (self, other)
    }

    fn fold_with<F>(self, folder: F) -> F
    where
        F: Folder<Self::Item>,
    {
        folder.consume_iter(self.iter)
    }
}

const LAYERS: usize = 4;

#[cfg(target_pointer_width = "64")]
pub const BITS: usize = 6;

#[cfg(target_pointer_width = "32")]
pub const BITS: usize = 5;