1use runmat_hir::{
2 BindingId, BindingStorage, EntrypointId, FunctionAbi, FunctionId, HirAssembly,
3 WorkspaceExportPolicy, WorkspaceVisibility,
4};
5use runmat_mir::{MirAssembly, MirLocalId};
6use serde::{Deserialize, Serialize};
7use std::collections::HashMap;
8
9#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
10pub struct VmSlotId(pub usize);
11
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13pub struct VmAssemblyLayout {
14 pub functions: HashMap<FunctionId, VmFunctionLayout>,
15 pub entrypoints: HashMap<EntrypointId, VmEntrypointLayout>,
16 pub storage_bindings: HashMap<BindingId, VmStorageBinding>,
17}
18
19#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
20pub struct VmFunctionLayout {
21 pub function: FunctionId,
22 pub display_name: String,
23 pub private_owner_scope: String,
24 pub frame_abi: VmFrameAbi,
25 pub binding_slots: HashMap<BindingId, VmSlotId>,
26 pub mir_local_slots: HashMap<MirLocalId, VmSlotId>,
27 pub captures: Vec<VmCaptureSlot>,
28 pub local_count: usize,
29}
30
31#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
32pub struct VmFrameAbi {
33 pub fixed_inputs: Vec<VmSlotId>,
34 pub varargin: Option<VmSlotId>,
35 pub fixed_outputs: Vec<VmSlotId>,
36 pub varargout: Option<VmSlotId>,
37 pub implicit_nargin: Option<VmSlotId>,
38 pub implicit_nargout: Option<VmSlotId>,
39}
40
41#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
42pub struct VmCaptureSlot {
43 pub binding: BindingId,
44 pub from_function: FunctionId,
45 pub slot: VmSlotId,
46}
47
48#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
49pub struct VmEntrypointLayout {
50 pub entrypoint: EntrypointId,
51 pub target: FunctionId,
52 pub workspace_export: WorkspaceExportPolicy,
53 pub exports: Vec<VmWorkspaceExport>,
54}
55
56#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
57pub struct VmWorkspaceExport {
58 pub binding: BindingId,
59 pub name: String,
60 pub slot: VmSlotId,
61 pub visibility: WorkspaceVisibility,
62}
63
64#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
65pub struct VmStorageBinding {
66 pub binding: BindingId,
67 pub name: String,
68 pub storage: BindingStorage,
69}
70
71#[derive(Debug, Clone, PartialEq, Eq)]
72pub enum LayoutError {
73 MissingFunction(FunctionId),
74 MissingMirBody(FunctionId),
75 MissingBinding(BindingId),
76 MissingBindingSlot {
77 function: FunctionId,
78 binding: BindingId,
79 },
80}
81
82pub fn derive_layout(
83 hir: &HirAssembly,
84 mir: &MirAssembly,
85) -> Result<VmAssemblyLayout, LayoutError> {
86 let mut functions = HashMap::new();
87 for function in &hir.functions {
88 let body = mir
89 .bodies
90 .get(&function.id)
91 .ok_or(LayoutError::MissingMirBody(function.id))?;
92 functions.insert(function.id, derive_function_layout(hir, function, body)?);
93 }
94
95 let mut entrypoints = HashMap::new();
96 for entrypoint in &hir.entrypoints {
97 let function_layout = functions
98 .get(&entrypoint.target)
99 .ok_or(LayoutError::MissingFunction(entrypoint.target))?;
100 let exports = match entrypoint.policy.workspace_export {
101 WorkspaceExportPolicy::ExportTopLevelBindings | WorkspaceExportPolicy::HostDefined => {
102 workspace_exports_for_function(hir, function_layout)?
103 }
104 WorkspaceExportPolicy::ReturnFunctionOutputs => Vec::new(),
105 };
106 entrypoints.insert(
107 entrypoint.id,
108 VmEntrypointLayout {
109 entrypoint: entrypoint.id,
110 target: entrypoint.target,
111 workspace_export: entrypoint.policy.workspace_export.clone(),
112 exports,
113 },
114 );
115 }
116
117 let storage_bindings = hir
118 .bindings
119 .iter()
120 .filter(|binding| binding.storage != BindingStorage::Lexical)
121 .map(|binding| {
122 (
123 binding.id,
124 VmStorageBinding {
125 binding: binding.id,
126 name: binding.name.0.clone(),
127 storage: binding.storage.clone(),
128 },
129 )
130 })
131 .collect();
132
133 Ok(VmAssemblyLayout {
134 functions,
135 entrypoints,
136 storage_bindings,
137 })
138}
139
140fn derive_function_layout(
141 hir: &HirAssembly,
142 function: &runmat_hir::HirFunction,
143 body: &runmat_mir::MirBody,
144) -> Result<VmFunctionLayout, LayoutError> {
145 let mut binding_slots = HashMap::new();
146 let mut next_slot = 0usize;
147
148 let frame_abi = allocate_frame_abi(&function.abi, &mut binding_slots, &mut next_slot);
149
150 for capture in &function.captures {
151 allocate_binding_slot(capture.binding, &mut binding_slots, &mut next_slot);
152 }
153 for binding in &function.locals {
154 allocate_binding_slot(*binding, &mut binding_slots, &mut next_slot);
155 }
156 for local in &body.locals {
157 if let Some(binding) = local.binding {
158 allocate_binding_slot(binding, &mut binding_slots, &mut next_slot);
159 }
160 }
161
162 let mut mir_local_slots = HashMap::new();
163 for local in &body.locals {
164 let slot = if let Some(binding) = local.binding {
165 *binding_slots
166 .get(&binding)
167 .ok_or(LayoutError::MissingBindingSlot {
168 function: function.id,
169 binding,
170 })?
171 } else {
172 let slot = VmSlotId(next_slot);
173 next_slot += 1;
174 slot
175 };
176 mir_local_slots.insert(local.id, slot);
177 }
178
179 let captures = function
180 .captures
181 .iter()
182 .map(|capture| {
183 let slot =
184 *binding_slots
185 .get(&capture.binding)
186 .ok_or(LayoutError::MissingBindingSlot {
187 function: function.id,
188 binding: capture.binding,
189 })?;
190 Ok(VmCaptureSlot {
191 binding: capture.binding,
192 from_function: capture.from_function,
193 slot,
194 })
195 })
196 .collect::<Result<_, LayoutError>>()?;
197
198 for binding in binding_slots.keys() {
199 if hir.bindings.get(binding.0).map(|b| b.id) != Some(*binding) {
200 return Err(LayoutError::MissingBinding(*binding));
201 }
202 }
203
204 Ok(VmFunctionLayout {
205 function: function.id,
206 display_name: function.name.0.clone(),
207 private_owner_scope: private_owner_scope_for_function(hir, function),
208 frame_abi,
209 binding_slots,
210 mir_local_slots,
211 captures,
212 local_count: next_slot,
213 })
214}
215
216fn private_owner_scope_from_name(name: &str) -> Option<String> {
217 if let Some((owner, _)) = name.split_once(".__private__.") {
218 return Some(owner.to_string());
219 }
220 name.rsplit_once('.').map(|(owner, _)| owner.to_string())
221}
222
223fn private_owner_scope_for_function(
224 hir: &HirAssembly,
225 function: &runmat_hir::HirFunction,
226) -> String {
227 if let Some(owner) =
228 private_owner_scope_from_name(&function.name.0).filter(|owner| !owner.is_empty())
229 {
230 return owner;
231 }
232
233 let mut parent = function.parent;
234 while let Some(parent_id) = parent {
235 let Some(parent_function) = hir
236 .functions
237 .iter()
238 .find(|candidate| candidate.id == parent_id)
239 else {
240 break;
241 };
242 if let Some(owner) =
243 private_owner_scope_from_name(&parent_function.name.0).filter(|owner| !owner.is_empty())
244 {
245 return owner;
246 }
247 parent = parent_function.parent;
248 }
249 String::new()
250}
251
252fn allocate_frame_abi(
253 abi: &FunctionAbi,
254 binding_slots: &mut HashMap<BindingId, VmSlotId>,
255 next_slot: &mut usize,
256) -> VmFrameAbi {
257 VmFrameAbi {
258 fixed_inputs: abi
259 .fixed_inputs
260 .iter()
261 .map(|binding| allocate_binding_slot(*binding, binding_slots, next_slot))
262 .collect(),
263 varargin: abi
264 .varargin
265 .map(|binding| allocate_binding_slot(binding, binding_slots, next_slot)),
266 fixed_outputs: abi
267 .fixed_outputs
268 .iter()
269 .map(|binding| allocate_binding_slot(*binding, binding_slots, next_slot))
270 .collect(),
271 varargout: abi
272 .varargout
273 .map(|binding| allocate_binding_slot(binding, binding_slots, next_slot)),
274 implicit_nargin: abi
275 .implicit_nargin
276 .map(|binding| allocate_binding_slot(binding, binding_slots, next_slot)),
277 implicit_nargout: abi
278 .implicit_nargout
279 .map(|binding| allocate_binding_slot(binding, binding_slots, next_slot)),
280 }
281}
282
283fn allocate_binding_slot(
284 binding: BindingId,
285 binding_slots: &mut HashMap<BindingId, VmSlotId>,
286 next_slot: &mut usize,
287) -> VmSlotId {
288 if let Some(slot) = binding_slots.get(&binding) {
289 return *slot;
290 }
291 let slot = VmSlotId(*next_slot);
292 *next_slot += 1;
293 binding_slots.insert(binding, slot);
294 slot
295}
296
297fn workspace_exports_for_function(
298 hir: &HirAssembly,
299 function_layout: &VmFunctionLayout,
300) -> Result<Vec<VmWorkspaceExport>, LayoutError> {
301 let mut exports = Vec::new();
302 for (binding, slot) in &function_layout.binding_slots {
303 let hir_binding = hir
304 .bindings
305 .get(binding.0)
306 .ok_or(LayoutError::MissingBinding(*binding))?;
307 match hir_binding.workspace_visibility {
308 WorkspaceVisibility::TopLevel
309 | WorkspaceVisibility::ModuleVisible
310 | WorkspaceVisibility::ImplicitAns => exports.push(VmWorkspaceExport {
311 binding: *binding,
312 name: hir_binding.name.0.clone(),
313 slot: *slot,
314 visibility: hir_binding.workspace_visibility.clone(),
315 }),
316 WorkspaceVisibility::Hidden => {}
317 }
318 }
319 exports.sort_by_key(|export| export.slot);
320 Ok(exports)
321}
322
323#[cfg(test)]
324mod tests {
325 use super::*;
326 use runmat_hir::{
327 BindingName, BindingOwner, BindingRole, CapturedBinding, EntrypointOrigin,
328 EntrypointPolicy, FunctionKind, FunctionModifiers, FunctionName, HirBinding, HirBlock,
329 HirEntrypoint, HirFunction, ModuleId, Span,
330 };
331 use runmat_mir::{MirBody, MirLocal, MirLocalKind};
332
333 #[test]
334 fn layout_reuses_shared_input_output_binding_slot() {
335 let function = FunctionId(0);
336 let binding = BindingId(0);
337 let assembly = HirAssembly {
338 functions: vec![HirFunction {
339 id: function,
340 module: ModuleId(0),
341 parent: None,
342 enclosing_class: None,
343 name: FunctionName("f".into()),
344 kind: FunctionKind::Named,
345 params: vec![binding],
346 outputs: vec![binding],
347 abi: FunctionAbi {
348 fixed_inputs: vec![binding],
349 varargin: None,
350 fixed_outputs: vec![binding],
351 varargout: None,
352 implicit_nargin: None,
353 implicit_nargout: None,
354 },
355 argument_validations: Vec::new(),
356 locals: Vec::new(),
357 captures: Vec::new(),
358 modifiers: FunctionModifiers::default(),
359 body: HirBlock { statements: vec![] },
360 span: Span::default(),
361 }],
362 bindings: vec![HirBinding {
363 id: binding,
364 owner: BindingOwner::Function(function),
365 name: BindingName("x".into()),
366 role: BindingRole::Parameter,
367 storage: BindingStorage::Lexical,
368 workspace_visibility: WorkspaceVisibility::Hidden,
369 declared_span: Span::default(),
370 }],
371 ..HirAssembly::default()
372 };
373 let mir = MirAssembly {
374 bodies: HashMap::from([(
375 function,
376 MirBody {
377 function,
378 abi: assembly.functions[0].abi.clone(),
379 locals: vec![MirLocal {
380 id: MirLocalId(0),
381 binding: Some(binding),
382 kind: MirLocalKind::Parameter,
383 span: Span::default(),
384 }],
385 blocks: vec![],
386 },
387 )]),
388 };
389
390 let layout = derive_layout(&assembly, &mir).expect("layout");
391 let function_layout = &layout.functions[&function];
392 assert_eq!(function_layout.frame_abi.fixed_inputs, vec![VmSlotId(0)]);
393 assert_eq!(function_layout.frame_abi.fixed_outputs, vec![VmSlotId(0)]);
394 assert_eq!(function_layout.binding_slots[&binding], VmSlotId(0));
395 assert_eq!(function_layout.mir_local_slots[&MirLocalId(0)], VmSlotId(0));
396 assert_eq!(function_layout.local_count, 1);
397 }
398
399 #[test]
400 fn entrypoint_exports_workspace_visible_bindings() {
401 let function = FunctionId(0);
402 let hidden = BindingId(0);
403 let visible = BindingId(1);
404 let entrypoint = EntrypointId(0);
405 let assembly = HirAssembly {
406 functions: vec![HirFunction {
407 id: function,
408 module: ModuleId(0),
409 parent: None,
410 enclosing_class: None,
411 name: FunctionName("entry".into()),
412 kind: FunctionKind::SyntheticEntrypoint,
413 params: Vec::new(),
414 outputs: Vec::new(),
415 abi: FunctionAbi {
416 fixed_inputs: Vec::new(),
417 varargin: None,
418 fixed_outputs: Vec::new(),
419 varargout: None,
420 implicit_nargin: None,
421 implicit_nargout: None,
422 },
423 argument_validations: Vec::new(),
424 locals: vec![hidden, visible],
425 captures: vec![CapturedBinding {
426 binding: hidden,
427 from_function: function,
428 }],
429 modifiers: FunctionModifiers::default(),
430 body: HirBlock { statements: vec![] },
431 span: Span::default(),
432 }],
433 bindings: vec![
434 HirBinding {
435 id: hidden,
436 owner: BindingOwner::Function(function),
437 name: BindingName("tmp".into()),
438 role: BindingRole::Local,
439 storage: BindingStorage::Lexical,
440 workspace_visibility: WorkspaceVisibility::Hidden,
441 declared_span: Span::default(),
442 },
443 HirBinding {
444 id: visible,
445 owner: BindingOwner::Function(function),
446 name: BindingName("x".into()),
447 role: BindingRole::Local,
448 storage: BindingStorage::Lexical,
449 workspace_visibility: WorkspaceVisibility::TopLevel,
450 declared_span: Span::default(),
451 },
452 ],
453 entrypoints: vec![HirEntrypoint {
454 id: entrypoint,
455 name: None,
456 target: function,
457 origin: EntrypointOrigin::HostSynthetic,
458 policy: EntrypointPolicy {
459 workspace_export: WorkspaceExportPolicy::ExportTopLevelBindings,
460 top_level_await: false,
461 },
462 }],
463 ..HirAssembly::default()
464 };
465 let mir = MirAssembly {
466 bodies: HashMap::from([(
467 function,
468 MirBody {
469 function,
470 abi: assembly.functions[0].abi.clone(),
471 locals: vec![
472 MirLocal {
473 id: MirLocalId(0),
474 binding: Some(hidden),
475 kind: MirLocalKind::Binding,
476 span: Span::default(),
477 },
478 MirLocal {
479 id: MirLocalId(1),
480 binding: Some(visible),
481 kind: MirLocalKind::Binding,
482 span: Span::default(),
483 },
484 ],
485 blocks: vec![],
486 },
487 )]),
488 };
489
490 let layout = derive_layout(&assembly, &mir).expect("layout");
491 let exports = &layout.entrypoints[&entrypoint].exports;
492 assert_eq!(exports.len(), 1);
493 assert_eq!(exports[0].binding, visible);
494 assert_eq!(exports[0].name, "x");
495 }
496}