flax 0.7.1

An ergonomic archetypical ECS
Documentation
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
mod as_deref;
mod cloned;
mod component;
mod component_mut;
mod copied;
mod entity_ref;
mod ext;
mod map;
mod maybe_mut;
mod opt;
mod read_only;
mod relations;
mod relations_mut;
mod satisfied;
mod source;
mod transform;

use crate::{
    archetype::{Archetype, ArchetypeId, Slice, Slot},
    filter::{RefFetch, StaticFilter},
    system::Access,
    util::Ptr,
    ArchetypeSearcher, Entity, World,
};
use alloc::vec::Vec;
use core::fmt::Debug;
use core::fmt::{self, Formatter};

pub use as_deref::*;
pub use cloned::*;
pub use component::*;
pub use component_mut::*;
pub use entity_ref::*;
pub use ext::FetchExt;
pub use map::Map;
pub use maybe_mut::{MaybeMut, MutGuard};
pub use opt::*;
pub use read_only::*;
pub use relations::{nth_relation, relations_like, NthRelation, Relations, RelationsIter};
pub use relations_mut::{relations_like_mut, RelationsIterMut, RelationsMut};
pub use satisfied::Satisfied;
pub use source::Source;
pub use transform::{Added, Modified, TransformFetch};

#[doc(hidden)]
pub struct FmtQuery<'r, Q>(pub &'r Q);

impl<'r, 'w, Q> Debug for FmtQuery<'r, Q>
where
    Q: Fetch<'w>,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        self.0.describe(f)
    }
}

/// Represents the world data necessary for declaring fetch access
#[derive(Copy, Clone)]
pub struct FetchAccessData<'w> {
    /// The current world
    pub world: &'w World,
    /// The iterated archetype to prepare for
    pub arch: &'w Archetype,
    /// The archetype id
    pub arch_id: ArchetypeId,
}

impl<'w> From<FetchPrepareData<'w>> for FetchAccessData<'w> {
    fn from(value: FetchPrepareData<'w>) -> Self {
        Self {
            world: value.world,
            arch: value.arch,
            arch_id: value.arch_id,
        }
    }
}

/// Represents the world data necessary for preparing a fetch
#[derive(Copy, Clone)]
pub struct FetchPrepareData<'w> {
    /// The current world
    pub world: &'w World,
    /// The iterated archetype to prepare for
    pub arch: &'w Archetype,
    /// The archetype id
    pub arch_id: ArchetypeId,
    /// The tick the previous time the query executed
    pub old_tick: u32,
    /// The new tick to write if query is mutable
    pub new_tick: u32,
}

/// Trait which gives an associated `Item` fetch type
pub trait FetchItem<'q> {
    /// The item yielded by the prepared fetch
    type Item;
}

/// A fetch describes a retrieval of data from the world and archetypes during a query.
///
/// A fetch is prepared, wherein borrows are acquired and a `PreparedFetch` is returned, which is
/// used to provide the query with values.
///
/// The PreparedFetch can in turn control the ranges of slots which are requested by the query,
/// e.g; filtering changed components
pub trait Fetch<'w>: for<'q> FetchItem<'q> {
    /// true if the fetch mutates any component and thus needs a change event
    const MUTABLE: bool;

    /// The prepared version of the fetch
    type Prepared: for<'x> PreparedFetch<'x, Item = <Self as FetchItem<'x>>::Item> + 'w;

    /// Prepares the fetch for an archetype by acquiring borrows.
    ///
    /// Returns `None` if the archetype does not match.
    fn prepare(&'w self, data: FetchPrepareData<'w>) -> Option<Self::Prepared>;

    /// Returns true if the archetype matches the fetch
    fn filter_arch(&self, data: FetchAccessData) -> bool;

    /// Returns which components and how will be accessed for an archetype.
    fn access(&self, data: FetchAccessData, dst: &mut Vec<Access>);

    /// Describes the fetch in a human-readable fashion
    fn describe(&self, f: &mut Formatter<'_>) -> fmt::Result;

    /// Returns the required component for the fetch.
    ///
    /// This is used for the query to determine which archetypes to visit
    #[inline]
    fn searcher(&self, _searcher: &mut ArchetypeSearcher) {}

    /// Convert the fetch to a reference type which works with `HRTB`
    #[inline]
    fn by_ref(&self) -> RefFetch<Self>
    where
        Self: Sized,
    {
        RefFetch(self)
    }
}

/// Borrowed state for a fetch
pub trait PreparedFetch<'q> {
    /// Item returned by fetch
    type Item: 'q;
    /// A chunk accessing a disjoint set of the borrow sequentially
    type Chunk: 'q;

    /// Indicates if the fetch will provide any filtering of slots.
    ///
    /// If `false`, this fetch will unconditionally yield.
    ///
    /// This is used to influence the default behavior of optional queries.
    const HAS_FILTER: bool;

    /// Creates a chunk to access a slice of the borrow
    ///
    /// # Safety
    ///
    /// `slots` must be disjoint to all other currently existing chunks
    unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk;

    /// Fetch the item from entity at the slot in the prepared storage.
    /// # Safety
    /// Must return non-aliased references to the underlying borrow of the
    /// prepared archetype.
    ///
    /// The callee is responsible for assuring disjoint calls.
    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item;

    #[inline]
    /// Filter the slots to visit
    /// Returns the leftmost subslice of `slots` which should be visited
    ///
    /// # Safety
    /// `slots` must not overlap any alive references returned by `fetch`
    unsafe fn filter_slots(&mut self, slots: Slice) -> Slice {
        slots
    }
}

/// Allows filtering the constituent parts of a fetch using a set union
pub trait UnionFilter {
    /// The union may not have the same filter behavior as a normal filter as sub-filters
    /// are combined using *or* instead of *and*. This means any non-filter will cause the
    /// whole tuple to always yield.
    const HAS_UNION_FILTER: bool;
    // Filter the slots using a union operation of the constituent part
    ///
    /// # Safety
    /// See: [`PreparedFetch::filter_slots`]
    unsafe fn filter_union(&mut self, slots: Slice) -> Slice;
}

impl<'q, F> PreparedFetch<'q> for &'q mut F
where
    F: PreparedFetch<'q>,
{
    type Item = F::Item;
    type Chunk = F::Chunk;

    const HAS_FILTER: bool = F::HAS_FILTER;

    unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk {
        (*self).create_chunk(slots)
    }

    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
        F::fetch_next(chunk)
    }

    unsafe fn filter_slots(&mut self, slots: Slice) -> Slice {
        (*self).filter_slots(slots)
    }
}

impl<'q> FetchItem<'q> for () {
    type Item = ();
}

impl UnionFilter for () {
    const HAS_UNION_FILTER: bool = false;

    unsafe fn filter_union(&mut self, slots: Slice) -> Slice {
        slots
    }
}

impl<'w> Fetch<'w> for () {
    const MUTABLE: bool = false;

    type Prepared = ();

    fn prepare(&self, _: FetchPrepareData<'w>) -> Option<Self::Prepared> {
        Some(())
    }

    fn filter_arch(&self, _: FetchAccessData) -> bool {
        true
    }

    fn describe(&self, f: &mut Formatter) -> fmt::Result {
        write!(f, "()")
    }

    #[inline]
    fn access(&self, _: FetchAccessData, _: &mut Vec<Access>) {}
}

impl<'q> RandomFetch<'q> for () {
    #[inline]
    unsafe fn fetch_shared(&'q self, _: Slot) -> Self::Item {}
    #[inline]
    unsafe fn fetch_shared_chunk(_: &Self::Chunk, _: Slot) -> Self::Item {}
}

impl<'q> PreparedFetch<'q> for () {
    type Item = ();

    type Chunk = ();

    const HAS_FILTER: bool = false;

    #[inline]
    unsafe fn create_chunk(&'q mut self, _: Slice) -> Self::Chunk {}

    #[inline]
    unsafe fn fetch_next(_: &mut Self::Chunk) -> Self::Item {}
}

// impl<'q, F> PreparedFetch<'q> for Option<F>
// where
//     F: PreparedFetch<'q>,
// {
//     type Item = Option<F::Item>;
//     type Chunk = Option<F::Chunk>;

//     unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk {
//         self.as_mut().map(|fetch| fetch.create_chunk(slots))
//     }

//     #[inline]
//     unsafe fn filter_slots(&mut self, slots: Slice) -> Slice {
//         if let Some(fetch) = self {
//             fetch.filter_slots(slots)
//         } else {
//             Slice::new(slots.end, slots.end)
//         }
//     }

//     unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
//         batch.as_mut().map(|fetch| F::fetch_next(fetch))
//     }
// }

#[derive(Debug, Clone)]
/// Returns the entity ids
pub struct EntityIds;
#[doc(hidden)]
pub struct ReadEntities<'a> {
    entities: &'a [Entity],
}

impl<'q> FetchItem<'q> for EntityIds {
    type Item = Entity;
}

impl<'w> Fetch<'w> for EntityIds {
    const MUTABLE: bool = false;

    type Prepared = ReadEntities<'w>;

    fn prepare(&self, data: FetchPrepareData<'w>) -> Option<Self::Prepared> {
        Some(ReadEntities {
            entities: data.arch.entities(),
        })
    }

    fn filter_arch(&self, _: FetchAccessData) -> bool {
        true
    }

    fn describe(&self, f: &mut Formatter) -> fmt::Result {
        f.write_str("entity_ids")
    }

    #[inline]
    fn access(&self, _: FetchAccessData, _: &mut Vec<Access>) {}
}

impl<'w, 'q> PreparedFetch<'q> for ReadEntities<'w> {
    type Item = Entity;
    type Chunk = Ptr<'q, Entity>;

    const HAS_FILTER: bool = false;

    unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk {
        Ptr::new(self.entities[slots.as_range()].as_ptr())
    }

    unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
        let old = chunk.as_ptr();
        *chunk = chunk.add(1);
        *old
    }
}

impl<'w, 'q> RandomFetch<'q> for ReadEntities<'w> {
    #[inline]
    unsafe fn fetch_shared(&self, slot: usize) -> Self::Item {
        self.entities[slot]
    }

    unsafe fn fetch_shared_chunk(chunk: &Self::Chunk, slot: Slot) -> Self::Item {
        *chunk.add(slot).as_ref()
    }
}

// Implement for tuples
macro_rules! tuple_impl {
    ($($idx: tt => $ty: ident),*) => {
        impl<'q, $($ty, )*> FetchItem<'q> for ($($ty,)*)
        where $($ty: FetchItem<'q>,)*
        {
            type Item = ($($ty::Item,)*);

        }

        impl<'q, $($ty, )*> RandomFetch<'q> for ($($ty,)*)
        where $($ty: RandomFetch<'q>,)*
        {

            #[inline(always)]
            unsafe fn fetch_shared(&'q self, slot: Slot) -> Self::Item {
                ($(
                    (self.$idx).fetch_shared(slot),
                )*)
            }

            #[inline(always)]
            unsafe fn fetch_shared_chunk(chunk: &Self::Chunk, slot: Slot) -> Self::Item {
                ($(
                    $ty::fetch_shared_chunk(&chunk.$idx, slot),
                )*)
            }
        }


        impl<'q, $($ty, )*> PreparedFetch<'q> for ($($ty,)*)
            where $($ty: PreparedFetch<'q>,)*
        {

            type Item = ($($ty::Item,)*);
            type Chunk = ($($ty::Chunk,)*);

            const HAS_FILTER: bool =  $($ty::HAS_FILTER )||*;

            #[inline]
            unsafe fn fetch_next(chunk: &mut Self::Chunk) -> Self::Item {
                ($(
                    $ty::fetch_next(&mut chunk.$idx),
                )*)
            }

            #[inline]
            unsafe fn create_chunk(&'q mut self, slots: Slice) -> Self::Chunk {
                ($((self.$idx).create_chunk(slots),)*)
            }

            #[inline]
            unsafe fn filter_slots(&mut self, mut slots: Slice) -> Slice {
                $(

                    slots = self.$idx.filter_slots(slots);
                )*

                slots
            }
        }

        impl<'q, $($ty, )*> UnionFilter for ($($ty,)*)
            where $($ty: PreparedFetch<'q>,)*
        {

            const HAS_UNION_FILTER: bool =  $($ty::HAS_FILTER )&&*;

            #[inline]
            unsafe fn filter_union(&mut self, slots: Slice) -> Slice {
                [
                    // Don't leak union into this
                    $( self.$idx.filter_slots(slots)),*
                ]
                .into_iter()
                .min()
                .unwrap_or_default()
            }
        }

        impl<'w, $($ty, )*> Fetch<'w> for ($($ty,)*)
        where $($ty: Fetch<'w>,)*
        {
            const MUTABLE: bool =  $($ty::MUTABLE )||*;
            type Prepared       = ($($ty::Prepared,)*);

            #[inline]
            fn prepare(&'w self, data: FetchPrepareData<'w>) -> Option<Self::Prepared> {
                Some( ($( (self.$idx).prepare(data)?,)*) )
            }

            #[inline]
            fn filter_arch(&self, data:FetchAccessData) -> bool {
                ( $((self.$idx).filter_arch(data)) && * )
            }

            #[inline]
            fn describe(&self, f: &mut Formatter) -> fmt::Result {
                Debug::fmt(&($(FmtQuery(&self.$idx),)*), f)
            }

            #[inline]
            fn access(&self, data: FetchAccessData, dst: &mut Vec<Access>) {
                $( (self.$idx).access(data, dst);)*
            }

            #[inline]
            fn searcher(&self, searcher: &mut ArchetypeSearcher) {
                $((self.$idx).searcher(searcher));*
            }
        }

        impl< $($ty: StaticFilter, )*> StaticFilter for ($($ty,)*)
        {
            fn filter_static(&self, arch: &Archetype) -> bool {
                ( $((self.$idx).filter_static(arch)) && * )
            }

        }
    };
}

tuple_impl! { 0 => A }
tuple_impl! { 0 => A, 1 => B }
tuple_impl! { 0 => A, 1 => B, 2 => C }
tuple_impl! { 0 => A, 1 => B, 2 => C, 3 => D }
tuple_impl! { 0 => A, 1 => B, 2 => C, 3 => D, 4 => E }
tuple_impl! { 0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F }
tuple_impl! { 0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => H }
tuple_impl! { 0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => H, 7 => I }
tuple_impl! { 0 => A, 1 => B, 2 => C, 3 => D, 4 => E, 5 => F, 6 => H, 7 => I, 8 => J }