sbpf_assembler/
ast.rs

1use crate::astnode::{ASTNode, ROData};
2use crate::instruction::Instruction;
3use crate::dynsym::{RelDynMap, DynamicSymbolMap, RelocationType};
4use crate::lexer::{ImmediateValue, Token};
5use crate::opcode::Opcode;
6use crate::parser::ParseResult;
7use crate::section::{DataSection, CodeSection};
8use crate::CompileError;
9
10use std::collections::HashMap;
11
12pub struct AST {
13    pub nodes: Vec<ASTNode>,
14    pub rodata_nodes: Vec<ASTNode>,
15
16    pub entry_label: Option<String>,
17    text_size: u64,
18    rodata_size: u64,
19}
20
21impl AST {
22    pub fn new() -> Self {
23        Self { nodes: Vec::new(), rodata_nodes: Vec::new(), entry_label: None, text_size: 0, rodata_size: 0 }
24    }
25    //
26    pub fn set_text_size(&mut self, text_size: u64) {
27        self.text_size = text_size;
28    }
29    //
30    pub fn set_rodata_size(&mut self, rodata_size: u64) {
31        self.rodata_size = rodata_size;
32    }
33    //
34    pub fn get_instruction_at_offset(&mut self, offset: u64) -> Option<&mut Instruction> {
35        self.nodes.iter_mut().find(|node| match node {
36            ASTNode::Instruction { instruction: _, offset: inst_offset, .. } => offset == *inst_offset,
37            _ => false,
38        }).map(|node| match node {
39            ASTNode::Instruction { instruction, .. } => instruction,
40            _ => panic!("Expected Instruction node"),
41        })
42    }
43    //
44    pub fn get_rodata_at_offset(&self, offset: u64) -> Option<&ROData> {
45        self.rodata_nodes.iter().find(|node| match node {
46            ASTNode::ROData { rodata: _, offset: rodata_offset, .. } => offset == *rodata_offset,
47            _ => false,
48        }).map(|node| match node {
49            ASTNode::ROData { rodata, .. } => rodata,
50            _ => panic!("Expected ROData node"),
51        })
52    }
53    //
54    pub fn build_program(&mut self) -> Result<ParseResult, Vec<CompileError>> {
55        let mut label_offset_map : HashMap<String, u64> = HashMap::new();
56
57        // iterate through text labels and rodata labels and find the pair
58        // of each label and offset
59        for node in &self.nodes {
60            match node {
61                ASTNode::Label { label, offset } => {
62                    label_offset_map.insert(label.name.clone(), *offset);
63                }
64                _ => {}
65            }
66        }
67
68        for node in &self.rodata_nodes {
69            match node {
70                ASTNode::ROData { rodata, offset } => {
71                    label_offset_map.insert(rodata.name.clone(), *offset + self.text_size);
72                }
73                _ => {}
74            }
75        }
76
77        // 1. resolve labels in the intruction nodes for lddw and jump
78        // 2. find relocation information
79
80        let mut program_is_static = true;
81        let mut relocations = RelDynMap::new();
82        let mut dynamic_symbols = DynamicSymbolMap::new();
83
84        let mut errors = Vec::new();
85
86        for node in &mut self.nodes {
87            match node {
88                ASTNode::Instruction { instruction: inst, offset, .. } => {
89                    // For jump/call instructions, replace label with relative offsets
90                    if inst.is_jump() || inst.opcode == Opcode::Call {
91                        if let Some(Token::Identifier(label, span)) = inst.operands.last() {
92                            let label = label.clone();
93                            if let Some(target_offset) = label_offset_map.get(&label) {
94                                let rel_offset = (*target_offset as i64 - *offset as i64) / 8 - 1;
95                                let last_idx = inst.operands.len() - 1;
96                                inst.operands[last_idx] = Token::ImmediateValue(ImmediateValue::Int(rel_offset), span.clone());
97                            } else if inst.is_jump() {
98                                // only error out unresolved jump labels, since call 
99                                // labels could exist externally
100                                errors.push(CompileError::UndefinedLabel { label: label.clone(), span: span.clone(), custom_label: None });
101                            }
102                        }
103                    }
104                    // This has to be done before resolving lddw labels since lddw 
105                    // operand needs to be absolute offset values
106                    if inst.needs_relocation() {
107                        program_is_static = false;
108                        let (reloc_type, label) = inst.get_relocation_info();
109                        relocations.add_rel_dyn(*offset, reloc_type, label.clone());
110                        if reloc_type == RelocationType::RSbfSyscall {
111                            dynamic_symbols.add_call_target(label.clone(), *offset);
112                        }
113                    }
114                    if inst.opcode == Opcode::Lddw {
115                        if let Some(Token::Identifier(name, span)) = inst.operands.last() {
116                            let label = name.clone();
117                            if let Some(target_offset) = label_offset_map.get(&label) {
118                                // actually lddw with label makes a program dynamic, so
119                                // we should be able to hard code ph_offset
120                                let ph_count = if program_is_static { 1 } else { 3 };
121                                let ph_offset = 64 + (ph_count as u64 * 56) as i64;
122                                let abs_offset = *target_offset as i64 + ph_offset;
123                                // Replace label with immediate value
124                                let last_idx = inst.operands.len() - 1;
125                                inst.operands[last_idx] = Token::ImmediateValue(ImmediateValue::Addr(abs_offset), span.clone());
126                            }  else {
127                                errors.push(CompileError::UndefinedLabel { label: name.clone(), span: span.clone(), custom_label: None });
128                            }
129                        }
130                    }
131                }
132                _ => {}
133            }
134        }
135
136        // Set entry point offset if an entry label was specified
137        if let Some(entry_label) = &self.entry_label {
138            if let Some(offset) = label_offset_map.get(entry_label) {
139                dynamic_symbols.add_entry_point(entry_label.clone(), *offset);
140            }
141        }
142
143        if !errors.is_empty() {
144            return Err(errors);
145        } else {
146            Ok(ParseResult {
147                code_section: CodeSection::new(std::mem::take(&mut self.nodes), self.text_size),
148                data_section: DataSection::new(std::mem::take(&mut self.rodata_nodes), self.rodata_size),
149                dynamic_symbols: dynamic_symbols,
150                relocation_data: relocations,
151                prog_is_static: program_is_static,
152            })
153        }
154    }
155}
sbpf_assembler/ast.rs

sbpf_assembler/
ast.rs
