misty_vm/
resources.rs

1use std::{
2    collections::HashMap,
3    fmt::Debug,
4    ops::Deref,
5    sync::{atomic::AtomicU64, Arc, RwLock, Weak},
6};
7
8use serde::{Deserialize, Serialize};
9
10pub enum ResourceUpdateAction {
11    Insert(MistyResourceId, Vec<u8>),
12    Remove(MistyResourceId),
13}
14
15#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
16pub struct MistyResourceId(u64);
17
18impl MistyResourceId {
19    pub fn wrap(id: u64) -> Self {
20        Self(id)
21    }
22    pub fn invalid() -> Self {
23        Self(0)
24    }
25}
26
27impl Default for MistyResourceId {
28    fn default() -> Self {
29        Self::invalid()
30    }
31}
32
33const _: () = {
34    static ALLOCATED: AtomicU64 = AtomicU64::new(1);
35    impl MistyResourceId {
36        pub fn alloc() -> Self {
37            let id = ALLOCATED.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
38            return Self(id);
39        }
40    }
41    impl Deref for MistyResourceId {
42        type Target = u64;
43
44        fn deref(&self) -> &Self::Target {
45            &self.0
46        }
47    }
48};
49
50#[derive(Debug)]
51enum ToFlushResourceAction {
52    Insert(MistyResourceHandle),
53    Remove,
54}
55
56impl PartialEq for ToFlushResourceAction {
57    fn eq(&self, other: &Self) -> bool {
58        match (self, other) {
59            (Self::Remove, Self::Remove) => true,
60            (Self::Insert(l0), Self::Insert(r0)) => l0.id() == r0.id(),
61            _ => false,
62        }
63    }
64}
65
66struct MistyResourceManagerStore {
67    pending_actions: Arc<RwLock<Option<HashMap<MistyResourceId, ToFlushResourceAction>>>>,
68    weak_map: RwLock<HashMap<MistyResourceId, Weak<MistyResourceHandleInner>>>,
69}
70
71pub struct MistyResourceManager {
72    store: Arc<MistyResourceManagerStore>,
73}
74
75impl Debug for MistyResourceManager {
76    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
77        f.debug_struct("ResourceManager").finish()
78    }
79}
80
81#[derive(Debug, Clone)]
82struct MistyResourceHandleInner {
83    id: MistyResourceId,
84    store_ref: Weak<MistyResourceManagerStore>,
85    buf: Vec<u8>,
86}
87
88#[derive(Debug, Clone)]
89pub struct MistyResourceHandle {
90    ptr: Arc<MistyResourceHandleInner>,
91}
92
93impl MistyResourceHandle {
94    pub fn id(&self) -> MistyResourceId {
95        self.ptr.id
96    }
97
98    pub fn load(&self) -> &Vec<u8> {
99        &self.ptr.buf
100    }
101}
102
103impl Drop for MistyResourceHandleInner {
104    fn drop(&mut self) {
105        let store_ref = self.store_ref.upgrade();
106        if store_ref.is_none() {
107            return;
108        }
109        let store_ref = store_ref.unwrap();
110
111        {
112            let mut writter = store_ref.pending_actions.write().unwrap();
113            let writer = writter.as_mut().unwrap();
114            if writer.contains_key(&self.id) {
115                debug_assert!(writer.get(&self.id).unwrap() != &ToFlushResourceAction::Remove);
116                writer.remove(&self.id);
117            } else {
118                writer.insert(self.id, ToFlushResourceAction::Remove);
119            }
120        }
121        {
122            let mut writter = store_ref.weak_map.write().unwrap();
123            writter.remove(&self.id);
124        }
125    }
126}
127
128impl MistyResourceManager {
129    pub fn new() -> Self {
130        Self {
131            store: Arc::new(MistyResourceManagerStore {
132                pending_actions: Arc::new(RwLock::new(Some(Default::default()))),
133                weak_map: Default::default(),
134            }),
135        }
136    }
137
138    pub fn get_handle(&self, id: MistyResourceId) -> Option<MistyResourceHandle> {
139        let reader = self.store.weak_map.read().unwrap();
140        reader
141            .get(&id)
142            .map(|ptr| ptr.upgrade())
143            .unwrap_or_default()
144            .map(|ptr| MistyResourceHandle { ptr })
145    }
146
147    pub fn insert(&self, buf: impl Into<Vec<u8>> + 'static) -> MistyResourceHandle {
148        let id = MistyResourceId::alloc();
149        let buf: Vec<u8> = buf.into();
150
151        let ptr = Arc::new(MistyResourceHandleInner {
152            id,
153            store_ref: Arc::downgrade(&self.store),
154            buf,
155        });
156        let handle = MistyResourceHandle { ptr: ptr.clone() };
157
158        {
159            let mut writter = self.store.pending_actions.write().unwrap();
160            let writer = writter.as_mut().unwrap();
161            writer.insert(id, ToFlushResourceAction::Insert(handle.clone()));
162        }
163        {
164            let mut writter = self.store.weak_map.write().unwrap();
165            writter.insert(id, Arc::downgrade(&ptr));
166        }
167
168        handle
169    }
170
171    pub(crate) fn take_all_actions(&self) -> Vec<ResourceUpdateAction> {
172        let pending_ids = {
173            let mut writer = self.store.pending_actions.write().unwrap();
174            let cloned = writer.take();
175            *writer = Some(Default::default());
176            cloned.unwrap()
177        };
178
179        let mut ret: Vec<ResourceUpdateAction> = Default::default();
180        for (id, action) in pending_ids {
181            match action {
182                ToFlushResourceAction::Insert(handle) => {
183                    ret.push(ResourceUpdateAction::Insert(id, handle.load().clone()));
184                }
185                ToFlushResourceAction::Remove => {
186                    ret.push(ResourceUpdateAction::Remove(id));
187                }
188            }
189        }
190        ret
191    }
192}