intuicio_framework_ecs/
prefab.rs

1use crate::{
2    archetype::{Archetype, ArchetypeColumnInfo, ArchetypeError},
3    bundle::Bundle,
4    entity::Entity,
5    processor::{WorldProcessor, WorldProcessorEntityMapping},
6    world::{Relation, World, WorldError},
7    Component,
8};
9use intuicio_core::{registry::Registry, types::TypeQuery};
10use intuicio_data::type_hash::TypeHash;
11use intuicio_framework_serde::{
12    from_intermediate, to_intermediate, Intermediate, IntermediateResult, SerializationRegistry,
13};
14use serde::{de::DeserializeOwned, Deserialize, Serialize};
15use std::{collections::HashMap, error::Error};
16
17#[derive(Debug)]
18pub enum PrefabError {
19    CouldNotFindType(TypeHash),
20    CouldNotSerializeType {
21        type_name: String,
22        module_name: Option<String>,
23    },
24    CouldNotDeserializeType {
25        type_name: String,
26        module_name: Option<String>,
27    },
28    World(WorldError),
29}
30
31impl std::fmt::Display for PrefabError {
32    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        match self {
34            Self::CouldNotFindType(type_hash) => {
35                write!(f, "Could not find type by hash: {:?}", type_hash)
36            }
37            Self::CouldNotSerializeType {
38                type_name,
39                module_name,
40            } => write!(
41                f,
42                "Could not serialize type: {}::{}",
43                module_name.as_deref().unwrap_or_default(),
44                type_name
45            ),
46            Self::CouldNotDeserializeType {
47                type_name,
48                module_name,
49            } => write!(
50                f,
51                "Could not deserialize type: {}::{}",
52                module_name.as_deref().unwrap_or_default(),
53                type_name
54            ),
55            Self::World(error) => write!(f, "World error: {}", error),
56        }
57    }
58}
59
60impl From<WorldError> for PrefabError {
61    fn from(value: WorldError) -> Self {
62        Self::World(value)
63    }
64}
65
66impl From<ArchetypeError> for PrefabError {
67    fn from(value: ArchetypeError) -> Self {
68        Self::World(WorldError::Archetype(value))
69    }
70}
71
72impl Error for PrefabError {}
73
74#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
75pub struct PrefabArchetypeColumn {
76    pub type_name: String,
77    pub module_name: Option<String>,
78    pub components: Vec<Intermediate>,
79}
80
81#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
82pub struct PrefabArchetype {
83    pub entities: Vec<Entity>,
84    pub columns: Vec<PrefabArchetypeColumn>,
85}
86
87#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
88pub struct Prefab {
89    pub archetypes: Vec<PrefabArchetype>,
90}
91
92impl Prefab {
93    pub fn register_relation_serializer<T: Serialize + DeserializeOwned + Component>(
94        serialization: &mut SerializationRegistry,
95    ) {
96        serialization.register::<Relation<T>>(
97            |data| {
98                Ok(Intermediate::Seq(
99                    data.iter()
100                        .map(|(item, entity)| {
101                            Ok(Intermediate::Tuple(vec![
102                                to_intermediate(item)?,
103                                to_intermediate(&entity)?,
104                            ]))
105                        })
106                        .collect::<IntermediateResult<Vec<_>>>()?,
107                ))
108            },
109            |data, value| {
110                let Intermediate::Seq(items) = value else {
111                    return Err("Expected intermediate sequence".into());
112                };
113                for tuple in items {
114                    let Intermediate::Tuple(tuple) = tuple else {
115                        return Err("Expected intermediate tuple".into());
116                    };
117                    if tuple.len() != 2 {
118                        return Err("Expected tuple to have 2 items".into());
119                    }
120                    data.add(from_intermediate(&tuple[0])?, from_intermediate(&tuple[1])?);
121                }
122                Ok(())
123            },
124        );
125    }
126
127    pub fn from_world<const LOCKING: bool>(
128        world: &World,
129        serialization: &SerializationRegistry,
130        registry: &Registry,
131    ) -> Result<Self, PrefabError> {
132        let archetypes = world
133            .archetypes()
134            .map(|archetype| {
135                let entities = archetype.entities().iter().collect();
136                let columns = archetype
137                    .columns()
138                    .map(|column| {
139                        let type_ = registry
140                            .find_type(TypeQuery {
141                                type_hash: Some(column.type_hash()),
142                                ..Default::default()
143                            })
144                            .ok_or_else(|| PrefabError::CouldNotFindType(column.type_hash()))?;
145                        let components =
146                            archetype.dynamic_column_iter::<LOCKING>(column.type_hash(), false)?;
147                        let components = components
148                            .map(|component| unsafe {
149                                serialization
150                                    .dynamic_serialize_from(column.type_hash(), component.data())
151                                    .map_err(|_| PrefabError::CouldNotSerializeType {
152                                        type_name: type_.type_name().to_owned(),
153                                        module_name: type_
154                                            .module_name()
155                                            .map(|name| name.to_owned()),
156                                    })
157                            })
158                            .collect::<Result<_, PrefabError>>()?;
159                        Ok(PrefabArchetypeColumn {
160                            type_name: type_.type_name().to_owned(),
161                            module_name: type_.module_name().map(|name| name.to_owned()),
162                            components,
163                        })
164                    })
165                    .collect::<Result<_, PrefabError>>()?;
166                Ok(PrefabArchetype { entities, columns })
167            })
168            .collect::<Result<_, PrefabError>>()?;
169        Ok(Self { archetypes })
170    }
171
172    pub fn from_entities<const LOCKING: bool>(
173        world: &World,
174        entities: impl IntoIterator<Item = Entity>,
175        processor: &WorldProcessor,
176        serialization: &SerializationRegistry,
177        registry: &Registry,
178    ) -> Result<Self, PrefabError> {
179        let mut total_entities = Vec::default();
180        processor.all_related_entities::<LOCKING>(world, entities, &mut total_entities)?;
181        let mut archetype_rows = HashMap::<u32, (&Archetype, Vec<usize>)>::new();
182        for entity in total_entities {
183            let id = world.entity_archetype_id(entity)?;
184            if let Some((archetype, rows)) = archetype_rows.get_mut(&id) {
185                rows.push(
186                    archetype
187                        .entities()
188                        .index_of(entity)
189                        .ok_or(WorldError::EntityDoesNotExists { entity })?,
190                );
191            } else {
192                let archetype = world.archetype_by_id(id)?;
193                archetype_rows.insert(
194                    id,
195                    (
196                        archetype,
197                        vec![archetype
198                            .entities()
199                            .index_of(entity)
200                            .ok_or(WorldError::EntityDoesNotExists { entity })?],
201                    ),
202                );
203            }
204        }
205        let archetypes = archetype_rows
206            .into_values()
207            .map(|(archetype, rows)| {
208                let entities = rows
209                    .iter()
210                    .map(|row| archetype.entities().get(*row).unwrap())
211                    .collect();
212                let columns = archetype
213                    .columns()
214                    .map(|column| {
215                        let type_ = registry
216                            .find_type(TypeQuery {
217                                type_hash: Some(column.type_hash()),
218                                ..Default::default()
219                            })
220                            .ok_or_else(|| PrefabError::CouldNotFindType(column.type_hash()))?;
221                        let components =
222                            archetype.dynamic_column::<LOCKING>(column.type_hash(), false)?;
223                        let components = rows
224                            .iter()
225                            .map(|row| unsafe {
226                                serialization
227                                    .dynamic_serialize_from(
228                                        column.type_hash(),
229                                        components.data(*row)?,
230                                    )
231                                    .map_err(|_| PrefabError::CouldNotSerializeType {
232                                        type_name: type_.type_name().to_owned(),
233                                        module_name: type_
234                                            .module_name()
235                                            .map(|name| name.to_owned()),
236                                    })
237                            })
238                            .collect::<Result<_, PrefabError>>()?;
239                        Ok(PrefabArchetypeColumn {
240                            type_name: type_.type_name().to_owned(),
241                            module_name: type_.module_name().map(|name| name.to_owned()),
242                            components,
243                        })
244                    })
245                    .collect::<Result<_, PrefabError>>()?;
246                Ok(PrefabArchetype { entities, columns })
247            })
248            .collect::<Result<_, PrefabError>>()?;
249        Ok(Self { archetypes })
250    }
251
252    pub fn to_world<const LOCKING: bool>(
253        &self,
254        processor: &WorldProcessor,
255        serialization: &SerializationRegistry,
256        registry: &Registry,
257        additional_components: impl Bundle + Clone,
258    ) -> Result<(World, HashMap<Entity, Entity>), PrefabError> {
259        let additional_columns = additional_components.columns();
260        let mut mappings = HashMap::<_, _>::default();
261        let mut world = World::default();
262        for archetype in &self.archetypes {
263            let column_types = archetype
264                .columns
265                .iter()
266                .map(|column| {
267                    registry
268                        .find_type(TypeQuery {
269                            name: Some(column.type_name.as_str().into()),
270                            module_name: column
271                                .module_name
272                                .as_ref()
273                                .map(|name| name.as_str().into()),
274                            ..Default::default()
275                        })
276                        .ok_or_else(|| PrefabError::CouldNotDeserializeType {
277                            type_name: column.type_name.to_owned(),
278                            module_name: column.module_name.to_owned(),
279                        })
280                })
281                .collect::<Result<Vec<_>, PrefabError>>()?;
282            let column_info = column_types
283                .iter()
284                .map(|type_| ArchetypeColumnInfo::from_type(type_))
285                .chain(additional_columns.iter().cloned())
286                .collect::<Vec<_>>();
287            let rows_count = archetype
288                .columns
289                .iter()
290                .map(|column| column.components.len())
291                .min()
292                .unwrap_or_default();
293            for index in 0..rows_count {
294                unsafe {
295                    let (entity, access) = world.spawn_uninitialized_raw(column_info.to_owned())?;
296                    for ((column, info), type_) in archetype
297                        .columns
298                        .iter()
299                        .zip(column_info.iter())
300                        .zip(column_types.iter())
301                    {
302                        let data = access.data(info.type_hash())?;
303                        let component = &column.components[index];
304                        access.initialize_raw(type_)?;
305                        serialization
306                            .dynamic_deserialize_to(info.type_hash(), data, component)
307                            .map_err(|_| PrefabError::CouldNotDeserializeType {
308                                type_name: column.type_name.to_owned(),
309                                module_name: column.module_name.to_owned(),
310                            })?;
311                    }
312                    additional_components.clone().initialize_into(&access);
313                    mappings.insert(archetype.entities[index], entity);
314                }
315            }
316        }
317        for archetype in world.archetypes() {
318            for column in archetype.columns() {
319                let access = archetype.dynamic_column::<LOCKING>(column.type_hash(), true)?;
320                for index in 0..archetype.len() {
321                    unsafe {
322                        processor.remap_entities_raw(
323                            column.type_hash(),
324                            access.data(index)?,
325                            WorldProcessorEntityMapping::new(&mappings),
326                        );
327                    }
328                }
329            }
330        }
331        Ok((world, mappings))
332    }
333
334    pub fn entities(&self) -> impl Iterator<Item = Entity> + '_ {
335        self.archetypes
336            .iter()
337            .flat_map(|archetype| archetype.entities.iter().copied())
338    }
339
340    pub fn rows(&self) -> impl Iterator<Item = PrefabRow> {
341        self.archetypes.iter().flat_map(|archetype| {
342            archetype
343                .entities
344                .iter()
345                .copied()
346                .enumerate()
347                .map(|(index, entity)| PrefabRow {
348                    entity,
349                    components: archetype
350                        .columns
351                        .iter()
352                        .map(|column| PrefabComponent {
353                            type_name: &column.type_name,
354                            module_name: column.module_name.as_deref(),
355                            data: &column.components[index],
356                        })
357                        .collect::<Vec<_>>(),
358                })
359        })
360    }
361}
362
363pub struct PrefabRow<'a> {
364    pub entity: Entity,
365    pub components: Vec<PrefabComponent<'a>>,
366}
367
368pub struct PrefabComponent<'a> {
369    pub type_name: &'a str,
370    pub module_name: Option<&'a str>,
371    pub data: &'a Intermediate,
372}
373
374#[cfg(test)]
375mod tests {
376    use super::*;
377
378    #[test]
379    fn test_prefab() {
380        let mut registry = Registry::default().with_basic_types();
381        Relation::<()>::install_to_registry(&mut registry);
382
383        let mut serialization = SerializationRegistry::default().with_basic_types();
384        Prefab::register_relation_serializer::<()>(&mut serialization);
385
386        let mut processor = WorldProcessor::default();
387        Relation::<()>::register_to_processor(&mut processor);
388
389        let mut world = World::default();
390        let a = world.spawn((42usize,)).unwrap();
391        let b = world.spawn((false, Relation::new((), a))).unwrap();
392        let c = world.spawn((true, Relation::new((), b))).unwrap();
393
394        {
395            let prefab = Prefab::from_world::<true>(&world, &serialization, &registry).unwrap();
396            let (world2, _) = prefab
397                .to_world::<true>(&processor, &serialization, &registry, (4.2f32,))
398                .unwrap();
399
400            let mut entities = world.entities().collect::<Vec<_>>();
401            let mut entities2 = world2.entities().collect::<Vec<_>>();
402            entities.sort();
403            entities2.sort();
404            assert_eq!(entities, entities2);
405
406            assert_eq!(*world2.component::<true, f32>(a).unwrap(), 4.2);
407            assert_eq!(*world2.component::<true, f32>(b).unwrap(), 4.2);
408            assert_eq!(*world2.component::<true, f32>(c).unwrap(), 4.2);
409
410            let old = world.component::<true, usize>(a).unwrap();
411            let new = world2.component::<true, usize>(a).unwrap();
412            assert_eq!(*old, *new);
413            let old = world.component::<true, bool>(b).unwrap();
414            let new = world2.component::<true, bool>(b).unwrap();
415            assert_eq!(*old, *new);
416            let old = world.component::<true, bool>(c).unwrap();
417            let new = world2.component::<true, bool>(c).unwrap();
418            assert_eq!(*old, *new);
419
420            assert!(world2.has_relation::<true, ()>(c, b));
421            assert!(world2.has_relation::<true, ()>(b, a));
422        }
423
424        {
425            let prefab =
426                Prefab::from_entities::<true>(&world, [c], &processor, &serialization, &registry)
427                    .unwrap();
428            let (world2, mappings) = prefab
429                .to_world::<true>(&processor, &serialization, &registry, ())
430                .unwrap();
431
432            let entities = world2.entities().collect::<Vec<_>>();
433            assert_eq!(entities.len(), 3);
434
435            let mappings = WorldProcessorEntityMapping::new(&mappings);
436            let a2 = mappings.remap(a);
437            let b2 = mappings.remap(b);
438            let c2 = mappings.remap(c);
439
440            let old = world.component::<true, usize>(a).unwrap();
441            let new = world2.component::<true, usize>(a2).unwrap();
442            assert_eq!(*old, *new);
443            let old = world.component::<true, bool>(b).unwrap();
444            let new = world2.component::<true, bool>(b2).unwrap();
445            assert_eq!(*old, *new);
446            let old = world.component::<true, bool>(c).unwrap();
447            let new = world2.component::<true, bool>(c2).unwrap();
448            assert_eq!(*old, *new);
449
450            assert!(world2.has_relation::<true, ()>(c2, b2));
451            assert!(world2.has_relation::<true, ()>(b2, a2));
452        }
453    }
454}