sbpf_assembler/
ast.rs

1use {
2    crate::{
3        CompileError, SbpfArch,
4        astnode::{ASTNode, ROData},
5        dynsym::{DynamicSymbolMap, RelDynMap, RelocationType},
6        parser::ParseResult,
7        section::{CodeSection, DataSection},
8    },
9    either::Either,
10    sbpf_common::{
11        inst_param::{Number, Register},
12        instruction::Instruction,
13        opcode::Opcode,
14    },
15    std::collections::HashMap,
16    syscall_map::murmur3_32,
17};
18
19#[derive(Default, Debug)]
20pub struct AST {
21    pub nodes: Vec<ASTNode>,
22    pub rodata_nodes: Vec<ASTNode>,
23
24    text_size: u64,
25    rodata_size: u64,
26}
27
28impl AST {
29    pub fn new() -> Self {
30        Self::default()
31    }
32
33    //
34    pub fn set_text_size(&mut self, text_size: u64) {
35        self.text_size = text_size;
36    }
37
38    //
39    pub fn set_rodata_size(&mut self, rodata_size: u64) {
40        self.rodata_size = rodata_size;
41    }
42
43    //
44    pub fn get_instruction_at_offset(&mut self, offset: u64) -> Option<&mut Instruction> {
45        self.nodes
46            .iter_mut()
47            .find(|node| match node {
48                ASTNode::Instruction {
49                    instruction: _,
50                    offset: inst_offset,
51                    ..
52                } => offset == *inst_offset,
53                _ => false,
54            })
55            .map(|node| match node {
56                ASTNode::Instruction { instruction, .. } => instruction,
57                _ => panic!("Expected Instruction node"),
58            })
59    }
60
61    //
62    pub fn get_rodata_at_offset(&self, offset: u64) -> Option<&ROData> {
63        self.rodata_nodes
64            .iter()
65            .find(|node| match node {
66                ASTNode::ROData {
67                    rodata: _,
68                    offset: rodata_offset,
69                    ..
70                } => offset == *rodata_offset,
71                _ => false,
72            })
73            .map(|node| match node {
74                ASTNode::ROData { rodata, .. } => rodata,
75                _ => panic!("Expected ROData node"),
76            })
77    }
78
79    /// Resolve numeric label references (like "2f" or "1b")
80    fn resolve_numeric_label(
81        label_ref: &str,
82        current_idx: usize,
83        numeric_labels: &[(String, u64, usize)],
84    ) -> Option<u64> {
85        if let Some(direction) = label_ref.chars().last()
86            && (direction == 'f' || direction == 'b')
87        {
88            let label_num = &label_ref[..label_ref.len() - 1];
89
90            if direction == 'f' {
91                // search forward from current position
92                for (name, offset, node_idx) in numeric_labels {
93                    if name == label_num && *node_idx > current_idx {
94                        return Some(*offset);
95                    }
96                }
97            } else {
98                // search backward from current position
99                for (name, offset, node_idx) in numeric_labels.iter().rev() {
100                    if name == label_num && *node_idx < current_idx {
101                        return Some(*offset);
102                    }
103                }
104            }
105        }
106        None
107    }
108
109    pub fn build_program(&mut self, arch: SbpfArch) -> Result<ParseResult, Vec<CompileError>> {
110        let mut label_offset_map: HashMap<String, u64> = HashMap::new();
111        let mut numeric_labels: Vec<(String, u64, usize)> = Vec::new();
112
113        // iterate through text labels and rodata labels and find the pair
114        // of each label and offset
115        for (idx, node) in self.nodes.iter().enumerate() {
116            if let ASTNode::Label { label, offset } = node {
117                label_offset_map.insert(label.name.clone(), *offset);
118                // Also track numeric labels separately for forward/backward resolution
119                numeric_labels.push((label.name.clone(), *offset, idx));
120            }
121        }
122
123        for node in &self.rodata_nodes {
124            if let ASTNode::ROData { rodata, offset } = node {
125                label_offset_map.insert(rodata.name.clone(), *offset + self.text_size);
126            }
127        }
128
129        // 1. resolve labels in the intruction nodes for lddw and jump
130        // 2. find relocation information
131
132        let mut relocations = RelDynMap::new();
133        let mut dynamic_symbols = DynamicSymbolMap::new();
134
135        let program_is_static = arch.is_v3()
136            || !self.nodes.iter().any(|node| {
137                matches!(node, ASTNode::Instruction { instruction: inst, .. }
138                    if inst.is_syscall()
139                    || (inst.opcode == Opcode::Lddw && matches!(&inst.imm, Some(Either::Left(_)))))
140            });
141
142        // Resolve both static and dynamic syscalls.
143        for node in self.nodes.iter_mut() {
144            if let ASTNode::Instruction {
145                instruction: inst,
146                offset,
147            } = node
148                && inst.is_syscall()
149                && let Some(Either::Left(syscall_name)) = &inst.imm
150            {
151                let syscall_name = syscall_name.clone();
152                if arch.is_v3() {
153                    // Static syscall: src = 0, imm = hash
154                    inst.src = Some(Register { n: 0 });
155                    inst.imm = Some(Either::Right(Number::Int(murmur3_32(&syscall_name) as i64)));
156                } else {
157                    // Dynamic syscall: src = 1, imm = -1
158                    inst.src = Some(Register { n: 1 });
159                    inst.imm = Some(Either::Right(Number::Int(-1)));
160
161                    // Add relocation for dynamic syscall
162                    relocations.add_rel_dyn(
163                        *offset,
164                        RelocationType::RSbfSyscall,
165                        syscall_name.clone(),
166                    );
167                    dynamic_symbols.add_call_target(syscall_name.clone(), *offset);
168                }
169            }
170        }
171
172        let mut errors = Vec::new();
173
174        for (idx, node) in self.nodes.iter_mut().enumerate() {
175            if let ASTNode::Instruction {
176                instruction: inst,
177                offset,
178                ..
179            } = node
180            {
181                // For jump/call instructions, replace label with relative offsets
182                if inst.is_jump()
183                    && let Some(Either::Left(label)) = &inst.off
184                {
185                    let target_offset = if let Some(offset) = label_offset_map.get(label) {
186                        Some(*offset)
187                    } else {
188                        // Handle numeric label references
189                        Self::resolve_numeric_label(label, idx, &numeric_labels)
190                    };
191
192                    if let Some(target_offset) = target_offset {
193                        let rel_offset = (target_offset as i64 - *offset as i64) / 8 - 1;
194                        inst.off = Some(Either::Right(rel_offset as i16));
195                    } else {
196                        errors.push(CompileError::UndefinedLabel {
197                            label: label.clone(),
198                            span: inst.span.clone(),
199                            custom_label: None,
200                        });
201                    }
202                } else if inst.opcode == Opcode::Call
203                    && let Some(Either::Left(label)) = &inst.imm
204                    && let Some(target_offset) = label_offset_map.get(label)
205                {
206                    let rel_offset = (*target_offset as i64 - *offset as i64) / 8 - 1;
207                    inst.src = Some(Register { n: 1 });
208                    inst.imm = Some(Either::Right(Number::Int(rel_offset)));
209                }
210
211                if inst.opcode == Opcode::Lddw
212                    && let Some(Either::Left(name)) = &inst.imm
213                {
214                    let label = name.clone();
215                    // Add relocation for lddw (only for v0)
216                    if !arch.is_v3() {
217                        relocations.add_rel_dyn(
218                            *offset,
219                            RelocationType::RSbf64Relative,
220                            label.clone(),
221                        );
222                    }
223
224                    if let Some(target_offset) = label_offset_map.get(&label) {
225                        let abs_offset = if arch.is_v3() {
226                            (*target_offset - self.text_size) as i64
227                        } else {
228                            let ph_count = if program_is_static { 1 } else { 3 };
229                            let ph_offset = 64 + (ph_count as u64 * 56) as i64;
230                            *target_offset as i64 + ph_offset
231                        };
232                        // Replace label with immediate value
233                        inst.imm = Some(Either::Right(Number::Addr(abs_offset)));
234                    } else {
235                        errors.push(CompileError::UndefinedLabel {
236                            label: name.clone(),
237                            span: inst.span.clone(),
238                            custom_label: None,
239                        });
240                    }
241                }
242            }
243        }
244
245        // Set entry point offset if a GlobalDecl was specified
246        let entry_label = self.nodes.iter().find_map(|node| {
247            if let ASTNode::GlobalDecl { global_decl } = node {
248                Some(global_decl.entry_label.clone())
249            } else {
250                None
251            }
252        });
253        if let Some(entry_label) = entry_label
254            && let Some(offset) = label_offset_map.get(&entry_label)
255        {
256            dynamic_symbols.add_entry_point(entry_label, *offset);
257        }
258
259        if !errors.is_empty() {
260            Err(errors)
261        } else {
262            Ok(ParseResult {
263                code_section: CodeSection::new(std::mem::take(&mut self.nodes), self.text_size),
264                data_section: DataSection::new(
265                    std::mem::take(&mut self.rodata_nodes),
266                    self.rodata_size,
267                ),
268                dynamic_symbols,
269                relocation_data: relocations,
270                prog_is_static: program_is_static,
271                arch,
272                debug_sections: Vec::default(),
273            })
274        }
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use {super::*, crate::parser::Token};
281
282    #[test]
283    fn test_ast_new() {
284        let ast = AST::new();
285        assert!(ast.nodes.is_empty());
286        assert!(ast.rodata_nodes.is_empty());
287        assert_eq!(ast.text_size, 0);
288        assert_eq!(ast.rodata_size, 0);
289    }
290
291    #[test]
292    fn test_ast_set_sizes() {
293        let mut ast = AST::new();
294        ast.set_text_size(100);
295        ast.set_rodata_size(50);
296        assert_eq!(ast.text_size, 100);
297        assert_eq!(ast.rodata_size, 50);
298    }
299
300    #[test]
301    fn test_get_instruction_at_offset() {
302        let mut ast = AST::new();
303        let inst = Instruction {
304            opcode: Opcode::Exit,
305            dst: None,
306            src: None,
307            off: None,
308            imm: None,
309            span: 0..4,
310        };
311        ast.nodes.push(ASTNode::Instruction {
312            instruction: inst,
313            offset: 0,
314        });
315
316        let found = ast.get_instruction_at_offset(0);
317        assert!(found.is_some());
318        assert_eq!(found.unwrap().opcode, Opcode::Exit);
319
320        let not_found = ast.get_instruction_at_offset(8);
321        assert!(not_found.is_none());
322    }
323
324    #[test]
325    fn test_get_rodata_at_offset() {
326        let mut ast = AST::new();
327        let rodata = ROData {
328            name: "data".to_string(),
329            args: vec![
330                Token::Directive("ascii".to_string(), 0..5),
331                Token::StringLiteral("test".to_string(), 6..12),
332            ],
333            span: 0..12,
334        };
335        ast.rodata_nodes.push(ASTNode::ROData {
336            rodata: rodata.clone(),
337            offset: 0,
338        });
339
340        let found = ast.get_rodata_at_offset(0);
341        assert!(found.is_some());
342        assert_eq!(found.unwrap().name, "data");
343    }
344
345    #[test]
346    fn test_resolve_numeric_label_forward() {
347        let numeric_labels = vec![("1".to_string(), 16, 2), ("2".to_string(), 32, 4)];
348
349        let result = AST::resolve_numeric_label("1f", 0, &numeric_labels);
350        assert_eq!(result, Some(16));
351
352        let result = AST::resolve_numeric_label("2f", 3, &numeric_labels);
353        assert_eq!(result, Some(32));
354    }
355
356    #[test]
357    fn test_resolve_numeric_label_backward() {
358        let numeric_labels = vec![("1".to_string(), 16, 2), ("2".to_string(), 32, 4)];
359
360        let result = AST::resolve_numeric_label("1b", 3, &numeric_labels);
361        assert_eq!(result, Some(16));
362
363        let result = AST::resolve_numeric_label("2b", 5, &numeric_labels);
364        assert_eq!(result, Some(32));
365    }
366
367    #[test]
368    fn test_build_program_simple() {
369        let mut ast = AST::new();
370        let inst = Instruction {
371            opcode: Opcode::Exit,
372            dst: None,
373            src: None,
374            off: None,
375            imm: None,
376            span: 0..4,
377        };
378        ast.nodes.push(ASTNode::Instruction {
379            instruction: inst,
380            offset: 0,
381        });
382        ast.set_text_size(8);
383        ast.set_rodata_size(0);
384
385        let result = ast.build_program(SbpfArch::V0);
386        assert!(result.is_ok());
387        let parse_result = result.unwrap();
388        assert!(parse_result.prog_is_static);
389    }
390
391    #[test]
392    fn test_build_program_undefined_label_error() {
393        let mut ast = AST::new();
394
395        // Jump to undefined label
396        let inst = Instruction {
397            opcode: Opcode::Ja,
398            dst: None,
399            src: None,
400            off: Some(Either::Left("undefined_label".to_string())),
401            imm: None,
402            span: 0..10,
403        };
404        ast.nodes.push(ASTNode::Instruction {
405            instruction: inst,
406            offset: 0,
407        });
408        ast.set_text_size(8);
409
410        let result = ast.build_program(SbpfArch::V0);
411        assert!(result.is_err());
412    }
413
414    #[test]
415    fn test_build_program_static_syscalls_no_relocation() {
416        let mut ast = AST::new();
417
418        let syscall_inst = Instruction {
419            opcode: Opcode::Call,
420            dst: None,
421            src: None,
422            off: None,
423            imm: Some(Either::Left("sol_log_".to_string())),
424            span: 0..8,
425        };
426        ast.nodes.push(ASTNode::Instruction {
427            instruction: syscall_inst,
428            offset: 0,
429        });
430
431        let exit_inst = Instruction {
432            opcode: Opcode::Exit,
433            dst: None,
434            src: None,
435            off: None,
436            imm: None,
437            span: 8..16,
438        };
439        ast.nodes.push(ASTNode::Instruction {
440            instruction: exit_inst,
441            offset: 8,
442        });
443
444        ast.set_text_size(16);
445        ast.set_rodata_size(0);
446
447        let result = ast.build_program(SbpfArch::V3);
448        assert!(result.is_ok());
449        let parse_result = result.unwrap();
450
451        assert!(parse_result.prog_is_static);
452        assert!(parse_result.relocation_data.get_rel_dyns().is_empty());
453    }
454
455    #[test]
456    fn test_build_program_dynamic_syscalls_with_relocation() {
457        let mut ast = AST::new();
458
459        let syscall_inst = Instruction {
460            opcode: Opcode::Call,
461            dst: None,
462            src: None,
463            off: None,
464            imm: Some(Either::Left("sol_log_".to_string())),
465            span: 0..8,
466        };
467        ast.nodes.push(ASTNode::Instruction {
468            instruction: syscall_inst,
469            offset: 0,
470        });
471
472        let exit_inst = Instruction {
473            opcode: Opcode::Exit,
474            dst: None,
475            src: None,
476            off: None,
477            imm: None,
478            span: 8..16,
479        };
480        ast.nodes.push(ASTNode::Instruction {
481            instruction: exit_inst,
482            offset: 8,
483        });
484
485        ast.set_text_size(16);
486        ast.set_rodata_size(0);
487
488        let result = ast.build_program(SbpfArch::V0);
489        assert!(result.is_ok());
490        let parse_result = result.unwrap();
491
492        assert!(!parse_result.prog_is_static);
493        assert!(!parse_result.relocation_data.get_rel_dyns().is_empty());
494    }
495}
sbpf_assembler/ast.rs

sbpf_assembler/
ast.rs
