runmat_core/session/
workspace.rs1use 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 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 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 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}