runmat_core/session/
workspace.rs

1use super::*;
2
3impl RunMatSession {
4    pub fn clear_variables(&mut self) {
5        self.variables.clear();
6        self.variable_array.clear();
7        self.variable_names.clear();
8        self.workspace_values.clear();
9        self.workspace_preview_tokens.clear();
10    }
11
12    pub async fn export_workspace_state(
13        &mut self,
14        mode: WorkspaceExportMode,
15    ) -> Result<Option<Vec<u8>>> {
16        if matches!(mode, WorkspaceExportMode::Off) {
17            return Ok(None);
18        }
19
20        let source_map = if self.workspace_values.is_empty() {
21            &self.variables
22        } else {
23            &self.workspace_values
24        };
25
26        let mut entries: Vec<(String, Value)> = Vec::with_capacity(source_map.len());
27        for (name, value) in source_map {
28            let gathered = gather_if_needed_async(value).await?;
29            entries.push((name.clone(), gathered));
30        }
31
32        if entries.is_empty() && matches!(mode, WorkspaceExportMode::Auto) {
33            return Ok(None);
34        }
35
36        let replay_mode = match mode {
37            WorkspaceExportMode::Auto => WorkspaceReplayMode::Auto,
38            WorkspaceExportMode::Force => WorkspaceReplayMode::Force,
39            WorkspaceExportMode::Off => WorkspaceReplayMode::Off,
40        };
41
42        runtime_export_workspace_state(&entries, replay_mode)
43            .await
44            .map_err(Into::into)
45    }
46
47    pub fn import_workspace_state(&mut self, bytes: &[u8]) -> Result<()> {
48        let entries = runtime_import_workspace_state(bytes)?;
49        self.clear_variables();
50
51        for (index, (name, value)) in entries.into_iter().enumerate() {
52            self.variable_names.insert(name.clone(), index);
53            self.variable_array.push(value.clone());
54            self.variables.insert(name.clone(), value.clone());
55            self.workspace_values.insert(name, value);
56        }
57
58        self.workspace_preview_tokens.clear();
59        self.workspace_version = self.workspace_version.wrapping_add(1);
60        Ok(())
61    }
62
63    pub fn workspace_snapshot(&mut self) -> WorkspaceSnapshot {
64        let source_map = if self.workspace_values.is_empty() {
65            &self.variables
66        } else {
67            &self.workspace_values
68        };
69
70        let mut entries: Vec<WorkspaceEntry> = source_map
71            .iter()
72            .map(|(name, value)| workspace_entry(name, value))
73            .collect();
74        entries.sort_by(|a, b| a.name.cmp(&b.name));
75
76        WorkspaceSnapshot {
77            full: true,
78            version: self.workspace_version,
79            values: self.attach_workspace_preview_tokens(entries),
80        }
81    }
82
83    /// Materialize a workspace variable for inspection (optionally identified by preview token).
84    pub async fn materialize_variable(
85        &mut self,
86        target: WorkspaceMaterializeTarget,
87        options: WorkspaceMaterializeOptions,
88    ) -> Result<MaterializedVariable> {
89        let name = match target {
90            WorkspaceMaterializeTarget::Name(name) => name,
91            WorkspaceMaterializeTarget::Token(id) => self
92                .workspace_preview_tokens
93                .get(&id)
94                .map(|ticket| ticket.name.clone())
95                .ok_or_else(|| anyhow::anyhow!("Unknown workspace preview token"))?,
96        };
97        let value = self
98            .workspace_values
99            .get(&name)
100            .or_else(|| self.variables.get(&name))
101            .cloned()
102            .ok_or_else(|| anyhow::anyhow!("Variable '{name}' not found in workspace"))?;
103
104        let is_gpu = matches!(value, Value::GpuTensor(_));
105        let residency = if is_gpu {
106            WorkspaceResidency::Gpu
107        } else {
108            WorkspaceResidency::Cpu
109        };
110        // For CPU values we can materialize directly. For GPU tensors, avoid downloading the
111        // entire buffer into wasm memory; gather only the requested preview/slice.
112        let host_value = value.clone();
113        let value_shape_vec = value_shape(&host_value).unwrap_or_default();
114        let mut preview = None;
115        if is_gpu {
116            if let Value::GpuTensor(handle) = &value {
117                if let Some((values, truncated)) =
118                    gather_gpu_preview_values(handle, &value_shape_vec, &options).await?
119                {
120                    preview = Some(WorkspacePreview { values, truncated });
121                }
122            }
123        } else {
124            if let Some(slice_opts) = options
125                .slice
126                .as_ref()
127                .and_then(|slice| slice.sanitized(&value_shape_vec))
128            {
129                let slice_elements = slice_opts.shape.iter().product::<usize>();
130                let slice_limit = slice_elements.clamp(1, MATERIALIZE_DEFAULT_LIMIT);
131                if let Some(slice_value) = slice_value_for_preview(&host_value, &slice_opts) {
132                    preview = preview_numeric_values(&slice_value, slice_limit)
133                        .map(|(values, truncated)| WorkspacePreview { values, truncated });
134                }
135            }
136            if preview.is_none() {
137                let max_elements = options.max_elements.clamp(1, MATERIALIZE_DEFAULT_LIMIT);
138                preview = preview_numeric_values(&host_value, max_elements)
139                    .map(|(values, truncated)| WorkspacePreview { values, truncated });
140            }
141        }
142        Ok(MaterializedVariable {
143            name,
144            class_name: matlab_class_name(&host_value),
145            dtype: if let Value::GpuTensor(handle) = &host_value {
146                gpu_dtype_label(handle).map(|label| label.to_string())
147            } else {
148                numeric_dtype_label(&host_value).map(|label| label.to_string())
149            },
150            shape: value_shape_vec,
151            is_gpu,
152            residency,
153            size_bytes: if let Value::GpuTensor(handle) = &host_value {
154                gpu_size_bytes(handle)
155            } else {
156                approximate_size_bytes(&host_value)
157            },
158            preview,
159            value: host_value,
160        })
161    }
162
163    /// Get a copy of current variables
164    pub fn get_variables(&self) -> &HashMap<String, Value> {
165        &self.variables
166    }
167
168    pub(crate) fn build_workspace_snapshot(
169        &mut self,
170        entries: Vec<WorkspaceEntry>,
171        full: bool,
172    ) -> WorkspaceSnapshot {
173        self.workspace_version = self.workspace_version.wrapping_add(1);
174        let version = self.workspace_version;
175        WorkspaceSnapshot {
176            full,
177            version,
178            values: self.attach_workspace_preview_tokens(entries),
179        }
180    }
181
182    fn attach_workspace_preview_tokens(
183        &mut self,
184        entries: Vec<WorkspaceEntry>,
185    ) -> Vec<WorkspaceEntry> {
186        self.workspace_preview_tokens.clear();
187        let mut values = Vec::with_capacity(entries.len());
188        for mut entry in entries {
189            let token = Uuid::new_v4();
190            self.workspace_preview_tokens.insert(
191                token,
192                WorkspaceMaterializeTicket {
193                    name: entry.name.clone(),
194                },
195            );
196            entry.preview_token = Some(token);
197            values.push(entry);
198        }
199        values
200    }
201}
runmat_core/session/workspace.rs

runmat_core/session/
workspace.rs
