anchor_yard_core/
lib.rs

1use std::{cell::RefCell, collections::HashMap, sync::Mutex};
2
3use once_cell::sync::Lazy;
4use serde::{Deserialize, Serialize};
5use shipyard::{AllStoragesViewMut, Component, EntityId, IntoIter, View, World};
6
7type Serializer = Box<dyn Fn(&World) -> HashMap<EntityId, Vec<u8>> + Send + Sync>;
8type Deserializer = Box<dyn Fn(&mut World, &HashMap<EntityId, Vec<u8>>) + Send + Sync>;
9
10#[derive(Default)]
11pub struct SnapshotRegistry {
12    components: HashMap<String, (Serializer, Deserializer)>,
13}
14
15impl SnapshotRegistry {
16    /// 스냅샷에 포함할 컴포넌트를 등록합니다.
17    pub fn register<T: Component + Serialize + for<'de> Deserialize<'de> + 'static>(&mut self) {
18        let type_name = std::any::type_name::<T>().to_string();
19
20        let serializer: Serializer = Box::new(|world| {
21            let mut data = HashMap::new();
22            if let Ok(view) = world.borrow::<View<T>>() {
23                for (id, component) in view.iter().with_id() {
24                    if let Ok(bytes) =
25                        bincode::serde::encode_to_vec(component, bincode::config::standard())
26                    {
27                        data.insert(id, bytes);
28                    }
29                }
30            }
31            data
32        });
33
34        let deserializer: Deserializer = Box::new(|world, data| {
35            let mut all_storages = world.borrow::<AllStoragesViewMut>().unwrap();
36            for bytes in data.values() {
37                if let Ok((c, _)) =
38                    bincode::serde::decode_from_slice::<T, _>(bytes, bincode::config::standard())
39                {
40                    all_storages.add_entity((c,));
41                }
42            }
43        });
44
45        self.components
46            .insert(type_name, (serializer, deserializer));
47    }
48}
49
50/// 전역 레지스트리. `main` 함수에서 컴포넌트를 등록할 때 사용합니다.
51pub static REGISTRY: Lazy<Mutex<SnapshotRegistry>> =
52    Lazy::new(|| Mutex::new(SnapshotRegistry::default()));
53
54thread_local! {
55    static CURRENT_WORLD: RefCell<Option<*const World>> = const { RefCell::new(None) };
56}
57
58#[derive(Serialize, Deserialize, Clone)]
59pub struct SystemSnapshot {
60    pub system_name: String,
61    pub execution_time_ms: u64,
62    pub timestamp: u64,
63    pub component_data: HashMap<String, HashMap<EntityId, Vec<u8>>>,
64}
65
66impl SystemSnapshot {
67    pub fn capture_world(world: &World, system_name: &str, execution_time_ms: u64) -> Option<Self> {
68        let mut component_data = HashMap::new();
69        let registry = REGISTRY.lock().unwrap();
70
71        for (name, (serializer, _)) in &registry.components {
72            component_data.insert(name.clone(), serializer(world));
73        }
74
75        Some(SystemSnapshot {
76            system_name: system_name.to_string(),
77            execution_time_ms,
78            timestamp: std::time::SystemTime::now()
79                .duration_since(std::time::UNIX_EPOCH)
80                .unwrap()
81                .as_secs(),
82            component_data,
83        })
84    }
85
86    pub fn save_to_file(&self) -> Result<(), anyhow::Error> {
87        std::fs::create_dir_all("snapshots")?;
88        let filename = format!("snapshots/{}_{}.snapshot", self.system_name, self.timestamp);
89        let data = bincode::serde::encode_to_vec(self, bincode::config::standard())?;
90        std::fs::write(filename, data)?;
91        Ok(())
92    }
93
94    pub fn load_from_file(path: &std::path::Path) -> Result<Self, anyhow::Error> {
95        let data = std::fs::read(path)?;
96        let (snapshot, _) = bincode::serde::decode_from_slice(&data, bincode::config::standard())?;
97        Ok(snapshot)
98    }
99
100    pub fn restore_world(&self) -> Result<World, anyhow::Error> {
101        let mut world = World::new();
102        let mut all_entities = std::collections::HashSet::new();
103        for component_map in self.component_data.values() {
104            for entity_id in component_map.keys() {
105                all_entities.insert(entity_id);
106            }
107        }
108
109        let registry = REGISTRY.lock().unwrap();
110        for (name, data) in &self.component_data {
111            if let Some((_, deserializer)) = registry.components.get(name) {
112                deserializer(&mut world, data);
113            }
114        }
115        Ok(world)
116    }
117}
118
119pub(crate) struct WorldGuard {
120    _private: (),
121}
122
123impl WorldGuard {
124    pub fn new(world: &World) -> Self {
125        CURRENT_WORLD.with(|current| {
126            *current.borrow_mut() = Some(world as *const World);
127        });
128        WorldGuard { _private: () }
129    }
130}
131
132impl Drop for WorldGuard {
133    fn drop(&mut self) {
134        CURRENT_WORLD.with(|current| {
135            *current.borrow_mut() = None;
136        });
137    }
138}
139
140pub fn with_current_world<F, R>(f: F) -> Option<R>
141where
142    F: FnOnce(&World) -> R,
143{
144    CURRENT_WORLD.with(|current| {
145        current.borrow().map(|ptr| {
146            let world = unsafe { &*ptr };
147            f(world)
148        })
149    })
150}
151
152pub trait WorldSnapshotExt {
153    fn run_default_workload_with_snapshot(&self) -> Result<(), shipyard::error::RunWorkload>;
154}
155
156impl WorldSnapshotExt for World {
157    fn run_default_workload_with_snapshot(&self) -> Result<(), shipyard::error::RunWorkload> {
158        let _guard = WorldGuard::new(self);
159        self.run_default_workload()
160    }
161}
anchor_yard_core/lib.rs

anchor_yard_core/
lib.rs
