hecs/
batch.rs

1use crate::alloc::collections::BinaryHeap;
2use core::{any::TypeId, fmt, mem::MaybeUninit, slice};
3
4use crate::{
5    archetype::{TypeIdMap, TypeInfo},
6    Archetype, Component,
7};
8
9/// A collection of component types
10#[derive(Debug, Clone, Default)]
11pub struct ColumnBatchType {
12    types: BinaryHeap<TypeInfo>,
13}
14
15impl ColumnBatchType {
16    /// Create an empty type
17    pub fn new() -> Self {
18        Self::default()
19    }
20
21    /// Update to include `T` components
22    pub fn add<T: Component>(&mut self) -> &mut Self {
23        self.types.push(TypeInfo::of::<T>());
24        self
25    }
26
27    /// Construct a [`ColumnBatchBuilder`] for *exactly* `size` entities with these components
28    pub fn into_batch(self, size: u32) -> ColumnBatchBuilder {
29        let mut types = self.types.into_sorted_vec();
30        types.dedup();
31        let fill = TypeIdMap::with_capacity_and_hasher(types.len(), Default::default());
32        let mut arch = Archetype::new(types);
33        arch.reserve(size);
34        ColumnBatchBuilder {
35            fill,
36            target_fill: size,
37            archetype: Some(arch),
38        }
39    }
40}
41
42/// An incomplete collection of component data for entities with the same component types
43pub struct ColumnBatchBuilder {
44    /// Number of components written so far for each component type
45    fill: TypeIdMap<u32>,
46    target_fill: u32,
47    pub(crate) archetype: Option<Archetype>,
48}
49
50unsafe impl Send for ColumnBatchBuilder {}
51unsafe impl Sync for ColumnBatchBuilder {}
52
53impl ColumnBatchBuilder {
54    /// Create a batch for *exactly* `size` entities with certain component types
55    pub fn new(ty: ColumnBatchType, size: u32) -> Self {
56        ty.into_batch(size)
57    }
58
59    /// Get a handle for inserting `T` components if `T` was in the [`ColumnBatchType`]
60    pub fn writer<T: Component>(&mut self) -> Option<BatchWriter<'_, T>> {
61        let archetype = self.archetype.as_mut().unwrap();
62        let state = archetype.get_state::<T>()?;
63        let base = archetype.get_base::<T>(state);
64        Some(BatchWriter {
65            fill: self.fill.entry(TypeId::of::<T>()).or_insert(0),
66            storage: unsafe {
67                slice::from_raw_parts_mut(base.as_ptr().cast(), self.target_fill as usize)
68                    .iter_mut()
69            },
70        })
71    }
72
73    /// Finish the batch, failing if any components are missing
74    pub fn build(mut self) -> Result<ColumnBatch, BatchIncomplete> {
75        let mut archetype = self.archetype.take().unwrap();
76        if archetype
77            .types()
78            .iter()
79            .any(|ty| self.fill.get(&ty.id()).copied().unwrap_or(0) != self.target_fill)
80        {
81            return Err(BatchIncomplete { _opaque: () });
82        }
83        unsafe {
84            archetype.set_len(self.target_fill);
85        }
86        Ok(ColumnBatch(archetype))
87    }
88}
89
90impl Drop for ColumnBatchBuilder {
91    fn drop(&mut self) {
92        if let Some(archetype) = self.archetype.take() {
93            for ty in archetype.types() {
94                let fill = self.fill.get(&ty.id()).copied().unwrap_or(0);
95                unsafe {
96                    let base = archetype.get_dynamic(ty.id(), 0, 0).unwrap();
97                    for i in 0..fill {
98                        base.as_ptr().add(i as usize).drop_in_place()
99                    }
100                }
101            }
102        }
103    }
104}
105
106/// A collection of component data for entities with the same component types
107pub struct ColumnBatch(pub(crate) Archetype);
108
109/// Handle for appending components
110pub struct BatchWriter<'a, T> {
111    fill: &'a mut u32,
112    storage: core::slice::IterMut<'a, MaybeUninit<T>>,
113}
114
115impl<T> BatchWriter<'_, T> {
116    /// Add a component if there's space remaining
117    pub fn push(&mut self, x: T) -> Result<(), T> {
118        match self.storage.next() {
119            None => Err(x),
120            Some(slot) => {
121                *slot = MaybeUninit::new(x);
122                *self.fill += 1;
123                Ok(())
124            }
125        }
126    }
127
128    /// How many components have been added so far
129    pub fn fill(&self) -> u32 {
130        *self.fill
131    }
132}
133
134/// Error indicating that a [`ColumnBatchBuilder`] was missing components
135#[derive(Debug, Clone, Eq, PartialEq, Hash)]
136pub struct BatchIncomplete {
137    _opaque: (),
138}
139
140#[cfg(feature = "std")]
141impl std::error::Error for BatchIncomplete {}
142
143impl fmt::Display for BatchIncomplete {
144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
145        f.write_str("batch incomplete")
146    }
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    #[test]
154    fn empty_batch() {
155        let mut types = ColumnBatchType::new();
156        types.add::<usize>();
157        let mut builder = types.into_batch(0);
158        let mut writer = builder.writer::<usize>().unwrap();
159        assert!(writer.push(42).is_err());
160    }
161}
hecs/batch.rs

hecs/
batch.rs
