1use runmat_builtins::Value;
2use runmat_thread_local::runmat_thread_local;
3use std::cell::RefCell;
4use std::collections::{HashMap, HashSet};
5
6#[derive(Debug, Clone)]
7enum SlotLifecycle {
8 Assigned(String),
9 Unassigned(String),
10}
11
12impl SlotLifecycle {
13 fn name(&self) -> &str {
14 match self {
15 SlotLifecycle::Assigned(name) | SlotLifecycle::Unassigned(name) => name,
16 }
17 }
18
19 fn is_assigned(&self) -> bool {
20 matches!(self, SlotLifecycle::Assigned(_))
21 }
22}
23
24struct WorkspaceState {
25 names: HashMap<String, usize>,
26 assigned: HashSet<String>,
27 assigned_names_this_execution: HashSet<String>,
28 assigned_ids_this_execution: HashSet<usize>,
29 removed_slots_this_execution: HashMap<usize, String>,
30 slot_lifecycle: HashMap<usize, SlotLifecycle>,
31 data_ptr: *const Value,
32 len: usize,
33}
34
35struct WorkspaceFrame {
36 state: WorkspaceState,
37 vars_ptr: *mut Vec<Value>,
38 vars_snapshot: Vec<Value>,
39 publish_on_drop: bool,
40}
41
42pub type WorkspaceSnapshot = (HashMap<String, usize>, HashSet<String>);
43
44#[derive(Debug, Clone)]
45pub struct WorkspaceValueSnapshot {
46 pub vars: Vec<Value>,
47 pub names: HashMap<String, usize>,
48 pub assigned: HashSet<String>,
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub enum WorkspaceTarget {
53 Current,
54 Caller,
55 Base,
56}
57
58#[derive(Debug, Clone)]
59pub struct WorkspaceTargetSnapshot {
60 pub names: HashMap<String, usize>,
61 pub assigned: HashSet<String>,
62 pub vars_ptr: *mut Vec<Value>,
63}
64
65#[derive(Debug, Clone)]
66pub struct WorkspaceAssignedReport {
67 pub ids: HashSet<usize>,
68 pub names: HashSet<String>,
69 pub removed_ids: HashSet<usize>,
70 pub removed_names: HashSet<String>,
71}
72
73runmat_thread_local! {
74 static WORKSPACE_STACK: RefCell<Vec<WorkspaceFrame>> = const { RefCell::new(Vec::new()) };
75 static PENDING_WORKSPACE: RefCell<Option<WorkspaceSnapshot>> = const { RefCell::new(None) };
76 static LAST_WORKSPACE_STATE: RefCell<Option<WorkspaceValueSnapshot>> = const { RefCell::new(None) };
77 static LAST_WORKSPACE_ASSIGNED_REPORT: RefCell<Option<WorkspaceAssignedReport>> = const { RefCell::new(None) };
78}
79
80fn mark_slot_unassigned(ws: &mut WorkspaceState, index: usize, name: String) {
81 ws.slot_lifecycle
82 .insert(index, SlotLifecycle::Unassigned(name.clone()));
83 ws.removed_slots_this_execution.insert(index, name);
84}
85
86fn mark_slot_assigned(ws: &mut WorkspaceState, index: usize, name: String) {
87 ws.slot_lifecycle
88 .insert(index, SlotLifecycle::Assigned(name.clone()));
89 ws.removed_slots_this_execution.remove(&index);
90}
91
92fn find_unassigned_slot_for_name(ws: &WorkspaceState, name: &str) -> Option<usize> {
93 ws.slot_lifecycle.iter().find_map(|(idx, state)| {
94 matches!(state, SlotLifecycle::Unassigned(slot_name) if slot_name == name).then_some(*idx)
95 })
96}
97
98fn upsert_slot_lifecycle_name(ws: &mut WorkspaceState, index: usize, name: &str) {
99 if let Some(existing_index) = ws.names.insert(name.to_string(), index) {
100 if existing_index != index {
101 mark_slot_unassigned(ws, existing_index, name.to_string());
102 ws.assigned.remove(name);
103 }
104 }
105
106 match ws.slot_lifecycle.get_mut(&index) {
107 Some(state) => {
108 let was_assigned = state.is_assigned();
109 let old_name = state.name().to_string();
110 if old_name != name {
111 if ws.names.get(&old_name).copied() == Some(index) {
112 ws.names.remove(&old_name);
113 }
114 ws.assigned.remove(&old_name);
115 }
116 *state = if was_assigned {
117 SlotLifecycle::Assigned(name.to_string())
118 } else {
119 SlotLifecycle::Unassigned(name.to_string())
120 };
121 }
122 None => {
123 ws.slot_lifecycle
124 .insert(index, SlotLifecycle::Unassigned(name.to_string()));
125 }
126 }
127}
128
129fn target_frame_index(len: usize, target: WorkspaceTarget) -> Option<usize> {
130 if len == 0 {
131 return None;
132 }
133 match target {
134 WorkspaceTarget::Current => Some(len - 1),
135 WorkspaceTarget::Caller => Some(len.saturating_sub(2)),
136 WorkspaceTarget::Base => Some(0),
137 }
138}
139
140fn lifecycle_from_names(
141 names: &HashMap<String, usize>,
142 assigned: &HashSet<String>,
143) -> HashMap<usize, SlotLifecycle> {
144 names
145 .iter()
146 .map(|(name, idx)| {
147 let lifecycle = if assigned.contains(name) {
148 SlotLifecycle::Assigned(name.clone())
149 } else {
150 SlotLifecycle::Unassigned(name.clone())
151 };
152 (*idx, lifecycle)
153 })
154 .collect()
155}
156
157pub struct WorkspaceStateGuard;
158
159impl Drop for WorkspaceStateGuard {
160 fn drop(&mut self) {
161 WORKSPACE_STACK.with(|stack| {
162 let mut stack = stack.borrow_mut();
163 if let Some(frame) = stack.pop() {
164 if !frame.publish_on_drop {
165 return;
166 }
167 let ws = frame.state;
168 let removed_ids = ws.removed_slots_this_execution.keys().copied().collect();
169 let removed_names = ws.removed_slots_this_execution.values().cloned().collect();
170 LAST_WORKSPACE_ASSIGNED_REPORT.with(|slot| {
171 *slot.borrow_mut() = Some(WorkspaceAssignedReport {
172 ids: ws.assigned_ids_this_execution,
173 names: ws.assigned_names_this_execution,
174 removed_ids,
175 removed_names,
176 });
177 });
178 LAST_WORKSPACE_STATE.with(|slot| {
179 *slot.borrow_mut() = Some(WorkspaceValueSnapshot {
180 vars: frame.vars_snapshot,
181 names: ws.names,
182 assigned: ws.assigned,
183 });
184 });
185 }
186 });
187 }
188}
189
190pub struct PendingWorkspaceGuard;
191
192impl Drop for PendingWorkspaceGuard {
193 fn drop(&mut self) {
194 PENDING_WORKSPACE.with(|slot| {
195 slot.borrow_mut().take();
196 });
197 }
198}
199
200pub fn push_pending_workspace(
201 names: HashMap<String, usize>,
202 assigned: HashSet<String>,
203) -> PendingWorkspaceGuard {
204 PENDING_WORKSPACE.with(|slot| {
205 *slot.borrow_mut() = Some((names, assigned));
206 });
207 PendingWorkspaceGuard
208}
209
210pub fn take_pending_workspace_state() -> Option<WorkspaceSnapshot> {
211 PENDING_WORKSPACE.with(|slot| slot.borrow_mut().take())
212}
213
214pub fn take_updated_workspace_state() -> Option<WorkspaceValueSnapshot> {
215 LAST_WORKSPACE_STATE.with(|slot| slot.borrow_mut().take())
216}
217
218pub fn take_updated_workspace_assigned_report() -> Option<WorkspaceAssignedReport> {
219 LAST_WORKSPACE_ASSIGNED_REPORT.with(|slot| slot.borrow_mut().take())
220}
221
222pub fn set_workspace_state(
223 names: HashMap<String, usize>,
224 assigned: HashSet<String>,
225 vars: &mut Vec<Value>,
226) -> WorkspaceStateGuard {
227 set_workspace_state_with_publish(names, assigned, vars, true)
228}
229
230pub fn set_transient_workspace_state(
231 names: HashMap<String, usize>,
232 assigned: HashSet<String>,
233 vars: &mut Vec<Value>,
234) -> WorkspaceStateGuard {
235 set_workspace_state_with_publish(names, assigned, vars, false)
236}
237
238fn set_workspace_state_with_publish(
239 names: HashMap<String, usize>,
240 assigned: HashSet<String>,
241 vars: &mut Vec<Value>,
242 publish_on_drop: bool,
243) -> WorkspaceStateGuard {
244 let mut slot_lifecycle = HashMap::new();
245 for (name, idx) in &names {
246 let lifecycle = if assigned.contains(name) {
247 SlotLifecycle::Assigned(name.clone())
248 } else {
249 SlotLifecycle::Unassigned(name.clone())
250 };
251 slot_lifecycle.insert(*idx, lifecycle);
252 }
253 let vars_ptr = vars as *mut Vec<Value>;
254 WORKSPACE_STACK.with(|stack| {
255 stack.borrow_mut().push(WorkspaceFrame {
256 state: WorkspaceState {
257 names,
258 assigned,
259 assigned_names_this_execution: HashSet::new(),
260 assigned_ids_this_execution: HashSet::new(),
261 removed_slots_this_execution: HashMap::new(),
262 slot_lifecycle,
263 data_ptr: vars.as_ptr(),
264 len: vars.len(),
265 },
266 vars_ptr,
267 vars_snapshot: vars.clone(),
268 publish_on_drop,
269 });
270 });
271 WorkspaceStateGuard
272}
273
274pub fn refresh_workspace_state(vars: &[Value]) {
275 WORKSPACE_STACK.with(|stack| {
276 if let Some(frame) = stack.borrow_mut().last_mut() {
277 frame.state.data_ptr = vars.as_ptr();
278 frame.state.len = vars.len();
279 frame.vars_snapshot = vars.to_vec();
280 }
281 });
282}
283
284pub fn workspace_lookup(name: &str) -> Option<Value> {
285 WORKSPACE_STACK.with(|stack| {
286 let stack = stack.borrow();
287 let frame = stack.last()?;
288 let ws = &frame.state;
289 let idx = ws.names.get(name)?;
290 if !ws.assigned.contains(name) {
291 return None;
292 }
293 if *idx >= ws.len {
294 return None;
295 }
296 unsafe {
297 let ptr = ws.data_ptr.add(*idx);
298 Some((*ptr).clone())
299 }
300 })
301}
302
303pub fn workspace_slot_assigned(index: usize) -> Option<bool> {
304 WORKSPACE_STACK.with(|stack| {
305 let stack = stack.borrow();
306 let frame = stack.last()?;
307 let ws = &frame.state;
308 ws.slot_lifecycle
309 .get(&index)
310 .map(SlotLifecycle::is_assigned)
311 })
312}
313
314pub fn workspace_slot_name(index: usize) -> Option<String> {
315 WORKSPACE_STACK.with(|stack| {
316 let stack = stack.borrow();
317 let frame = stack.last()?;
318 let ws = &frame.state;
319 ws.slot_lifecycle
320 .get(&index)
321 .map(|state| state.name().to_string())
322 })
323}
324
325pub fn workspace_assign(name: &str, value: Value) -> Result<(), String> {
326 workspace_assign_target(WorkspaceTarget::Current, name, value)
327}
328
329pub fn workspace_assign_target(
330 target: WorkspaceTarget,
331 name: &str,
332 value: Value,
333) -> Result<(), String> {
334 WORKSPACE_STACK.with(|stack| {
335 let mut stack = stack.borrow_mut();
336 let index = target_frame_index(stack.len(), target)
337 .ok_or_else(|| "load: workspace state unavailable".to_string())?;
338 let frame = stack
339 .get_mut(index)
340 .ok_or_else(|| "load: workspace state unavailable".to_string())?;
341 set_workspace_variable_in_frame(frame, name, value)
342 })
343}
344
345pub fn workspace_clear() -> Result<(), String> {
346 WORKSPACE_STACK.with(|stack| {
347 let mut stack = stack.borrow_mut();
348 let frame = stack
349 .last_mut()
350 .ok_or_else(|| "clear: workspace state unavailable".to_string())?;
351 let ws = &mut frame.state;
352 let vars = unsafe { &mut *frame.vars_ptr };
353 for (name, idx) in ws.names.clone() {
354 if let Some(value) = vars.get_mut(idx) {
355 *value = Value::Num(0.0);
356 }
357 mark_slot_unassigned(ws, idx, name);
358 }
359 ws.names.clear();
360 ws.assigned.clear();
361 frame.vars_snapshot = vars.clone();
362 ws.data_ptr = vars.as_ptr();
363 ws.len = vars.len();
364 Ok(())
365 })
366}
367
368pub fn workspace_remove(name: &str) -> Result<(), String> {
369 WORKSPACE_STACK.with(|stack| {
370 let mut stack = stack.borrow_mut();
371 let frame = stack
372 .last_mut()
373 .ok_or_else(|| "clear: workspace state unavailable".to_string())?;
374 let ws = &mut frame.state;
375 let vars = unsafe { &mut *frame.vars_ptr };
376 if let Some(idx) = ws.names.remove(name) {
377 if let Some(value) = vars.get_mut(idx) {
378 *value = Value::Num(0.0);
379 }
380 ws.assigned.remove(name);
381 mark_slot_unassigned(ws, idx, name.to_string());
382 frame.vars_snapshot = vars.clone();
383 ws.data_ptr = vars.as_ptr();
384 ws.len = vars.len();
385 }
386 Ok(())
387 })
388}
389
390pub fn workspace_snapshot() -> Vec<(String, Value)> {
391 WORKSPACE_STACK.with(|stack| {
392 let stack = stack.borrow();
393 if let Some(frame) = stack.last() {
394 let ws = &frame.state;
395 let mut entries: Vec<(String, Value)> = ws
396 .names
397 .iter()
398 .filter_map(|(name, idx)| {
399 if *idx >= ws.len {
400 return None;
401 }
402 if !ws.assigned.contains(name) {
403 return None;
404 }
405 unsafe {
406 let ptr = ws.data_ptr.add(*idx);
407 Some((name.clone(), (*ptr).clone()))
408 }
409 })
410 .collect();
411 entries.sort_by(|a, b| a.0.cmp(&b.0));
412 entries
413 } else {
414 Vec::new()
415 }
416 })
417}
418
419pub fn workspace_target_snapshot(
420 target: WorkspaceTarget,
421) -> Result<WorkspaceTargetSnapshot, String> {
422 WORKSPACE_STACK.with(|stack| {
423 let stack = stack.borrow();
424 let index = target_frame_index(stack.len(), target)
425 .ok_or_else(|| "workspace state unavailable".to_string())?;
426 let frame = stack
427 .get(index)
428 .ok_or_else(|| "workspace state unavailable".to_string())?;
429 Ok(WorkspaceTargetSnapshot {
430 names: frame.state.names.clone(),
431 assigned: frame.state.assigned.clone(),
432 vars_ptr: frame.vars_ptr,
433 })
434 })
435}
436
437pub fn replace_workspace_target_vars_and_state(
438 target: WorkspaceTarget,
439 vars: Vec<Value>,
440 names: HashMap<String, usize>,
441 assigned: HashSet<String>,
442) -> Result<(), String> {
443 WORKSPACE_STACK.with(|stack| {
444 let mut stack = stack.borrow_mut();
445 let index = target_frame_index(stack.len(), target)
446 .ok_or_else(|| "workspace state unavailable".to_string())?;
447 let frame = stack
448 .get_mut(index)
449 .ok_or_else(|| "workspace state unavailable".to_string())?;
450 let target_vars = unsafe { &mut *frame.vars_ptr };
451 *target_vars = vars;
452 frame.vars_snapshot = target_vars.clone();
453 frame.state.names = names;
454 frame.state.assigned = assigned;
455 frame.state.slot_lifecycle =
456 lifecycle_from_names(&frame.state.names, &frame.state.assigned);
457 frame.state.data_ptr = target_vars.as_ptr();
458 frame.state.len = target_vars.len();
459 Ok(())
460 })
461}
462
463#[cfg(test)]
464pub fn set_workspace_variable(
465 name: &str,
466 value: Value,
467 vars: &mut Vec<Value>,
468) -> Result<(), String> {
469 WORKSPACE_STACK.with(|stack| {
470 let mut stack = stack.borrow_mut();
471 match stack.last_mut() {
472 Some(frame) => set_workspace_variable_in_frame(frame, name, value),
473 None => Err("load: workspace state unavailable".to_string()),
474 }
475 })?;
476 let _ = vars;
477 Ok(())
478}
479
480fn set_workspace_variable_in_frame(
481 frame: &mut WorkspaceFrame,
482 name: &str,
483 value: Value,
484) -> Result<(), String> {
485 let vars = unsafe { &mut *frame.vars_ptr };
486 let ws = &mut frame.state;
487 let idx = if let Some(idx) = ws.names.get(name).copied() {
488 idx
489 } else if let Some(idx) = find_unassigned_slot_for_name(ws, name) {
490 ws.names.insert(name.to_string(), idx);
491 idx
492 } else {
493 let idx = vars.len();
494 ws.names.insert(name.to_string(), idx);
495 idx
496 };
497 if idx >= vars.len() {
498 vars.resize(idx + 1, Value::Num(0.0));
499 }
500 vars[idx] = value;
501 frame.vars_snapshot = vars.clone();
502 ws.data_ptr = vars.as_ptr();
503 ws.len = vars.len();
504 ws.assigned.insert(name.to_string());
505 ws.assigned_names_this_execution.insert(name.to_string());
506 ws.assigned_ids_this_execution.insert(idx);
507 mark_slot_assigned(ws, idx, name.to_string());
508 Ok(())
509}
510
511pub fn ensure_workspace_slot_name(index: usize, name: &str) {
512 WORKSPACE_STACK.with(|stack| {
513 if let Some(frame) = stack.borrow_mut().last_mut() {
514 upsert_slot_lifecycle_name(&mut frame.state, index, name);
515 }
516 });
517}
518
519pub fn mark_workspace_assigned(index: usize) {
520 WORKSPACE_STACK.with(|stack| {
521 if let Some(frame) = stack.borrow_mut().last_mut() {
522 let ws = &mut frame.state;
523 if let Some(name) = ws
524 .slot_lifecycle
525 .get(&index)
526 .map(|slot| slot.name().to_string())
527 {
528 ws.assigned.insert(name.clone());
529 ws.assigned_names_this_execution.insert(name.clone());
530 ws.assigned_ids_this_execution.insert(index);
531 mark_slot_assigned(ws, index, name);
532 let vars = unsafe { &*frame.vars_ptr };
533 frame.vars_snapshot = vars.clone();
534 }
535 }
536 });
537}
538
539pub(crate) fn reset_thread_state_for_tests() {
540 WORKSPACE_STACK.with(|stack| stack.borrow_mut().clear());
541 PENDING_WORKSPACE.with(|slot| {
542 slot.borrow_mut().take();
543 });
544 LAST_WORKSPACE_STATE.with(|slot| {
545 slot.borrow_mut().take();
546 });
547 LAST_WORKSPACE_ASSIGNED_REPORT.with(|slot| {
548 slot.borrow_mut().take();
549 });
550}
551
552#[cfg(test)]
553mod tests {
554 use super::*;
555
556 fn take_report_after(f: impl FnOnce(&mut Vec<Value>)) -> WorkspaceAssignedReport {
557 let _ = take_updated_workspace_assigned_report();
558 let _ = take_updated_workspace_state();
559
560 let mut vars = Vec::new();
561 {
562 let _guard = set_workspace_state(HashMap::new(), HashSet::new(), &mut vars);
563 f(&mut vars);
564 }
565
566 take_updated_workspace_assigned_report().expect("workspace report should be recorded")
567 }
568
569 #[test]
570 fn remove_preserves_assignment_report_and_records_removal() {
571 let report = take_report_after(|vars| {
572 set_workspace_variable("x", Value::Num(1.0), vars).unwrap();
573 workspace_remove("x").unwrap();
574 });
575
576 assert!(report.names.contains("x"));
577 assert!(report.ids.contains(&0));
578 assert!(report.removed_names.contains("x"));
579 assert!(report.removed_ids.contains(&0));
580 }
581
582 #[test]
583 fn clear_preserves_assignment_report_and_records_removal() {
584 let report = take_report_after(|vars| {
585 set_workspace_variable("x", Value::Num(1.0), vars).unwrap();
586 workspace_clear().unwrap();
587 });
588
589 assert!(report.names.contains("x"));
590 assert!(report.ids.contains(&0));
591 assert!(report.removed_names.contains("x"));
592 assert!(report.removed_ids.contains(&0));
593 }
594
595 #[test]
596 fn clear_preserves_active_frame_slots() {
597 let _ = take_updated_workspace_state();
598 let mut vars = vec![Value::Num(11.0), Value::Num(0.0), Value::Num(0.0)];
599 let mut names = HashMap::new();
600 names.insert("x".to_string(), 0);
601 let assigned = HashSet::from(["x".to_string()]);
602
603 {
604 let _guard = set_workspace_state(names, assigned, &mut vars);
605 workspace_clear().unwrap();
606
607 assert_eq!(vars.len(), 3);
608 assert_eq!(vars[0], Value::Num(0.0));
609 assert_eq!(vars[1], Value::Num(0.0));
610 assert_eq!(vars[2], Value::Num(0.0));
611 assert_eq!(workspace_lookup("x"), None);
612 assert_eq!(workspace_slot_assigned(0), Some(false));
613 }
614
615 let snapshot = take_updated_workspace_state().expect("workspace state should be recorded");
616 assert_eq!(snapshot.vars.len(), 3);
617 assert!(snapshot.names.is_empty());
618 assert!(snapshot.assigned.is_empty());
619 }
620
621 #[test]
622 fn assignment_after_clear_clears_final_removal_marker() {
623 let report = take_report_after(|vars| {
624 set_workspace_variable("x", Value::Num(1.0), vars).unwrap();
625 workspace_clear().unwrap();
626 set_workspace_variable("x", Value::Num(2.0), vars).unwrap();
627 });
628
629 assert!(report.names.contains("x"));
630 assert!(report.removed_names.is_empty());
631 assert!(report.removed_ids.is_empty());
632 }
633
634 #[test]
635 fn assignment_after_remove_reuses_previous_slot() {
636 let mut vars = Vec::new();
637 let _ = take_updated_workspace_state();
638 {
639 let _guard = set_workspace_state(HashMap::new(), HashSet::new(), &mut vars);
640 set_workspace_variable("x", Value::Num(1.0), &mut vars).unwrap();
641 set_workspace_variable("z", Value::Num(9.0), &mut vars).unwrap();
642 workspace_remove("x").unwrap();
643 set_workspace_variable("x", Value::Num(42.0), &mut vars).unwrap();
644 assert_eq!(workspace_lookup("x"), Some(Value::Num(42.0)));
645 assert_eq!(vars[0], Value::Num(42.0));
646 }
647
648 let snapshot = take_updated_workspace_state().expect("workspace state should be recorded");
649 assert_eq!(snapshot.names.get("x"), Some(&0));
650 assert_eq!(snapshot.names.get("z"), Some(&1));
651 assert!(snapshot.assigned.contains("x"));
652 assert!(snapshot.assigned.contains("z"));
653 assert_eq!(snapshot.vars[0], Value::Num(42.0));
654 }
655}