1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

//! Defines common traits

use hecs::{Entity, Fetch, Query, QueryBorrow};

/// Traits for types which represent a view or subset of some other type.
pub trait View<'a> {
    /// The type which View comes from
    type Superset;
    /// Splits from the containing superset
    fn split(orig: Self::Superset) -> Self;
}

// Implement view for self. A set is always its own subset

impl<'a, T> View<'a> for &'a T {
    type Superset = Self;

    fn split(orig: Self::Superset) -> Self {
        orig
    }
}

impl<'a, T> View<'a> for &'a mut T {
    type Superset = Self;

    fn split(orig: Self::Superset) -> Self {
        orig
    }
}

/// Extends the queries for additional paralell operation
pub trait QueryExt {
    /// Item returned by the query
    type Item;
    /// Execute a function for each item of the query in pararell using rayon.
    #[cfg(feature = "parallel")]
    fn par_for_each(self, batch_size: u32, func: impl Fn((Entity, Self::Item)) + Send + Sync);
    /// Fallible version of [`QueryBorrow::par_for_each`]
    #[cfg(feature = "parallel")]
    fn try_par_for_each<E: Send>(
        self,
        batch_size: u32,
        func: impl Fn((Entity, Self::Item)) -> Result<(), E> + Send + Sync,
    ) -> Result<(), E>;
}

impl<'w, 'q, Q> QueryExt for &'q mut QueryBorrow<'w, Q>
where
    Q: Query,
    <Q::Fetch as Fetch<'q>>::Item: Send,
{
    type Item = <Q::Fetch as Fetch<'q>>::Item;

    #[cfg(feature = "parallel")]
    fn par_for_each(self, batch_size: u32, func: impl Fn((Entity, Self::Item)) + Send + Sync) {
        use rayon::iter::{ParallelBridge, ParallelIterator};
        self.iter_batched(batch_size)
            .par_bridge()
            .for_each(|batch| batch.for_each(&func))
    }

    #[cfg(feature = "parallel")]
    fn try_par_for_each<E: Send>(
        self,
        batch_size: u32,
        func: impl Fn((Entity, Self::Item)) -> Result<(), E> + Send + Sync,
    ) -> Result<(), E> {
        use rayon::iter::{ParallelBridge, ParallelIterator};
        self.iter_batched(batch_size)
            .par_bridge()
            .try_for_each(|mut batch| batch.try_for_each(&func))
    }
}