sbpf_assembler/
ast.rs

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

sbpf_assembler/
ast.rs
