use super::*;
impl RunMatSession {
pub fn workspace_handle(&self) -> crate::abi::WorkspaceHandle {
self.abi_workspace_handle
}
pub(crate) fn workspace_binding_key(&self, name: &str) -> crate::abi::WorkspaceBindingKey {
let binding = runmat_hir::BindingName(name.to_string());
if let Some(source) = self.active_source_identity.clone() {
crate::abi::WorkspaceBindingKey::SourceBinding {
source,
def_path: source_binding_def_path(self.current_source_name()),
binding,
}
} else {
crate::abi::WorkspaceBindingKey::Interactive {
session: self.abi_workspace_handle.0,
name: binding,
}
}
}
pub(crate) fn bind_workspace_slot(&mut self, name: String, slot: usize) {
let key = self.workspace_binding_key(&name);
self.workspace_bindings
.insert(name, SessionWorkspaceBinding { key, slot });
}
pub(crate) fn lowering_workspace_bindings(&self) -> HashMap<String, usize> {
self.workspace_bindings
.iter()
.map(|(name, binding)| (name.clone(), binding.slot))
.collect()
}
pub fn clear_variables(&mut self) {
self.variable_array.clear();
self.workspace_bindings.clear();
self.workspace_values.clear();
self.workspace_preview_tokens.clear();
}
pub async fn export_workspace_state(
&mut self,
mode: WorkspaceExportMode,
) -> Result<Option<Vec<u8>>> {
if matches!(mode, WorkspaceExportMode::Off) {
return Ok(None);
}
let mut entries: Vec<(String, Value)> = Vec::with_capacity(self.workspace_values.len());
for (name, value) in &self.workspace_values {
let gathered = gather_if_needed_async(value).await?;
entries.push((name.clone(), gathered));
}
if entries.is_empty() && matches!(mode, WorkspaceExportMode::Auto) {
return Ok(None);
}
let replay_mode = match mode {
WorkspaceExportMode::Auto => WorkspaceReplayMode::Auto,
WorkspaceExportMode::Force => WorkspaceReplayMode::Force,
WorkspaceExportMode::Off => WorkspaceReplayMode::Off,
};
runtime_export_workspace_state(&entries, replay_mode)
.await
.map_err(Into::into)
}
pub fn import_workspace_state(&mut self, bytes: &[u8]) -> Result<()> {
let entries = runtime_import_workspace_state(bytes)?;
self.clear_variables();
for (index, (name, value)) in entries.into_iter().enumerate() {
self.bind_workspace_slot(name.clone(), index);
self.variable_array.push(value.clone());
self.workspace_values.insert(name, value);
}
self.workspace_preview_tokens.clear();
self.workspace_version = self.workspace_version.wrapping_add(1);
Ok(())
}
pub fn workspace_snapshot(&mut self) -> WorkspaceSnapshot {
let mut entries: Vec<WorkspaceEntry> = self
.workspace_values
.iter()
.map(|(name, value)| workspace_entry(name, value))
.collect();
entries.sort_by(|a, b| a.name.cmp(&b.name));
WorkspaceSnapshot {
full: true,
version: self.workspace_version,
values: self.attach_workspace_preview_tokens(entries),
}
}
pub async fn materialize_variable(
&mut self,
target: WorkspaceMaterializeTarget,
options: WorkspaceMaterializeOptions,
) -> Result<MaterializedVariable> {
let name = match target {
WorkspaceMaterializeTarget::Name(name) => name,
WorkspaceMaterializeTarget::Token(id) => self
.workspace_preview_tokens
.get(&id)
.map(|ticket| ticket.name.clone())
.ok_or_else(|| anyhow::anyhow!("Unknown workspace preview token"))?,
};
let value = self
.workspace_values
.get(&name)
.cloned()
.ok_or_else(|| anyhow::anyhow!("Variable '{name}' not found in workspace"))?;
let is_gpu = matches!(value, Value::GpuTensor(_));
let residency = if is_gpu {
WorkspaceResidency::Gpu
} else {
WorkspaceResidency::Cpu
};
let host_value = value.clone();
let value_shape_vec = value_shape(&host_value).unwrap_or_default();
let mut preview = None;
if is_gpu {
if let Value::GpuTensor(handle) = &value {
if let Some((values, truncated)) =
gather_gpu_preview_values(handle, &value_shape_vec, &options).await?
{
preview = Some(WorkspacePreview { values, truncated });
}
}
} else {
if let Some(slice_opts) = options
.slice
.as_ref()
.and_then(|slice| slice.sanitized(&value_shape_vec))
{
let slice_elements = slice_opts.shape.iter().product::<usize>();
let slice_limit = slice_elements.clamp(1, MATERIALIZE_DEFAULT_LIMIT);
if let Some(slice_value) = slice_value_for_preview(&host_value, &slice_opts) {
preview = preview_numeric_values(&slice_value, slice_limit)
.map(|(values, truncated)| WorkspacePreview { values, truncated });
}
}
if preview.is_none() {
let max_elements = options.max_elements.clamp(1, MATERIALIZE_DEFAULT_LIMIT);
preview = preview_numeric_values(&host_value, max_elements)
.map(|(values, truncated)| WorkspacePreview { values, truncated });
}
}
Ok(MaterializedVariable {
name,
class_name: matlab_class_name(&host_value),
dtype: if let Value::GpuTensor(handle) = &host_value {
gpu_dtype_label(handle).map(|label| label.to_string())
} else {
numeric_dtype_label(&host_value).map(|label| label.to_string())
},
shape: value_shape_vec,
is_gpu,
residency,
size_bytes: if let Value::GpuTensor(handle) = &host_value {
gpu_size_bytes(handle)
} else {
approximate_size_bytes(&host_value)
},
preview,
value: host_value,
})
}
pub fn get_variables(&self) -> &HashMap<String, Value> {
&self.workspace_values
}
pub(crate) fn build_workspace_snapshot(
&mut self,
entries: Vec<WorkspaceEntry>,
full: bool,
) -> WorkspaceSnapshot {
self.workspace_version = self.workspace_version.wrapping_add(1);
let version = self.workspace_version;
WorkspaceSnapshot {
full,
version,
values: self.attach_workspace_preview_tokens(entries),
}
}
fn attach_workspace_preview_tokens(
&mut self,
entries: Vec<WorkspaceEntry>,
) -> Vec<WorkspaceEntry> {
self.workspace_preview_tokens.clear();
let mut values = Vec::with_capacity(entries.len());
for mut entry in entries {
let token = Uuid::new_v4();
self.workspace_preview_tokens.insert(
token,
WorkspaceMaterializeTicket {
name: entry.name.clone(),
},
);
entry.preview_token = Some(token);
values.push(entry);
}
values
}
}
fn source_binding_def_path(source_name: &str) -> runmat_hir::DefPath {
use std::path::Path;
let stem = Path::new(source_name)
.file_stem()
.and_then(|stem| stem.to_str())
.filter(|stem| !stem.trim().is_empty())
.unwrap_or("entrypoint")
.to_string();
runmat_hir::DefPath {
package: runmat_hir::PackageName("workspace".to_string()),
module: runmat_hir::QualifiedName(vec![runmat_hir::SymbolName(stem.clone())]),
item: vec![runmat_hir::DefPathSegment::Function(
runmat_hir::SymbolName(stem),
)],
}
}