1use std::collections::HashMap;
8
9use rayon::prelude::*;
10
11use crate::bytecode::{
12 BytecodeProgram, Constant, DebugInfo, Function, FunctionBlob, FunctionHash, Instruction,
13 LinkedFunction, LinkedProgram, Operand, Program, SourceMap,
14};
15use shape_abi_v1::PermissionSet;
16use shape_value::{FunctionId, StringId};
17
18#[derive(Debug, thiserror::Error)]
23pub enum LinkError {
24 #[error("Missing function blob: {0}")]
25 MissingBlob(FunctionHash),
26 #[error("Circular dependency detected")]
27 CircularDependency,
28 #[error("Constant pool overflow: {0} constants exceeds u16 max")]
29 ConstantPoolOverflow(usize),
30 #[error("String pool overflow: {0} strings exceeds u32 max")]
31 StringPoolOverflow(usize),
32}
33
34fn topo_sort(program: &Program) -> Result<Vec<FunctionHash>, LinkError> {
42 let mut state: HashMap<FunctionHash, u8> = HashMap::new();
44 let mut order: Vec<FunctionHash> = Vec::with_capacity(program.function_store.len());
45
46 fn visit(
47 hash: FunctionHash,
48 program: &Program,
49 state: &mut HashMap<FunctionHash, u8>,
50 order: &mut Vec<FunctionHash>,
51 ) -> Result<(), LinkError> {
52 match state.get(&hash).copied().unwrap_or(0) {
53 2 => return Ok(()), 1 => return Err(LinkError::CircularDependency),
55 _ => {}
56 }
57 state.insert(hash, 1); let blob = program
60 .function_store
61 .get(&hash)
62 .ok_or(LinkError::MissingBlob(hash))?;
63
64 for dep in &blob.dependencies {
65 if *dep == FunctionHash::ZERO {
68 continue;
69 }
70 visit(*dep, program, state, order)?;
71 }
72
73 state.insert(hash, 2); order.push(hash);
75 Ok(())
76 }
77
78 visit(program.entry, program, &mut state, &mut order)?;
82
83 let remaining: Vec<FunctionHash> = program
86 .function_store
87 .keys()
88 .copied()
89 .filter(|h| state.get(h).copied().unwrap_or(0) != 2)
90 .collect();
91 for hash in remaining {
92 visit(hash, program, &mut state, &mut order)?;
93 }
94
95 Ok(order)
96}
97
98fn remap_fid(
106 dep_idx: u16,
107 blob: &FunctionBlob,
108 current_function_id: usize,
109 hash_to_id: &HashMap<FunctionHash, usize>,
110 name_to_id: &HashMap<&str, usize>,
111) -> u16 {
112 if let Some(dep_hash) = blob.dependencies.get(dep_idx as usize) {
113 if *dep_hash == FunctionHash::ZERO {
114 if let Some(callee_name) = blob.callee_names.get(dep_idx as usize) {
116 if callee_name != &blob.name {
117 if let Some(target_id) = name_to_id.get(callee_name.as_str()) {
119 *target_id as u16
120 } else {
121 current_function_id as u16
123 }
124 } else {
125 current_function_id as u16
127 }
128 } else {
129 current_function_id as u16
131 }
132 } else {
133 hash_to_id[dep_hash] as u16
134 }
135 } else {
136 dep_idx
138 }
139}
140
141fn remap_operand(
144 operand: Operand,
145 const_base: usize,
146 string_base: usize,
147 blob: &FunctionBlob,
148 current_function_id: usize,
149 hash_to_id: &HashMap<FunctionHash, usize>,
150 name_to_id: &HashMap<&str, usize>,
151) -> Operand {
152 match operand {
153 Operand::Const(i) => Operand::Const((const_base + i as usize) as u16),
154 Operand::Property(i) => Operand::Property((string_base + i as usize) as u16),
155 Operand::Name(StringId(i)) => Operand::Name(StringId((string_base + i as usize) as u32)),
156 Operand::Function(FunctionId(dep_idx)) => {
157 Operand::Function(FunctionId(remap_fid(
158 dep_idx,
159 blob,
160 current_function_id,
161 hash_to_id,
162 name_to_id,
163 )))
164 }
165 Operand::ClosureAlloc { fid: FunctionId(dep_idx), escapes } => {
166 Operand::ClosureAlloc {
167 fid: FunctionId(remap_fid(
168 dep_idx,
169 blob,
170 current_function_id,
171 hash_to_id,
172 name_to_id,
173 )),
174 escapes,
175 }
176 }
177 Operand::TypedMethodCall {
178 method_id,
179 arg_count,
180 string_id,
181 receiver_type_tag,
182 } => Operand::TypedMethodCall {
183 method_id,
184 arg_count,
185 string_id: (string_base + string_id as usize) as u16,
186 receiver_type_tag,
187 },
188 Operand::Offset(_)
190 | Operand::Local(_)
191 | Operand::ModuleBinding(_)
192 | Operand::Builtin(_)
193 | Operand::Count(_)
194 | Operand::ColumnIndex(_)
195 | Operand::TypedField { .. }
196 | Operand::TypedObjectAlloc { .. }
197 | Operand::TypedMerge { .. }
198 | Operand::ColumnAccess { .. }
199 | Operand::ForeignFunction(_)
200 | Operand::MatrixDims { .. }
201 | Operand::Width(_)
202 | Operand::TypedLocal(_, _)
203 | Operand::TypedModuleBinding(_, _)
204 | Operand::FieldOffset(_) => operand,
205 }
206}
207
208fn remap_constant(
216 constant: &Constant,
217 blob: &FunctionBlob,
218 current_function_id: usize,
219 hash_to_id: &HashMap<FunctionHash, usize>,
220 name_to_id: &HashMap<&str, usize>,
221) -> Constant {
222 match constant {
223 Constant::Function(dep_idx) => {
224 let dep_idx = *dep_idx as usize;
225 if dep_idx < blob.dependencies.len() {
226 let dep_hash = blob.dependencies[dep_idx];
227 if dep_hash == FunctionHash::ZERO {
228 if let Some(callee_name) = blob.callee_names.get(dep_idx) {
230 if callee_name != &blob.name {
231 if let Some(target_id) = name_to_id.get(callee_name.as_str()) {
233 Constant::Function(*target_id as u16)
234 } else {
235 Constant::Function(current_function_id as u16)
236 }
237 } else {
238 Constant::Function(current_function_id as u16)
239 }
240 } else {
241 Constant::Function(current_function_id as u16)
242 }
243 } else {
244 let linked_id = hash_to_id[&dep_hash];
245 Constant::Function(linked_id as u16)
246 }
247 } else {
248 constant.clone()
250 }
251 }
252 other => other.clone(),
253 }
254}
255
256const PARALLEL_THRESHOLD: usize = 50;
263
264struct BlobOffsets {
266 instruction_base: usize,
267 const_base: usize,
268 string_base: usize,
269}
270
271pub fn link(program: &Program) -> Result<LinkedProgram, LinkError> {
282 let sorted = topo_sort(program)?;
283
284 let blobs: Vec<&FunctionBlob> = sorted
286 .iter()
287 .map(|h| {
288 program
289 .function_store
290 .get(h)
291 .ok_or(LinkError::MissingBlob(*h))
292 })
293 .collect::<Result<Vec<_>, _>>()?;
294
295 let mut offsets: Vec<BlobOffsets> = Vec::with_capacity(blobs.len());
299 let mut hash_to_id: HashMap<FunctionHash, usize> = HashMap::with_capacity(blobs.len());
300 let mut name_to_id: HashMap<&str, usize> = HashMap::with_capacity(blobs.len());
301
302 let mut total_instructions: usize = 0;
303 let mut total_constants: usize = 0;
304 let mut total_strings: usize = 0;
305
306 for (i, blob) in blobs.iter().enumerate() {
307 offsets.push(BlobOffsets {
308 instruction_base: total_instructions,
309 const_base: total_constants,
310 string_base: total_strings,
311 });
312 hash_to_id.insert(blob.content_hash, i);
313 name_to_id.insert(&blob.name, i);
314
315 total_instructions += blob.instructions.len();
316 total_constants += blob.constants.len();
317 total_strings += blob.strings.len();
318 }
319
320 if total_constants > u16::MAX as usize + 1 {
322 return Err(LinkError::ConstantPoolOverflow(total_constants));
323 }
324 if total_strings > u32::MAX as usize + 1 {
325 return Err(LinkError::StringPoolOverflow(total_strings));
326 }
327
328 let total_required_permissions = blobs.iter().fold(PermissionSet::pure(), |acc, blob| {
330 acc.union(&blob.required_permissions)
331 });
332
333 let use_parallel = blobs.len() > PARALLEL_THRESHOLD;
337
338 let mut instructions: Vec<Instruction> = Vec::with_capacity(total_instructions);
340 let mut constants: Vec<Constant> = Vec::with_capacity(total_constants);
341 let mut strings: Vec<String> = Vec::with_capacity(total_strings);
342
343 if use_parallel {
344 struct BlobResult {
358 instructions: Vec<Instruction>,
359 constants: Vec<Constant>,
360 strings: Vec<String>,
361 source_map: Vec<(usize, u16, u32)>,
362 }
363
364 let results: Vec<BlobResult> = blobs
365 .par_iter()
366 .zip(offsets.par_iter())
367 .enumerate()
368 .map(|(function_id, (blob, off))| {
369 let remapped_instrs: Vec<Instruction> = blob
370 .instructions
371 .iter()
372 .map(|instr| {
373 let remapped_operand = instr.operand.map(|op| {
374 remap_operand(
375 op,
376 off.const_base,
377 off.string_base,
378 blob,
379 function_id,
380 &hash_to_id,
381 &name_to_id,
382 )
383 });
384 Instruction {
385 opcode: instr.opcode,
386 operand: remapped_operand,
387 }
388 })
389 .collect();
390
391 let remapped_consts: Vec<Constant> = blob
392 .constants
393 .iter()
394 .map(|c| remap_constant(c, blob, function_id, &hash_to_id, &name_to_id))
395 .collect();
396
397 let cloned_strings: Vec<String> = blob.strings.clone();
398
399 let source_entries: Vec<(usize, u16, u32)> = blob
400 .source_map
401 .iter()
402 .map(|&(local_offset, file_id, line)| {
403 (off.instruction_base + local_offset, file_id as u16, line)
404 })
405 .collect();
406
407 BlobResult {
408 instructions: remapped_instrs,
409 constants: remapped_consts,
410 strings: cloned_strings,
411 source_map: source_entries,
412 }
413 })
414 .collect();
415
416 let mut merged_line_numbers: Vec<(usize, u16, u32)> = Vec::new();
419 for result in results {
420 instructions.extend(result.instructions);
421 constants.extend(result.constants);
422 strings.extend(result.strings);
423 merged_line_numbers.extend(result.source_map);
424 }
425
426 merged_line_numbers.sort_by_key(|&(offset, _, _)| offset);
427
428 let functions: Vec<LinkedFunction> = blobs
429 .iter()
430 .zip(offsets.iter())
431 .map(|(blob, off)| LinkedFunction {
432 blob_hash: blob.content_hash,
433 entry_point: off.instruction_base,
434 body_length: blob.instructions.len(),
435 name: blob.name.clone(),
436 arity: blob.arity,
437 param_names: blob.param_names.clone(),
438 locals_count: blob.locals_count,
439 is_closure: blob.is_closure,
440 captures_count: blob.captures_count,
441 is_async: blob.is_async,
442 ref_params: blob.ref_params.clone(),
443 ref_mutates: blob.ref_mutates.clone(),
444 mutable_captures: blob.mutable_captures.clone(),
445 frame_descriptor: blob.frame_descriptor.clone(),
446 })
447 .collect();
448
449 let debug_info = DebugInfo {
450 source_map: SourceMap {
451 files: program.debug_info.source_map.files.clone(),
452 source_texts: program.debug_info.source_map.source_texts.clone(),
453 },
454 line_numbers: merged_line_numbers,
455 variable_names: program.debug_info.variable_names.clone(),
456 source_text: String::new(),
457 };
458
459 return Ok(LinkedProgram {
460 entry: program.entry,
461 instructions,
462 constants,
463 strings,
464 functions,
465 hash_to_id,
466 debug_info,
467 data_schema: program.data_schema.clone(),
468 module_binding_names: program.module_binding_names.clone(),
469 top_level_locals_count: program.top_level_locals_count,
470 top_level_local_storage_hints: program.top_level_local_storage_hints.clone(),
471 type_schema_registry: program.type_schema_registry.clone(),
472 module_binding_storage_hints: program.module_binding_storage_hints.clone(),
473 function_local_storage_hints: program.function_local_storage_hints.clone(),
474 top_level_frame: program.top_level_frame.clone(),
475 top_level_local_concrete_types: program.top_level_local_concrete_types.clone(),
476 function_local_concrete_types: program.function_local_concrete_types.clone(),
477 function_return_concrete_types: program.function_return_concrete_types.clone(),
478 monomorphized_method_call_sites:
479 program.monomorphized_method_call_sites.clone(),
480 value_call_return_concrete_types:
485 program.value_call_return_concrete_types.clone(),
486 operator_trait_dispatch_sites:
491 program.operator_trait_dispatch_sites.clone(),
492 trait_method_symbols: program.trait_method_symbols.clone(),
493 foreign_functions: program.foreign_functions.clone(),
494 native_struct_layouts: program.native_struct_layouts.clone(),
495 total_required_permissions: total_required_permissions.clone(),
496 closure_function_layouts: remap_closure_function_layouts(
497 program,
498 &blobs,
499 ),
500 trait_vtables: program.trait_vtables.clone(),
501 has_imported_const_inline: program.has_imported_const_inline,
502 has_w17_marshal_residual: program.has_w17_marshal_residual,
503 });
504 }
505
506 let mut merged_line_numbers: Vec<(usize, u16, u32)> = Vec::new();
510
511 for (function_id, (blob, off)) in blobs.iter().zip(offsets.iter()).enumerate() {
512 for instr in &blob.instructions {
514 let remapped_operand = instr.operand.map(|op| {
515 remap_operand(
516 op,
517 off.const_base,
518 off.string_base,
519 blob,
520 function_id,
521 &hash_to_id,
522 &name_to_id,
523 )
524 });
525 instructions.push(Instruction {
526 opcode: instr.opcode,
527 operand: remapped_operand,
528 });
529 }
530
531 for c in &blob.constants {
533 constants.push(remap_constant(
534 c,
535 blob,
536 function_id,
537 &hash_to_id,
538 &name_to_id,
539 ));
540 }
541
542 strings.extend(blob.strings.iter().cloned());
544
545 for &(local_offset, file_id, line) in &blob.source_map {
547 let global_offset = off.instruction_base + local_offset;
548 merged_line_numbers.push((global_offset, file_id as u16, line));
549 }
550 }
551
552 merged_line_numbers.sort_by_key(|&(offset, _, _)| offset);
554
555 let functions: Vec<LinkedFunction> = blobs
556 .iter()
557 .zip(offsets.iter())
558 .map(|(blob, off)| LinkedFunction {
559 blob_hash: blob.content_hash,
560 entry_point: off.instruction_base,
561 body_length: blob.instructions.len(),
562 name: blob.name.clone(),
563 arity: blob.arity,
564 param_names: blob.param_names.clone(),
565 locals_count: blob.locals_count,
566 is_closure: blob.is_closure,
567 captures_count: blob.captures_count,
568 is_async: blob.is_async,
569 ref_params: blob.ref_params.clone(),
570 ref_mutates: blob.ref_mutates.clone(),
571 mutable_captures: blob.mutable_captures.clone(),
572 frame_descriptor: blob.frame_descriptor.clone(),
573 })
574 .collect();
575
576 let debug_info = DebugInfo {
577 source_map: SourceMap {
578 files: program.debug_info.source_map.files.clone(),
579 source_texts: program.debug_info.source_map.source_texts.clone(),
580 },
581 line_numbers: merged_line_numbers,
582 variable_names: program.debug_info.variable_names.clone(),
583 source_text: String::new(),
584 };
585
586 Ok(LinkedProgram {
587 entry: program.entry,
588 instructions,
589 constants,
590 strings,
591 functions,
592 hash_to_id,
593 debug_info,
594 data_schema: program.data_schema.clone(),
595 module_binding_names: program.module_binding_names.clone(),
596 top_level_locals_count: program.top_level_locals_count,
597 top_level_local_storage_hints: program.top_level_local_storage_hints.clone(),
598 type_schema_registry: program.type_schema_registry.clone(),
599 module_binding_storage_hints: program.module_binding_storage_hints.clone(),
600 function_local_storage_hints: program.function_local_storage_hints.clone(),
601 top_level_frame: program.top_level_frame.clone(),
602 top_level_local_concrete_types: program.top_level_local_concrete_types.clone(),
603 function_local_concrete_types: program.function_local_concrete_types.clone(),
604 function_return_concrete_types: program.function_return_concrete_types.clone(),
605 monomorphized_method_call_sites:
606 program.monomorphized_method_call_sites.clone(),
607 value_call_return_concrete_types:
612 program.value_call_return_concrete_types.clone(),
613 operator_trait_dispatch_sites:
618 program.operator_trait_dispatch_sites.clone(),
619 trait_method_symbols: program.trait_method_symbols.clone(),
620 foreign_functions: program.foreign_functions.clone(),
621 native_struct_layouts: program.native_struct_layouts.clone(),
622 total_required_permissions,
623 closure_function_layouts: remap_closure_function_layouts(program, &blobs),
624 trait_vtables: program.trait_vtables.clone(),
625 has_imported_const_inline: program.has_imported_const_inline,
626 has_w17_marshal_residual: program.has_w17_marshal_residual,
627 })
628}
629
630fn remap_closure_function_layouts(
636 program: &Program,
637 blobs: &[&FunctionBlob],
638) -> Vec<Option<std::sync::Arc<shape_value::v2::closure_layout::ClosureLayout>>> {
639 if program.closure_function_layouts_by_name.is_empty() {
640 return Vec::new();
641 }
642 blobs
643 .iter()
644 .map(|blob| {
645 program
646 .closure_function_layouts_by_name
647 .get(&blob.name)
648 .cloned()
649 })
650 .collect()
651}
652
653pub fn linked_to_bytecode_program(linked: &LinkedProgram) -> BytecodeProgram {
660 let functions: Vec<Function> = linked
661 .functions
662 .iter()
663 .map(|lf| Function {
664 name: lf.name.clone(),
665 arity: lf.arity,
666 param_names: lf.param_names.clone(),
667 locals_count: lf.locals_count,
668 entry_point: lf.entry_point,
669 body_length: lf.body_length,
670 is_closure: lf.is_closure,
671 captures_count: lf.captures_count,
672 is_async: lf.is_async,
673 ref_params: lf.ref_params.clone(),
674 ref_mutates: lf.ref_mutates.clone(),
675 mutable_captures: lf.mutable_captures.clone(),
676 frame_descriptor: lf.frame_descriptor.clone(),
677 osr_entry_points: Vec::new(),
678 mir_data: None,
679 })
680 .collect();
681
682 BytecodeProgram {
683 instructions: linked.instructions.clone(),
684 constants: linked.constants.clone(),
685 strings: linked.strings.clone(),
686 functions,
687 has_imported_const_inline: linked.has_imported_const_inline,
688 has_w17_marshal_residual: linked.has_w17_marshal_residual,
689 debug_info: linked.debug_info.clone(),
690 data_schema: linked.data_schema.clone(),
691 module_binding_names: linked.module_binding_names.clone(),
692 top_level_locals_count: linked.top_level_locals_count,
693 top_level_local_storage_hints: linked.top_level_local_storage_hints.clone(),
694 type_schema_registry: linked.type_schema_registry.clone(),
695 module_binding_storage_hints: linked.module_binding_storage_hints.clone(),
696 function_local_storage_hints: linked.function_local_storage_hints.clone(),
697 top_level_frame: linked.top_level_frame.clone(),
698 top_level_local_concrete_types: linked.top_level_local_concrete_types.clone(),
699 function_local_concrete_types: linked.function_local_concrete_types.clone(),
700 function_return_concrete_types: linked.function_return_concrete_types.clone(),
701 monomorphized_method_call_sites:
702 linked.monomorphized_method_call_sites.clone(),
703 value_call_return_concrete_types:
708 linked.value_call_return_concrete_types.clone(),
709 operator_trait_dispatch_sites:
714 linked.operator_trait_dispatch_sites.clone(),
715 top_level_mir: None,
716 compiled_annotations: HashMap::new(),
717 trait_method_symbols: linked.trait_method_symbols.clone(),
718 expanded_function_defs: HashMap::new(),
719 string_index: HashMap::new(),
720 foreign_functions: linked.foreign_functions.clone(),
721 native_struct_layouts: linked.native_struct_layouts.clone(),
722 content_addressed: None,
723 function_blob_hashes: linked
724 .functions
725 .iter()
726 .map(|lf| {
727 if lf.blob_hash == FunctionHash::ZERO {
728 None
729 } else {
730 Some(lf.blob_hash)
731 }
732 })
733 .collect(),
734 monomorphization_keys: Vec::new(),
735 closure_function_layouts: linked.closure_function_layouts.clone(),
736 trait_vtables: linked.trait_vtables.clone(),
737 }
738}
739
740#[cfg(test)]
745#[path = "linker_tests.rs"]
746mod tests;