sbpf_assembler/
astnode.rs

1use crate::opcode::Opcode;
2use crate::instruction::Instruction;
3use crate::lexer::{Token, ImmediateValue};
4use crate::debuginfo::{DebugInfo, RegisterHint, RegisterType};
5use std::collections::HashMap;
6use std::ops::Range;
7use crate::errors::CompileError;
8
9#[derive(Debug, Clone)]
10pub enum ASTNode {
11    // only present in the AST
12    Directive {
13        directive: Directive,
14    },
15    GlobalDecl {
16        global_decl: GlobalDecl,
17    },
18    EquDecl {
19        equ_decl: EquDecl,
20    },
21    ExternDecl {
22        extern_decl: ExternDecl,
23    },
24    RodataDecl {
25        rodata_decl: RodataDecl,
26    },
27    Label {
28        label: Label,
29        offset: u64,
30    },
31    // present in the bytecode
32    ROData {
33        rodata: ROData,
34        offset: u64,
35    },
36    Instruction {
37        instruction: Instruction,
38        offset: u64,
39    },
40
41}
42
43#[derive(Debug, Clone)]
44pub struct Directive {
45    pub name: String,
46    pub args: Vec<Token>,
47    pub span: Range<usize>,
48}
49
50#[derive(Debug, Clone)]
51pub struct GlobalDecl {
52    pub entry_label: String,
53    pub span: Range<usize>,
54}
55
56impl GlobalDecl {
57    pub fn get_entry_label(&self) -> String {
58        self.entry_label.clone()
59    }
60}
61
62#[derive(Debug, Clone)]
63pub struct EquDecl {
64    pub name: String,
65    pub value: Token,
66    pub span: Range<usize>,
67}
68
69impl EquDecl {
70    pub fn get_name(&self) -> String {
71        self.name.clone()
72    }
73    pub fn get_val(&self) -> ImmediateValue {
74        match &self.value {
75            Token::ImmediateValue(val, _) => val.clone(),
76            _ => panic!("Invalid Equ declaration"),
77        }
78    }
79}
80
81#[derive(Debug, Clone)]
82pub struct ExternDecl {
83    pub args: Vec<Token>,
84    pub span: Range<usize>,
85}
86
87#[derive(Debug, Clone)]
88pub struct RodataDecl {
89    pub span: Range<usize>,
90}
91
92#[derive(Debug, Clone)]
93pub struct Label {
94    pub name: String,
95    pub span: Range<usize>,
96}
97
98#[derive(Debug, Clone)]
99pub struct ROData {
100    pub name: String,
101    pub args: Vec<Token>,
102    pub span: Range<usize>,
103}
104
105impl ROData {
106    /// Validates that an immediate value is within the specified range
107    fn validate_immediate_range(
108        value: &ImmediateValue,
109        min: i64,
110        max: i64,
111        span: Range<usize>,
112    ) -> Result<(), CompileError> {
113        match value {
114            ImmediateValue::Int(val) => {
115                if *val < min || *val > max {
116                    return Err(CompileError::OutOfRangeLiteral { span, custom_label: None });
117                }
118            }
119            ImmediateValue::Addr(val) => {
120                if *val < min || *val > max {
121                    return Err(CompileError::OutOfRangeLiteral { span, custom_label: None });
122                }
123            }
124        }
125        Ok(())
126    }
127
128    pub fn get_size(&self) -> u64 {
129        let size: u64;
130        match (
131            &self.args[0],
132            &self.args[1],
133        ) {
134            (Token::Directive(_, _), Token::StringLiteral(s, _)) => {
135                size = s.len() as u64;
136            }
137            (Token::Directive(directive, _), Token::VectorLiteral(values, _)) => {
138                match directive.as_str() {
139                    "byte" => {
140                        size = values.len() as u64 * 1;
141                    }
142                    "short" => {
143                        size = values.len() as u64 * 2;
144                    }
145                    "int" | "long" => {
146                        size = values.len() as u64 * 4;
147                    }
148                    "quad" => {
149                        size = values.len() as u64 * 8;
150                    }
151                    _ => panic!("Invalid ROData declaration"),
152                }
153            }
154            _ => panic!("Invalid ROData declaration"),
155        }
156        size
157    }
158    pub fn verify(&self) -> Result<(), CompileError> {
159        match (
160            &self.args[0],
161            &self.args[1],
162        ) {
163            (Token::Directive(directive, directive_span), Token::StringLiteral(_, _)) => {
164                if directive.as_str() != "ascii" {
165                    return Err(CompileError::InvalidRODataDirective { span: directive_span.clone(), custom_label: None });
166                }
167            }
168            (Token::Directive(directive, directive_span), Token::VectorLiteral(values, vector_literal_span)) => {
169                match directive.as_str() {
170                    "byte" => {
171                        for value in values {
172                            Self::validate_immediate_range(value, i8::MIN as i64, i8::MAX as i64, vector_literal_span.clone())?;
173                        }
174                    }
175                    "short" => {
176                        for value in values {
177                            Self::validate_immediate_range(value, i16::MIN as i64, i16::MAX as i64, vector_literal_span.clone())?;
178                        }
179                    }
180                    "int" | "long" => {
181                        for value in values {
182                            Self::validate_immediate_range(value, i32::MIN as i64, i32::MAX as i64, vector_literal_span.clone())?;
183                        }
184                    }
185                    "quad" => {
186                        for value in values {
187                            Self::validate_immediate_range(value, i64::MIN as i64, i64::MAX as i64, vector_literal_span.clone())?;
188                        }
189                    }
190                _ => {
191                        return Err(CompileError::InvalidRODataDirective { span: directive_span.clone(), custom_label: None });
192                    }
193                }
194            }
195            _ => {
196                return Err(CompileError::InvalidRodataDecl { span: self.span.clone(), custom_label: None });
197            }
198        }
199        Ok(())
200    }
201}
202
203impl ASTNode {
204    pub fn bytecode_with_debug_map(&self) -> Option<(Vec<u8>, HashMap<u64, DebugInfo>)> {
205        match self {
206            ASTNode::Instruction { instruction: Instruction { opcode, operands, span }, offset } => {
207                let mut bytes = Vec::new();
208                let mut debug_map = HashMap::new();
209                let mut debug_info = DebugInfo::new(span.clone());
210                bytes.push(opcode.to_bytecode());  // 1 byte opcode
211                
212                if *opcode == Opcode::Call {
213                    bytes.extend_from_slice(&[0x10, 0x00, 0x00]);
214                    if let Some(Token::ImmediateValue(imm, _)) = operands.last() {
215                        let imm32 = match imm {
216                            ImmediateValue::Int(val) => *val as i32,
217                            ImmediateValue::Addr(val) => *val as i32,
218                        };
219                        bytes.extend_from_slice(&imm32.to_le_bytes());
220                    } else {
221                        // external calls
222                        bytes.extend_from_slice(&[0xFF, 0xFF, 0xFF, 0xFF]);
223                    } 
224                } else if *opcode == Opcode::Lddw {
225                    match &operands[..] {
226                        [Token::Register(reg, _), Token::ImmediateValue(imm, _)] => {
227                            // 1 byte register number (strip 'r' prefix)
228                            bytes.push(*reg);
229                            
230                            // 2 bytes of zeros (offset/reserved)
231                            bytes.extend_from_slice(&[0, 0]);
232
233                            // 8 bytes immediate value in little-endian
234                            let imm64 = match imm {
235                                ImmediateValue::Int(val) => *val as i64,
236                                ImmediateValue::Addr(val) => *val as i64,
237                            };
238                            bytes.extend_from_slice(&imm64.to_le_bytes()[..4]);
239                            bytes.extend_from_slice(&[0, 0, 0, 0]);
240                            bytes.extend_from_slice(&imm64.to_le_bytes()[4..8]);
241                        }
242                        _ => {}
243                    }
244                } else {
245                    match &operands[..] {
246                        [Token::ImmediateValue(imm, _)] => {
247                            // 1 byte of zeros (no register)
248                            bytes.push(0);
249                            
250                            if *opcode == Opcode::Ja {
251                                // 2 bytes immediate value in little-endian for 'ja'
252                                let imm16 = match imm {
253                                    ImmediateValue::Int(val) => *val as i16,
254                                    ImmediateValue::Addr(val) => *val as i16,
255                                };
256                                bytes.extend_from_slice(&imm16.to_le_bytes());
257                            } else {
258                                // 4 bytes immediate value in little-endian
259                                let imm32 = match imm {
260                                    ImmediateValue::Int(val) => *val as i32,
261                                    ImmediateValue::Addr(val) => *val as i32,
262                                };
263                                bytes.extend_from_slice(&imm32.to_le_bytes());
264                            }
265                        },
266
267                        [Token::Register(reg, _)] => {
268                            bytes.push(*reg);
269                            bytes.extend_from_slice(&[0, 0, 0, 0, 0, 0]);
270                        },
271
272                        [Token::Register(reg, _), Token::ImmediateValue(imm, _)] => {
273                            // 1 byte register number (strip 'r' prefix)
274                            bytes.push(*reg);
275                            
276                            // 2 bytes of zeros (offset/reserved)
277                            bytes.extend_from_slice(&[0, 0]);
278                            
279                            // 4 bytes immediate value in little-endian
280                            let imm32 = match imm {
281                                ImmediateValue::Int(val) => *val as i32,
282                                ImmediateValue::Addr(val) => {
283                                    debug_info.register_hint = RegisterHint {
284                                        register: *reg as usize,
285                                        register_type: RegisterType::Addr
286                                    };
287                                    *val as i32
288                                }
289                            };
290                            bytes.extend_from_slice(&imm32.to_le_bytes());
291                        },
292
293                        [Token::Register(reg, _), Token::ImmediateValue(imm, _), Token::ImmediateValue(offset, _)] => {
294                            // 1 byte register number (strip 'r' prefix)
295                            bytes.push(*reg);
296                            
297                            // 2 bytes of offset in little-endian
298                            let offset16 = match offset {
299                                ImmediateValue::Int(val) => *val as u16,
300                                ImmediateValue::Addr(val) => *val as u16,
301                            };
302                            bytes.extend_from_slice(&offset16.to_le_bytes());
303                            
304                            // 4 bytes immediate value in little-endianß
305                            let imm32 = match imm {
306                                ImmediateValue::Int(val) => *val as i32,
307                                ImmediateValue::Addr(val) => {
308                                    debug_info.register_hint = RegisterHint {
309                                        register: *reg as usize,
310                                        register_type: RegisterType::Addr
311                                    };
312                                    *val as i32
313                                }
314                            };
315                            bytes.extend_from_slice(&imm32.to_le_bytes());
316                        },                    
317                        
318                        [Token::Register(dst, _), Token::Register(src, _)] => {
319                            // Convert register strings to numbers
320                            let dst_num = dst;
321                            let src_num = src;
322                            
323                            // Combine src and dst into a single byte (src in high nibble, dst in low nibble)
324                            let reg_byte = (src_num << 4) | dst_num;
325                            bytes.push(reg_byte);
326                        },
327                        [Token::Register(dst, _), Token::Register(reg, _), Token::ImmediateValue(offset, _)] => {
328                            // Combine base register and destination register into a single byte
329                            let reg_byte = (reg << 4) | dst;
330                            bytes.push(reg_byte);
331                            
332                            // Add the offset as a 16-bit value in little-endian
333                            let offset16 = match offset {
334                                ImmediateValue::Int(val) => *val as u16,
335                                ImmediateValue::Addr(val) => *val as u16,
336                            };
337                            bytes.extend_from_slice(&offset16.to_le_bytes());
338                        },
339                        [Token::Register(reg, _), Token::ImmediateValue(offset, _), Token::Register(dst, _)] => {
340                            // Combine base register and destination register into a single byte
341                            let reg_byte = (dst << 4) | reg;
342                            bytes.push(reg_byte);
343                            
344                            // Add the offset as a 16-bit value in little-endian
345                            let offset16 = match offset {
346                                ImmediateValue::Int(val) => *val as u16,
347                                ImmediateValue::Addr(val) => *val as u16,
348                            };
349                            bytes.extend_from_slice(&offset16.to_le_bytes());
350                        }
351                        
352                        _ => {}
353                    }
354                }
355
356                // Add padding to make it 8 or 16 bytes depending on opcode
357                let target_len = if *opcode == Opcode::Lddw { 16 } else { 8 };
358                while bytes.len() < target_len {
359                    bytes.push(0);
360                }
361
362                debug_map.insert(*offset, debug_info);
363                
364                Some((bytes, debug_map))
365            },
366            ASTNode::ROData { rodata: ROData { name: _, args, .. }, .. } => {
367                let mut bytes = Vec::new();
368                let debug_map = HashMap::<u64, DebugInfo>::new();
369                match (
370                    &args[0],
371                    &args[1],
372                ) {
373                    (Token::Directive(_, _), Token::StringLiteral(str_literal, _)) => {
374                        let str_bytes = str_literal.as_bytes().to_vec();
375                        bytes.extend(str_bytes);
376                    } 
377                    (Token::Directive(directive, _), Token::VectorLiteral(values, _)) => {
378                        if directive == "byte" {
379                            for value in values {
380                                let imm8 = match value {
381                                    ImmediateValue::Int(val) => *val as i8,
382                                    ImmediateValue::Addr(val) => *val as i8,
383                                };
384                                bytes.extend(imm8.to_le_bytes());
385                            }
386                        } else if directive == "short" {
387                            for value in values {
388                                let imm16 = match value {
389                                    ImmediateValue::Int(val) => *val as i16,
390                                    ImmediateValue::Addr(val) => *val as i16,
391                                };
392                                bytes.extend(imm16.to_le_bytes());
393                            }
394                        } else if directive == "int" || directive == "long" {
395                            for value in values {
396                                let imm32 = match value {
397                                    ImmediateValue::Int(val) => *val as i32,
398                                    ImmediateValue::Addr(val) => *val as i32,
399                                };
400                                bytes.extend(imm32.to_le_bytes());
401                            }
402                        } else if directive == "quad" {
403                            for value in values {
404                                let imm64 = match value {
405                                    ImmediateValue::Int(val) => *val as i64,
406                                    ImmediateValue::Addr(val) => *val as i64,
407                                };
408                                bytes.extend(imm64.to_le_bytes());
409                            }
410                        } else {
411                            panic!("Invalid ROData declaration");
412                        }
413                    }
414
415                    _ => panic!("Invalid ROData declaration"),
416                }
417                Some((bytes, debug_map))
418            },
419            _ => None
420        }
421    }
422
423    // Keep the old bytecode method for backward compatibility
424    pub fn bytecode(&self) -> Option<Vec<u8>> {
425        self.bytecode_with_debug_map().map(|(bytes, _)| bytes)
426    }
427}
sbpf_assembler/astnode.rs

sbpf_assembler/
astnode.rs
