Skip to main content

runmat_core/session/
workspace.rs

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