Skip to main content

runmat_vm/runtime/
workspace.rs

1use runmat_builtins::Value;
2use runmat_thread_local::runmat_thread_local;
3use std::cell::RefCell;
4use std::collections::{HashMap, HashSet};
5
6struct WorkspaceState {
7    names: HashMap<String, usize>,
8    assigned: HashSet<String>,
9    assigned_names_this_execution: HashSet<String>,
10    assigned_ids_this_execution: HashSet<usize>,
11    removed_this_execution: HashMap<String, usize>,
12    idx_to_name: HashMap<usize, String>,
13    data_ptr: *const Value,
14    len: usize,
15}
16
17pub type WorkspaceSnapshot = (HashMap<String, usize>, HashSet<String>);
18
19#[derive(Debug, Clone)]
20pub struct WorkspaceAssignedReport {
21    pub ids: HashSet<usize>,
22    pub names: HashSet<String>,
23    pub removed_ids: HashSet<usize>,
24    pub removed_names: HashSet<String>,
25}
26
27runmat_thread_local! {
28    static WORKSPACE_STATE: RefCell<Option<WorkspaceState>> = const { RefCell::new(None) };
29    static PENDING_WORKSPACE: RefCell<Option<WorkspaceSnapshot>> = const { RefCell::new(None) };
30    static LAST_WORKSPACE_STATE: RefCell<Option<WorkspaceSnapshot>> = const { RefCell::new(None) };
31    static LAST_WORKSPACE_ASSIGNED_REPORT: RefCell<Option<WorkspaceAssignedReport>> = const { RefCell::new(None) };
32    static WORKSPACE_VARS: RefCell<Option<*mut Vec<Value>>> = const { RefCell::new(None) };
33}
34
35pub struct WorkspaceStateGuard;
36
37impl Drop for WorkspaceStateGuard {
38    fn drop(&mut self) {
39        WORKSPACE_STATE.with(|state| {
40            let mut state_mut = state.borrow_mut();
41            if let Some(ws) = state_mut.take() {
42                let removed_ids = ws.removed_this_execution.values().copied().collect();
43                let removed_names = ws.removed_this_execution.keys().cloned().collect();
44                LAST_WORKSPACE_ASSIGNED_REPORT.with(|slot| {
45                    *slot.borrow_mut() = Some(WorkspaceAssignedReport {
46                        ids: ws.assigned_ids_this_execution,
47                        names: ws.assigned_names_this_execution,
48                        removed_ids,
49                        removed_names,
50                    });
51                });
52                LAST_WORKSPACE_STATE.with(|slot| {
53                    *slot.borrow_mut() = Some((ws.names, ws.assigned));
54                });
55            }
56        });
57        WORKSPACE_VARS.with(|slot| {
58            slot.borrow_mut().take();
59        });
60    }
61}
62
63pub struct PendingWorkspaceGuard;
64
65impl Drop for PendingWorkspaceGuard {
66    fn drop(&mut self) {
67        PENDING_WORKSPACE.with(|slot| {
68            slot.borrow_mut().take();
69        });
70    }
71}
72
73pub fn push_pending_workspace(
74    names: HashMap<String, usize>,
75    assigned: HashSet<String>,
76) -> PendingWorkspaceGuard {
77    PENDING_WORKSPACE.with(|slot| {
78        *slot.borrow_mut() = Some((names, assigned));
79    });
80    PendingWorkspaceGuard
81}
82
83pub fn take_pending_workspace_state() -> Option<WorkspaceSnapshot> {
84    PENDING_WORKSPACE.with(|slot| slot.borrow_mut().take())
85}
86
87pub fn clone_pending_workspace_state() -> Option<WorkspaceSnapshot> {
88    PENDING_WORKSPACE.with(|slot| slot.borrow().clone())
89}
90
91pub fn restore_pending_workspace_state(snapshot: WorkspaceSnapshot) {
92    PENDING_WORKSPACE.with(|slot| {
93        *slot.borrow_mut() = Some(snapshot);
94    });
95}
96
97pub fn take_updated_workspace_state() -> Option<WorkspaceSnapshot> {
98    LAST_WORKSPACE_STATE.with(|slot| slot.borrow_mut().take())
99}
100
101pub fn take_updated_workspace_assigned_report() -> Option<WorkspaceAssignedReport> {
102    LAST_WORKSPACE_ASSIGNED_REPORT.with(|slot| slot.borrow_mut().take())
103}
104
105pub fn set_workspace_state(
106    names: HashMap<String, usize>,
107    assigned: HashSet<String>,
108    vars: &mut Vec<Value>,
109) -> WorkspaceStateGuard {
110    let idx_to_name: HashMap<usize, String> = names.iter().map(|(k, &v)| (v, k.clone())).collect();
111    WORKSPACE_STATE.with(|state| {
112        *state.borrow_mut() = Some(WorkspaceState {
113            names,
114            assigned,
115            assigned_names_this_execution: HashSet::new(),
116            assigned_ids_this_execution: HashSet::new(),
117            removed_this_execution: HashMap::new(),
118            idx_to_name,
119            data_ptr: vars.as_ptr(),
120            len: vars.len(),
121        });
122    });
123    let vars_ptr = vars as *mut Vec<Value>;
124    WORKSPACE_VARS.with(|slot| {
125        *slot.borrow_mut() = Some(vars_ptr);
126    });
127    WorkspaceStateGuard
128}
129
130pub fn refresh_workspace_state(vars: &[Value]) {
131    WORKSPACE_STATE.with(|state| {
132        if let Some(ws) = state.borrow_mut().as_mut() {
133            ws.data_ptr = vars.as_ptr();
134            ws.len = vars.len();
135        }
136    });
137}
138
139pub fn workspace_lookup(name: &str) -> Option<Value> {
140    WORKSPACE_STATE.with(|state| {
141        let state_ref = state.borrow();
142        let ws = state_ref.as_ref()?;
143        let idx = ws.names.get(name)?;
144        if !ws.assigned.contains(name) {
145            return None;
146        }
147        if *idx >= ws.len {
148            return None;
149        }
150        unsafe {
151            let ptr = ws.data_ptr.add(*idx);
152            Some((*ptr).clone())
153        }
154    })
155}
156
157pub fn workspace_slot_assigned(index: usize) -> Option<bool> {
158    WORKSPACE_STATE.with(|state| {
159        let state_ref = state.borrow();
160        let ws = state_ref.as_ref()?;
161        let name = ws.idx_to_name.get(&index)?;
162        Some(ws.assigned.contains(name))
163    })
164}
165
166pub fn workspace_state_available() -> bool {
167    WORKSPACE_STATE.with(|state| state.borrow().is_some())
168}
169
170pub fn workspace_assign(name: &str, value: Value) -> Result<(), String> {
171    let vars_ptr = WORKSPACE_VARS.with(|slot| *slot.borrow());
172    let Some(vars_ptr) = vars_ptr else {
173        return Err("load: workspace state unavailable".to_string());
174    };
175    let vars = unsafe { &mut *vars_ptr };
176    set_workspace_variable(name, value, vars)
177}
178
179pub fn workspace_clear() -> Result<(), String> {
180    let vars_ptr = WORKSPACE_VARS.with(|slot| *slot.borrow());
181    let Some(vars_ptr) = vars_ptr else {
182        return Err("clear: workspace state unavailable".to_string());
183    };
184    let vars = unsafe { &mut *vars_ptr };
185
186    WORKSPACE_STATE.with(|state| {
187        let mut state_mut = state.borrow_mut();
188        let Some(ws) = state_mut.as_mut() else {
189            return Err("clear: workspace state unavailable".to_string());
190        };
191        vars.clear();
192        for (name, idx) in &ws.names {
193            if ws.assigned.contains(name) {
194                ws.removed_this_execution.insert(name.clone(), *idx);
195            }
196        }
197        ws.names.clear();
198        ws.assigned.clear();
199        ws.idx_to_name.clear();
200        ws.data_ptr = vars.as_ptr();
201        ws.len = vars.len();
202        Ok(())
203    })
204}
205
206pub fn workspace_remove(name: &str) -> Result<(), String> {
207    let vars_ptr = WORKSPACE_VARS.with(|slot| *slot.borrow());
208    let Some(vars_ptr) = vars_ptr else {
209        return Err("clear: workspace state unavailable".to_string());
210    };
211    let vars = unsafe { &mut *vars_ptr };
212
213    WORKSPACE_STATE.with(|state| {
214        let mut state_mut = state.borrow_mut();
215        let Some(ws) = state_mut.as_mut() else {
216            return Err("clear: workspace state unavailable".to_string());
217        };
218        if let Some(idx) = ws.names.remove(name) {
219            if idx < vars.len() {
220                vars[idx] = Value::Num(0.0);
221            }
222            if ws.assigned.contains(name) {
223                ws.removed_this_execution.insert(name.to_string(), idx);
224            }
225            ws.assigned.remove(name);
226            ws.idx_to_name.remove(&idx);
227            ws.data_ptr = vars.as_ptr();
228            ws.len = vars.len();
229        }
230        Ok(())
231    })
232}
233
234pub fn workspace_snapshot() -> Vec<(String, Value)> {
235    WORKSPACE_STATE.with(|state| {
236        if let Some(ws) = state.borrow().as_ref() {
237            let mut entries: Vec<(String, Value)> = ws
238                .names
239                .iter()
240                .filter_map(|(name, idx)| {
241                    if *idx >= ws.len {
242                        return None;
243                    }
244                    if !ws.assigned.contains(name) {
245                        return None;
246                    }
247                    unsafe {
248                        let ptr = ws.data_ptr.add(*idx);
249                        Some((name.clone(), (*ptr).clone()))
250                    }
251                })
252                .collect();
253            entries.sort_by(|a, b| a.0.cmp(&b.0));
254            entries
255        } else {
256            Vec::new()
257        }
258    })
259}
260
261pub fn set_workspace_variable(
262    name: &str,
263    value: Value,
264    vars: &mut Vec<Value>,
265) -> Result<(), String> {
266    let mut result = Ok(());
267    WORKSPACE_STATE.with(|state| {
268        let mut state_mut = state.borrow_mut();
269        match state_mut.as_mut() {
270            Some(ws) => {
271                let idx = if let Some(idx) = ws.names.get(name).copied() {
272                    idx
273                } else {
274                    let idx = vars.len();
275                    ws.names.insert(name.to_string(), idx);
276                    ws.idx_to_name.insert(idx, name.to_string());
277                    idx
278                };
279                if idx >= vars.len() {
280                    vars.resize(idx + 1, Value::Num(0.0));
281                }
282                vars[idx] = value;
283                ws.data_ptr = vars.as_ptr();
284                ws.len = vars.len();
285                ws.assigned.insert(name.to_string());
286                ws.assigned_names_this_execution.insert(name.to_string());
287                ws.assigned_ids_this_execution.insert(idx);
288                ws.removed_this_execution.remove(name);
289            }
290            None => {
291                result = Err("load: workspace state unavailable".to_string());
292            }
293        }
294    });
295    result
296}
297
298pub fn ensure_workspace_slot_name(index: usize, name: &str) {
299    WORKSPACE_STATE.with(|state| {
300        if let Some(ws) = state.borrow_mut().as_mut() {
301            ws.names.entry(name.to_string()).or_insert(index);
302            ws.idx_to_name
303                .entry(index)
304                .or_insert_with(|| name.to_string());
305        }
306    });
307}
308
309pub fn mark_workspace_assigned(index: usize) {
310    WORKSPACE_STATE.with(|state| {
311        if let Some(ws) = state.borrow_mut().as_mut() {
312            if let Some(name) = ws.idx_to_name.get(&index).cloned() {
313                ws.assigned.insert(name.clone());
314                ws.assigned_names_this_execution.insert(name.clone());
315                ws.assigned_ids_this_execution.insert(index);
316                ws.removed_this_execution.remove(&name);
317            }
318        }
319    });
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325
326    fn take_report_after(f: impl FnOnce(&mut Vec<Value>)) -> WorkspaceAssignedReport {
327        let _ = take_updated_workspace_assigned_report();
328        let _ = take_updated_workspace_state();
329
330        let mut vars = Vec::new();
331        {
332            let _guard = set_workspace_state(HashMap::new(), HashSet::new(), &mut vars);
333            f(&mut vars);
334        }
335
336        take_updated_workspace_assigned_report().expect("workspace report should be recorded")
337    }
338
339    #[test]
340    fn remove_preserves_assignment_report_and_records_removal() {
341        let report = take_report_after(|vars| {
342            set_workspace_variable("x", Value::Num(1.0), vars).unwrap();
343            workspace_remove("x").unwrap();
344        });
345
346        assert!(report.names.contains("x"));
347        assert!(report.ids.contains(&0));
348        assert!(report.removed_names.contains("x"));
349        assert!(report.removed_ids.contains(&0));
350    }
351
352    #[test]
353    fn clear_preserves_assignment_report_and_records_removal() {
354        let report = take_report_after(|vars| {
355            set_workspace_variable("x", Value::Num(1.0), vars).unwrap();
356            workspace_clear().unwrap();
357        });
358
359        assert!(report.names.contains("x"));
360        assert!(report.ids.contains(&0));
361        assert!(report.removed_names.contains("x"));
362        assert!(report.removed_ids.contains(&0));
363    }
364
365    #[test]
366    fn assignment_after_clear_clears_final_removal_marker() {
367        let report = take_report_after(|vars| {
368            set_workspace_variable("x", Value::Num(1.0), vars).unwrap();
369            workspace_clear().unwrap();
370            set_workspace_variable("x", Value::Num(2.0), vars).unwrap();
371        });
372
373        assert!(report.names.contains("x"));
374        assert!(report.removed_names.is_empty());
375        assert!(report.removed_ids.is_empty());
376    }
377}