1use {
2 crate::{
3 CompileError, SbpfArch,
4 astnode::{ASTNode, ROData},
5 dynsym::{DynamicSymbolMap, RelDynMap, RelocationType},
6 header::ProgramHeader,
7 optimizer,
8 parser::ProgramLayout,
9 section::{CodeSection, DataSection},
10 },
11 either::Either,
12 sbpf_common::{
13 inst_param::{Number, Register},
14 instruction::Instruction,
15 opcode::Opcode,
16 },
17 std::{
18 collections::{HashMap, HashSet},
19 path::PathBuf,
20 },
21 syscall_map::murmur3_32,
22};
23
24type LabelOffsetMap = HashMap<String, u64>;
25type NumericLabel = (String, u64, usize);
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub enum OptimizationConfig {
29 Disabled,
30 Enabled { cfg_dump_dir: Option<PathBuf> },
31}
32
33impl Default for OptimizationConfig {
34 fn default() -> Self {
35 Self::disabled()
36 }
37}
38
39impl OptimizationConfig {
40 pub fn disabled() -> Self {
41 Self::Disabled
42 }
43
44 pub fn enabled() -> Self {
45 Self::Enabled { cfg_dump_dir: None }
46 }
47
48 pub fn with_cfg_dump_dir(self, path: impl Into<PathBuf>) -> Self {
49 match self {
50 Self::Enabled { .. } => Self::Enabled {
51 cfg_dump_dir: Some(path.into()),
52 },
53 Self::Disabled => Self::Disabled,
54 }
55 }
56}
57
58#[derive(Default, Debug)]
59pub struct AST {
60 pub nodes: Vec<ASTNode>,
61 pub rodata_nodes: Vec<ASTNode>,
62
63 function_entries: HashSet<String>,
64 text_size: u64,
65 rodata_size: u64,
66}
67
68impl AST {
69 pub fn new() -> Self {
70 Self::default()
71 }
72
73 pub fn add_function_entry(&mut self, name: String) {
74 self.function_entries.insert(name);
75 }
76
77 pub(crate) fn function_entries(&self) -> &HashSet<String> {
78 &self.function_entries
79 }
80
81 pub fn set_text_size(&mut self, text_size: u64) {
83 self.text_size = text_size;
84 }
85
86 pub fn set_rodata_size(&mut self, rodata_size: u64) {
88 self.rodata_size = rodata_size;
89 }
90
91 pub fn get_instruction_at_offset(&mut self, offset: u64) -> Option<&mut Instruction> {
93 self.nodes
94 .iter_mut()
95 .find(|node| match node {
96 ASTNode::Instruction {
97 instruction: _,
98 offset: inst_offset,
99 ..
100 } => offset == *inst_offset,
101 _ => false,
102 })
103 .map(|node| match node {
104 ASTNode::Instruction { instruction, .. } => instruction,
105 _ => panic!("Expected Instruction node"),
106 })
107 }
108
109 pub fn get_rodata_at_offset(&self, offset: u64) -> Option<&ROData> {
111 self.rodata_nodes
112 .iter()
113 .find(|node| match node {
114 ASTNode::ROData {
115 rodata: _,
116 offset: rodata_offset,
117 ..
118 } => offset == *rodata_offset,
119 _ => false,
120 })
121 .map(|node| match node {
122 ASTNode::ROData { rodata, .. } => rodata,
123 _ => panic!("Expected ROData node"),
124 })
125 }
126
127 pub(crate) fn resolve_numeric_label(
129 label_ref: &str,
130 current_idx: usize,
131 numeric_labels: &[NumericLabel],
132 ) -> Option<u64> {
133 if let Some(direction) = label_ref.chars().last()
134 && (direction == 'f' || direction == 'b')
135 {
136 let label_num = &label_ref[..label_ref.len() - 1];
137
138 if direction == 'f' {
139 for (name, offset, node_idx) in numeric_labels {
141 if name == label_num && *node_idx > current_idx {
142 return Some(*offset);
143 }
144 }
145 } else {
146 for (name, offset, node_idx) in numeric_labels.iter().rev() {
148 if name == label_num && *node_idx < current_idx {
149 return Some(*offset);
150 }
151 }
152 }
153 }
154 None
155 }
156}
157
158pub fn build_program(
159 mut ast: AST,
160 arch: SbpfArch,
161 optimization: OptimizationConfig,
162) -> Result<ProgramLayout, Vec<CompileError>> {
163 let optimization = run_optimizations(&mut ast, &optimization);
164 let mut errors = optimization.errors;
165
166 let (label_offset_map, numeric_labels) = label_offset_map(&ast);
167 let program_is_static = arch.is_v3()
168 || !ast.nodes.iter().any(|node| {
169 matches!(node, ASTNode::Instruction { instruction: inst, .. }
170 if inst.is_syscall()
171 || (inst.opcode == Opcode::Lddw && matches!(&inst.imm, Some(Either::Left(_)))))
172 });
173
174 let label_resolution = resolve_label_references(
175 &mut ast,
176 arch,
177 program_is_static,
178 &label_offset_map,
179 &numeric_labels,
180 );
181 errors.extend(label_resolution.errors);
182
183 optimizer::remove_temp_control_flow_target_labels(
184 &mut ast.nodes,
185 &optimization.labels_to_remove,
186 );
187
188 if !errors.is_empty() {
189 Err(errors)
190 } else {
191 Ok(ProgramLayout {
192 code_section: CodeSection::new(std::mem::take(&mut ast.nodes), ast.text_size),
193 data_section: DataSection::new(std::mem::take(&mut ast.rodata_nodes), ast.rodata_size),
194 dynamic_symbols: label_resolution.dynamic_symbols,
195 relocation_data: label_resolution.relocations,
196 prog_is_static: program_is_static,
197 arch,
198 debug_sections: Vec::default(),
199 })
200 }
201}
202
203#[derive(Default)]
204struct OptimizationOutcome {
205 labels_to_remove: HashSet<String>,
206 errors: Vec<CompileError>,
207}
208
209fn run_optimizations(ast: &mut AST, config: &OptimizationConfig) -> OptimizationOutcome {
210 let OptimizationConfig::Enabled { cfg_dump_dir } = config else {
211 return OptimizationOutcome::default();
212 };
213
214 let canonicalized_targets = optimizer::canonicalize_control_flow_targets(&mut ast.nodes);
218 let labels_to_remove = canonicalized_targets.labels_to_remove;
219 let mut errors = Vec::new();
220
221 if canonicalized_targets.errors.is_empty() {
222 if let Some(dump_dir) = cfg_dump_dir.as_deref() {
223 let mut dump_errors = Vec::new();
224 if let Err(error) = std::fs::create_dir_all(dump_dir) {
225 dump_errors.push((dump_dir.to_path_buf(), error));
226 optimizer::eliminate_unreachable_functions(ast);
227 } else {
228 optimizer::eliminate_unreachable_functions_with_observer(ast, |stage, cfg| {
229 let path = dump_dir.join(stage.file_name());
230 if let Err(error) = std::fs::write(&path, sbpf_analyze::dump_cfg(cfg)) {
231 dump_errors.push((path, error));
232 }
233 });
234 }
235 for (path, error) in dump_errors {
236 errors.push(CompileError::BytecodeError {
237 error: format!("failed to write CFG dump '{}': {error}", path.display()),
238 span: 0..0,
239 custom_label: None,
240 });
241 }
242 } else {
243 optimizer::eliminate_unreachable_functions(ast);
244 }
245 }
246
247 OptimizationOutcome {
248 labels_to_remove,
249 errors,
250 }
251}
252
253#[derive(Default)]
254struct LabelResolution {
255 dynamic_symbols: DynamicSymbolMap,
256 relocations: RelDynMap,
257 errors: Vec<CompileError>,
258}
259
260fn resolve_label_references(
261 ast: &mut AST,
262 arch: SbpfArch,
263 program_is_static: bool,
264 label_offset_map: &LabelOffsetMap,
265 numeric_labels: &[NumericLabel],
266) -> LabelResolution {
267 let mut relocations = RelDynMap::new();
268 let mut dynamic_symbols = DynamicSymbolMap::new();
269 let mut errors = Vec::new();
270
271 for node in ast.nodes.iter_mut() {
273 if let ASTNode::Instruction {
274 instruction: inst,
275 offset,
276 } = node
277 && inst.is_syscall()
278 && let Some(Either::Left(syscall_name)) = &inst.imm
279 {
280 let syscall_name = syscall_name.clone();
281 if arch.is_v3() {
282 inst.src = Some(Register { n: 0 });
284 inst.imm = Some(Either::Right(Number::Int(murmur3_32(&syscall_name) as i64)));
285 } else {
286 inst.src = Some(Register { n: 1 });
288 inst.imm = Some(Either::Right(Number::Int(-1)));
289
290 relocations.add_rel_dyn(*offset, RelocationType::RSbfSyscall, syscall_name.clone());
292 dynamic_symbols.add_call_target(syscall_name.clone(), *offset);
293 }
294 }
295 }
296
297 for (idx, node) in ast.nodes.iter_mut().enumerate() {
298 if let ASTNode::Instruction {
299 instruction: inst,
300 offset,
301 ..
302 } = node
303 {
304 if inst.is_jump()
306 && let Some(Either::Left(label)) = &inst.off
307 {
308 let target_offset = if let Some(offset) = label_offset_map.get(label) {
309 Some(*offset)
310 } else {
311 AST::resolve_numeric_label(label, idx, numeric_labels)
313 };
314
315 if let Some(target_offset) = target_offset {
316 let rel_offset = (target_offset as i64 - *offset as i64) / 8 - 1;
317 inst.off = Some(Either::Right(rel_offset as i16));
318 } else {
319 errors.push(CompileError::UndefinedLabel {
320 label: label.clone(),
321 span: inst.span.clone(),
322 custom_label: None,
323 });
324 }
325 } else if inst.opcode == Opcode::Call
326 && let Some(Either::Left(label)) = &inst.imm
327 && let Some(target_offset) = label_offset_map.get(label)
328 {
329 let rel_offset = (*target_offset as i64 - *offset as i64) / 8 - 1;
330 inst.src = Some(Register { n: 1 });
331 inst.imm = Some(Either::Right(Number::Int(rel_offset)));
332 }
333
334 if inst.opcode == Opcode::Lddw
335 && let Some(Either::Left(name)) = &inst.imm
336 {
337 let label = name.clone();
338 if !arch.is_v3() {
340 relocations.add_rel_dyn(*offset, RelocationType::RSbf64Relative, label.clone());
341 }
342
343 if let Some(target_offset) = label_offset_map.get(&label) {
344 let abs_offset = if arch.is_v3() {
345 if *target_offset >= ast.text_size {
346 (ProgramHeader::V3_RODATA_VADDR + *target_offset - ast.text_size) as i64
347 } else {
348 (ProgramHeader::V3_BYTECODE_VADDR + *target_offset) as i64
349 }
350 } else {
351 let ph_count = if program_is_static { 1 } else { 3 };
352 let ph_offset = 64 + (ph_count as u64 * 56) as i64;
353 *target_offset as i64 + ph_offset
354 };
355 inst.imm = Some(Either::Right(Number::Addr(abs_offset)));
357 } else {
358 errors.push(CompileError::UndefinedLabel {
359 label: name.clone(),
360 span: inst.span.clone(),
361 custom_label: None,
362 });
363 }
364 }
365 }
366 }
367
368 let entry_label = ast.nodes.iter().find_map(|node| {
370 if let ASTNode::GlobalDecl { global_decl } = node {
371 Some(global_decl.entry_label.clone())
372 } else {
373 None
374 }
375 });
376 if let Some(entry_label) = entry_label
377 && let Some(offset) = label_offset_map.get(&entry_label)
378 {
379 dynamic_symbols.add_entry_point(entry_label, *offset);
380 }
381
382 LabelResolution {
383 dynamic_symbols,
384 relocations,
385 errors,
386 }
387}
388
389fn label_offset_map(ast: &AST) -> (LabelOffsetMap, Vec<NumericLabel>) {
390 let mut label_offset_map = HashMap::new();
391 let mut numeric_labels = Vec::new();
392
393 for (idx, node) in ast.nodes.iter().enumerate() {
394 if let ASTNode::Label { label, offset } = node {
395 label_offset_map.insert(label.name.clone(), *offset);
396 numeric_labels.push((label.name.clone(), *offset, idx));
397 }
398 }
399
400 for node in &ast.rodata_nodes {
401 if let ASTNode::ROData { rodata, offset } = node {
402 label_offset_map.insert(rodata.name.clone(), *offset + ast.text_size);
403 }
404 }
405
406 (label_offset_map, numeric_labels)
407}
408
409#[cfg(test)]
410mod tests {
411 use {
412 super::*,
413 crate::{astnode::Label, parser::Token},
414 };
415
416 #[test]
417 fn test_ast_new() {
418 let ast = AST::new();
419 assert!(ast.nodes.is_empty());
420 assert!(ast.rodata_nodes.is_empty());
421 assert_eq!(ast.text_size, 0);
422 assert_eq!(ast.rodata_size, 0);
423 }
424
425 #[test]
426 fn test_ast_set_sizes() {
427 let mut ast = AST::new();
428 ast.set_text_size(100);
429 ast.set_rodata_size(50);
430 assert_eq!(ast.text_size, 100);
431 assert_eq!(ast.rodata_size, 50);
432 }
433
434 #[test]
435 fn test_get_instruction_at_offset() {
436 let mut ast = AST::new();
437 ast.nodes
438 .push(instruction_node(Opcode::Exit, 0, None, None));
439
440 let found = ast.get_instruction_at_offset(0);
441 assert!(found.is_some());
442 assert_eq!(found.unwrap().opcode, Opcode::Exit);
443
444 let not_found = ast.get_instruction_at_offset(8);
445 assert!(not_found.is_none());
446 }
447
448 #[test]
449 fn test_get_rodata_at_offset() {
450 let mut ast = AST::new();
451 let rodata = ROData {
452 name: "data".to_string(),
453 args: vec![
454 Token::Directive("ascii".to_string(), 0..5),
455 Token::StringLiteral("test".to_string(), 6..12),
456 ],
457 span: 0..12,
458 };
459 ast.rodata_nodes.push(ASTNode::ROData {
460 rodata: rodata.clone(),
461 offset: 0,
462 });
463
464 let found = ast.get_rodata_at_offset(0);
465 assert!(found.is_some());
466 assert_eq!(found.unwrap().name, "data");
467 }
468
469 #[test]
470 fn test_resolve_numeric_label_forward() {
471 let numeric_labels = vec![("1".to_string(), 16, 2), ("2".to_string(), 32, 4)];
472
473 let result = AST::resolve_numeric_label("1f", 0, &numeric_labels);
474 assert_eq!(result, Some(16));
475
476 let result = AST::resolve_numeric_label("2f", 3, &numeric_labels);
477 assert_eq!(result, Some(32));
478 }
479
480 #[test]
481 fn test_resolve_numeric_label_backward() {
482 let numeric_labels = vec![("1".to_string(), 16, 2), ("2".to_string(), 32, 4)];
483
484 let result = AST::resolve_numeric_label("1b", 3, &numeric_labels);
485 assert_eq!(result, Some(16));
486
487 let result = AST::resolve_numeric_label("2b", 5, &numeric_labels);
488 assert_eq!(result, Some(32));
489 }
490
491 #[test]
492 fn test_canonicalize_numeric_jump_target_to_label() {
493 let mut nodes = vec![
494 instruction_node(Opcode::Ja, 0, Some(Either::Left("1f".to_string())), None),
495 label_node("1", 8),
496 instruction_node(Opcode::Exit, 8, None, None),
497 ];
498
499 let canonicalized = optimizer::canonicalize_control_flow_targets(&mut nodes);
500
501 assert!(canonicalized.errors.is_empty());
502 assert!(canonicalized.labels_to_remove.is_empty());
503 let ASTNode::Instruction { instruction, .. } = &nodes[0] else {
504 panic!("expected instruction");
505 };
506 assert_eq!(instruction.off, Some(Either::Left("1".to_string())));
507 }
508
509 #[test]
510 fn test_canonicalize_relative_jump_target_to_synthetic_label() {
511 let mut nodes = vec![
512 instruction_node(Opcode::Ja, 0, Some(Either::Right(1)), None),
513 instruction_node(Opcode::Exit, 8, None, None),
514 instruction_node(Opcode::Exit, 16, None, None),
515 ];
516
517 let canonicalized = optimizer::canonicalize_control_flow_targets(&mut nodes);
518
519 assert!(canonicalized.errors.is_empty());
520 assert_eq!(canonicalized.labels_to_remove.len(), 1);
521
522 let ASTNode::Instruction { instruction, .. } = &nodes[0] else {
523 panic!("expected instruction");
524 };
525 let Some(Either::Left(label)) = &instruction.off else {
526 panic!("expected canonical label target");
527 };
528 assert!(label.starts_with("temp_"));
529 assert!(canonicalized.labels_to_remove.contains(label));
530 assert!(
531 matches!(&nodes[2], ASTNode::Label { label: target_label, offset }
532 if target_label.name == *label && *offset == 16)
533 );
534 }
535
536 #[test]
537 fn test_canonicalize_rejects_invalid_relative_jump_target() {
538 let mut nodes = vec![
539 instruction_node(Opcode::Ja, 0, Some(Either::Right(1)), None),
540 instruction_node(Opcode::Exit, 8, None, None),
541 ];
542
543 let canonicalized = optimizer::canonicalize_control_flow_targets(&mut nodes);
544
545 assert_eq!(canonicalized.errors.len(), 1);
546 assert!(canonicalized.labels_to_remove.is_empty());
547 assert_eq!(nodes.len(), 2);
548 let ASTNode::Instruction { instruction, .. } = &nodes[0] else {
549 panic!("expected instruction");
550 };
551 assert_eq!(instruction.off, Some(Either::Right(1)));
552 }
553
554 #[test]
555 fn test_dce_recomputes_relative_call_target_after_removing_code() {
556 let mut ast = AST::new();
557 ast.add_function_entry("entrypoint".to_string());
558 ast.add_function_entry("dead".to_string());
559 ast.add_function_entry("target".to_string());
560 ast.nodes = vec![
561 label_node("entrypoint", 0),
562 internal_call_node(0, 2),
563 instruction_node(Opcode::Exit, 8, None, None),
564 label_node("dead", 16),
565 instruction_node(Opcode::Exit, 16, None, None),
566 label_node("target", 24),
567 instruction_node(Opcode::Exit, 24, None, None),
568 ];
569 ast.set_text_size(32);
570
571 let program_layout =
572 build_program(ast, SbpfArch::V0, OptimizationConfig::enabled()).unwrap();
573 let nodes = program_layout.code_section.get_nodes();
574
575 assert_eq!(
576 nodes
577 .iter()
578 .filter(|node| matches!(node, ASTNode::Instruction { .. }))
579 .count(),
580 3
581 );
582 assert!(matches!(
583 &nodes[1],
584 ASTNode::Instruction { instruction, offset }
585 if instruction.opcode == Opcode::Call
586 && instruction.src == Some(Register { n: 1 })
587 && instruction.imm == Some(Either::Right(Number::Int(1)))
588 && *offset == 0
589 ));
590 assert!(matches!(
591 &nodes[3],
592 ASTNode::Label { label, offset } if label.name == "target" && *offset == 16
593 ));
594 }
595
596 #[test]
597 fn test_optimize_ast_removes_temp_jump_target_labels() {
598 let mut ast = AST::new();
599 ast.nodes = vec![
600 instruction_node(Opcode::Ja, 0, Some(Either::Right(1)), None),
601 instruction_node(Opcode::Exit, 8, None, None),
602 instruction_node(Opcode::Exit, 16, None, None),
603 ];
604 ast.set_text_size(24);
605 ast.set_rodata_size(0);
606
607 let result = build_program(ast, SbpfArch::V0, OptimizationConfig::enabled());
608
609 assert!(result.is_ok());
610 let program_layout = result.unwrap();
611 assert!(!program_layout.code_section.get_nodes().iter().any(|node| {
612 matches!(node, ASTNode::Label { label, .. }
613 if label
614 .name
615 .starts_with("temp_"))
616 }));
617 let ASTNode::Instruction { instruction, .. } = &program_layout.code_section.get_nodes()[0]
618 else {
619 panic!("expected instruction");
620 };
621 assert_eq!(instruction.off, Some(Either::Right(1)));
622 }
623
624 #[test]
625 fn test_build_program_simple() {
626 for arch in [SbpfArch::V0, SbpfArch::V3] {
627 let mut ast = AST::new();
628 ast.nodes
629 .push(instruction_node(Opcode::Exit, 0, None, None));
630 ast.set_text_size(8);
631 ast.set_rodata_size(0);
632
633 let result = build_program(ast, arch, OptimizationConfig::default());
634 assert!(result.is_ok());
635 let parse_result = result.unwrap();
636 assert!(parse_result.prog_is_static);
637 }
638 }
639
640 #[test]
641 fn test_build_program_undefined_label_error() {
642 for arch in [SbpfArch::V0, SbpfArch::V3] {
643 let mut ast = AST::new();
644 ast.nodes.push(instruction_node(
645 Opcode::Ja,
646 0,
647 Some(Either::Left("undefined_label".to_string())),
648 None,
649 ));
650 ast.set_text_size(8);
651
652 let result = build_program(ast, arch, OptimizationConfig::default());
653 assert!(result.is_err());
654 }
655 }
656
657 #[test]
658 fn test_build_program_static_syscalls_no_relocation() {
659 let mut ast = AST::new();
660
661 ast.nodes.push(instruction_node(
662 Opcode::Call,
663 0,
664 None,
665 Some(Either::Left("sol_log_".to_string())),
666 ));
667 ast.nodes
668 .push(instruction_node(Opcode::Exit, 8, None, None));
669
670 ast.set_text_size(16);
671 ast.set_rodata_size(0);
672
673 let result = build_program(ast, SbpfArch::V3, OptimizationConfig::default());
674 assert!(result.is_ok());
675 let parse_result = result.unwrap();
676
677 assert!(parse_result.prog_is_static);
678 assert!(parse_result.relocation_data.get_rel_dyns().is_empty());
679 }
680
681 #[test]
682 fn test_build_program_dynamic_syscalls_with_relocation() {
683 let mut ast = AST::new();
684
685 ast.nodes.push(instruction_node(
686 Opcode::Call,
687 0,
688 None,
689 Some(Either::Left("sol_log_".to_string())),
690 ));
691 ast.nodes
692 .push(instruction_node(Opcode::Exit, 8, None, None));
693
694 ast.set_text_size(16);
695 ast.set_rodata_size(0);
696
697 let result = build_program(ast, SbpfArch::V0, OptimizationConfig::default());
698 assert!(result.is_ok());
699 let parse_result = result.unwrap();
700
701 assert!(!parse_result.prog_is_static);
702 assert!(!parse_result.relocation_data.get_rel_dyns().is_empty());
703 }
704
705 fn label_node(name: &str, offset: u64) -> ASTNode {
706 ASTNode::Label {
707 label: Label {
708 name: name.to_string(),
709 span: 0..0,
710 },
711 offset,
712 }
713 }
714
715 fn instruction_node(
716 opcode: Opcode,
717 offset: u64,
718 off: Option<Either<String, i16>>,
719 imm: Option<Either<String, Number>>,
720 ) -> ASTNode {
721 ASTNode::Instruction {
722 instruction: Instruction {
723 opcode,
724 dst: None,
725 src: None,
726 off,
727 imm,
728 span: 0..0,
729 },
730 offset,
731 }
732 }
733
734 fn internal_call_node(offset: u64, relative_offset: i64) -> ASTNode {
735 ASTNode::Instruction {
736 instruction: Instruction {
737 opcode: Opcode::Call,
738 dst: None,
739 src: Some(Register { n: 1 }),
740 off: None,
741 imm: Some(Either::Right(Number::Int(relative_offset))),
742 span: 0..0,
743 },
744 offset,
745 }
746 }
747}