Skip to main content

sbpf_assembler/
lib.rs

1use {anyhow::Result, codespan::Files};
2
3// Parser
4pub mod parser;
5
6// Preprocessor (include + macro expansion)
7pub mod preprocessor;
8
9// Error handling and diagnostics
10pub mod errors;
11pub mod macros;
12
13// Intermediate Representation
14pub mod ast;
15pub mod astnode;
16pub mod dynsym;
17pub mod optimizer;
18
19// ELF header, program, section
20pub mod header;
21pub mod program;
22pub mod section;
23
24// Debug info
25pub mod debug;
26
27// WASM bindings
28#[cfg(target_arch = "wasm32")]
29pub mod wasm;
30
31pub use self::{
32    ast::OptimizationConfig,
33    astnode::ASTNode,
34    debug::DebugData,
35    errors::CompileError,
36    parser::{ProgramLayout, Token, parse, parse_with_optimization},
37    preprocessor::{
38        FileResolver, FsFileResolver, MockFileResolver, PreprocessResult, preprocess,
39        source_map::{FileRegistry, SourceMap, SourceOrigin},
40    },
41    program::Program,
42};
43
44/// sBPF target architecture
45#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
46pub enum SbpfArch {
47    V0,
48    #[default]
49    V3,
50}
51
52impl SbpfArch {
53    pub fn is_v3(&self) -> bool {
54        matches!(self, SbpfArch::V3)
55    }
56
57    pub fn e_flags(&self) -> u32 {
58        match self {
59            SbpfArch::V0 => 0,
60            SbpfArch::V3 => 3,
61        }
62    }
63}
64
65/// Debug mode configuration for the assembler
66#[derive(Debug, Clone)]
67pub struct DebugMode {
68    /// Source filename for debug info
69    pub filename: String,
70    /// Source directory for debug info
71    pub directory: String,
72}
73
74/// Options for the assembler
75#[derive(Debug, Clone, Default)]
76pub struct AssemblerOption {
77    /// sBPF target architecture
78    pub arch: SbpfArch,
79    /// Optional debug mode configuration
80    pub debug_mode: Option<DebugMode>,
81    /// Optional optimization and CFG diagnostic configuration
82    pub optimization: OptimizationConfig,
83}
84
85impl AssemblerOption {
86    /// Set the target architecture
87    pub fn with_arch(mut self, arch: SbpfArch) -> Self {
88        self.arch = arch;
89        self
90    }
91
92    /// Enable debug mode with the given config
93    pub fn with_debug_mode(mut self, debug_mode: DebugMode) -> Self {
94        self.debug_mode = Some(debug_mode);
95        self
96    }
97}
98
99/// An error enriched with source location information from preprocessing.
100/// Wraps a `CompileError` with the resolved original source location.
101#[derive(Debug)]
102pub struct AssemblerError {
103    pub error: CompileError,
104    pub origin: Option<SourceOrigin>,
105    /// Column offset (0-based) within the original line, if known.
106    pub column: Option<usize>,
107}
108
109impl std::fmt::Display for AssemblerError {
110    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
111        write!(f, "{}", self.error)
112    }
113}
114
115impl std::error::Error for AssemblerError {
116    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
117        Some(&self.error)
118    }
119}
120
121/// Returned when assembly with preprocessing fails.
122/// Contains the errors and the file registry so callers can render
123/// diagnostics against the original source files.
124#[derive(Debug)]
125pub struct AssembleErrors {
126    pub errors: Vec<AssemblerError>,
127    pub file_registry: FileRegistry,
128}
129
130/// Assembler for SBPF assembly code
131#[derive(Debug, Clone)]
132pub struct Assembler {
133    options: AssemblerOption,
134}
135
136impl Assembler {
137    /// Create a new Assembler with the given options
138    pub fn new(options: AssemblerOption) -> Self {
139        Self { options }
140    }
141
142    /// Assemble source code directly (no preprocessing).
143    /// This is the original API -- macros and includes are not supported.
144    pub fn assemble(&self, source: &str) -> Result<Vec<u8>, Vec<CompileError>> {
145        let parse_result = match parse_with_optimization(
146            source,
147            self.options.arch,
148            self.options.optimization.clone(),
149        ) {
150            Ok(result) => result,
151            Err(errors) => {
152                return Err(errors);
153            }
154        };
155
156        // Build debug data if debug mode is enabled
157        let debug_data = if let Some(ref debug_mode) = self.options.debug_mode {
158            let (lines, labels) = collect_line_and_label_entries(source, &parse_result);
159            let code_end = parse_result.code_section.get_size();
160
161            Some(DebugData {
162                filename: debug_mode.filename.clone(),
163                directory: debug_mode.directory.clone(),
164                lines,
165                labels,
166                code_start: 0,
167                code_end,
168            })
169        } else {
170            None
171        };
172
173        let program = Program::from_parse_result(parse_result, debug_data);
174        let bytecode = program.emit_bytecode();
175        Ok(bytecode)
176    }
177
178    /// Assemble with preprocessing: resolves `.include` and expands `.macro` directives
179    /// before parsing. Errors include source location information from the source map,
180    /// and the file registry is returned so callers can render diagnostics against
181    /// original source files.
182    pub fn assemble_with_preprocess(
183        &self,
184        source: &str,
185        source_path: &str,
186        resolver: Option<&dyn FileResolver>,
187    ) -> Result<Vec<u8>, AssembleErrors> {
188        // Run preprocessor
189        let preprocess_result =
190            preprocess(source, source_path, resolver).map_err(|failure| AssembleErrors {
191                errors: failure
192                    .errors
193                    .into_iter()
194                    .map(|e| AssemblerError {
195                        error: e.error,
196                        origin: e.origin,
197                        column: None,
198                    })
199                    .collect(),
200                file_registry: failure.file_registry,
201            })?;
202
203        let expanded = &preprocess_result.expanded_source;
204        let source_map = &preprocess_result.source_map;
205
206        // Parse the expanded source
207        let parse_result = match parse_with_optimization(
208            expanded,
209            self.options.arch,
210            self.options.optimization.clone(),
211        ) {
212            Ok(result) => result,
213            Err(errors) => {
214                // Extract file registry from source map before moving errors
215                let file_registry = source_map.file_registry.clone();
216                return Err(AssembleErrors {
217                    errors: errors
218                        .into_iter()
219                        .map(|e| {
220                            let span = e.span();
221                            let origin = source_map.resolve_span(span, expanded).clone();
222                            // Compute column offset within the line from the
223                            // expanded source so we can highlight the right token.
224                            let col = expanded[..span.start]
225                                .rfind('\n')
226                                .map(|nl| span.start - nl - 1)
227                                .unwrap_or(span.start);
228                            AssemblerError {
229                                error: e,
230                                column: Some(col),
231                                origin: Some(origin),
232                            }
233                        })
234                        .collect(),
235                    file_registry,
236                });
237            }
238        };
239
240        // Build debug data if debug mode is enabled
241        let debug_data = if let Some(ref debug_mode) = self.options.debug_mode {
242            let (lines, labels) = collect_line_and_label_entries(expanded, &parse_result);
243            let code_end = parse_result.code_section.get_size();
244
245            Some(DebugData {
246                filename: debug_mode.filename.clone(),
247                directory: debug_mode.directory.clone(),
248                lines,
249                labels,
250                code_start: 0,
251                code_end,
252            })
253        } else {
254            None
255        };
256
257        let program = Program::from_parse_result(parse_result, debug_data);
258        let bytecode = program.emit_bytecode();
259        Ok(bytecode)
260    }
261
262    /// Convenience method: read a file from disk and assemble with full preprocessing.
263    pub fn assemble_file(&self, path: &std::path::Path) -> Result<Vec<u8>, AssembleErrors> {
264        let source = std::fs::read_to_string(path).map_err(|e| AssembleErrors {
265            errors: vec![AssemblerError {
266                error: CompileError::IncludeReadError {
267                    path: path.display().to_string(),
268                    reason: e.to_string(),
269                    span: 0..0,
270                    custom_label: Some("Failed to read source file".to_string()),
271                },
272                origin: None,
273                column: None,
274            }],
275            file_registry: FileRegistry::new(),
276        })?;
277
278        let source_path = path.to_string_lossy();
279        let resolver = FsFileResolver::new();
280        self.assemble_with_preprocess(&source, &source_path, Some(&resolver))
281    }
282}
283
284type LineEntry = (u64, u32); // (offset, line)
285type LabelEntry = (String, u64, u32); // (label, offset, line)
286
287/// Helper function to collect line and label entries
288fn collect_line_and_label_entries(
289    source: &str,
290    parse_result: &ProgramLayout,
291) -> (Vec<LineEntry>, Vec<LabelEntry>) {
292    let mut files: Files<&str> = Files::new();
293    let file_id = files.add("source", source);
294
295    let mut line_entries = Vec::new();
296    let mut label_entries = Vec::new();
297
298    for node in parse_result.code_section.get_nodes() {
299        match node {
300            ASTNode::Instruction {
301                instruction,
302                offset,
303            } => {
304                let line_index = files.line_index(file_id, instruction.span.start as u32);
305                let line_number = (line_index.to_usize() + 1) as u32;
306                line_entries.push((*offset, line_number));
307            }
308            ASTNode::Label { label, offset } => {
309                let line_index = files.line_index(file_id, label.span.start as u32);
310                let line_number = (line_index.to_usize() + 1) as u32;
311                label_entries.push((label.name.clone(), *offset, line_number));
312            }
313            _ => {}
314        }
315    }
316
317    for node in parse_result.data_section.get_nodes() {
318        if let ASTNode::ROData { rodata, offset } = node {
319            let line_index = files.line_index(file_id, rodata.span.start as u32);
320            let line_number = (line_index.to_usize() + 1) as u32;
321            label_entries.push((rodata.name.clone(), *offset, line_number));
322        }
323    }
324
325    (line_entries, label_entries)
326}
327
328#[cfg(test)]
329pub fn assemble(source: &str) -> Result<Vec<u8>, Vec<CompileError>> {
330    let options = AssemblerOption::default();
331    let assembler = Assembler::new(options);
332    assembler.assemble(source)
333}
334
335#[cfg(test)]
336pub fn assemble_with_debug_data(
337    source: &str,
338    filename: &str,
339    directory: &str,
340) -> Result<Vec<u8>, Vec<CompileError>> {
341    let options = AssemblerOption::default().with_debug_mode(DebugMode {
342        filename: filename.to_string(),
343        directory: directory.to_string(),
344    });
345    let assembler = Assembler::new(options);
346    assembler.assemble(source)
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352
353    #[test]
354    fn test_assemble_success() {
355        let source = "exit";
356        let result = assemble(source);
357        assert!(result.is_ok());
358        let bytecode = result.unwrap();
359        assert!(!bytecode.is_empty());
360    }
361
362    #[test]
363    fn test_assemble_parse_error() {
364        let source = "invalid_xyz";
365        let result = assemble(source);
366        assert!(result.is_err());
367    }
368
369    #[test]
370    fn test_parse_error_has_correct_span() {
371        // Verify parse errors point to the actual error position,
372        // not span 0..len (which would lose position info for source maps).
373        let source = ".rodata\n    thing2: .bogus 5\n.text\nentrypoint:\n    exit\n";
374        let result = assemble(source);
375        assert!(result.is_err());
376        let errors = result.unwrap_err();
377        let span = errors[0].span();
378        // The error should point somewhere near ".bogus" (byte ~20), NOT 0..len
379        assert!(
380            span.start > 0,
381            "Error span should not start at 0, got {:?}",
382            span
383        );
384        assert!(
385            span.end < source.len(),
386            "Error span should not cover entire source, got {:?}",
387            span
388        );
389        // Verify the span actually points at ".bogus", not at "thing2"
390        let error_text = &source[span.start..span.end.min(source.len())];
391        assert!(
392            error_text.starts_with('.'),
393            "Error should point at '.bogus', but points at: '{}'",
394            &source[span.start..span.start + 10.min(source.len() - span.start)]
395        );
396    }
397
398    #[test]
399    fn test_assemble_with_equ_directive() {
400        let source = r#"
401        .globl entrypoint
402        .equ MY_CONST, 42
403        entrypoint:
404            mov64 r1, MY_CONST
405            exit
406        "#;
407        let result = assemble(source);
408        assert!(result.is_ok());
409    }
410
411    #[test]
412    fn test_assemble_const_expr_overflow_errors() {
413        for expr in ["BIG + 1", "BIG * 2"] {
414            let source = format!(
415                r#"
416                .globl entrypoint
417                .equ BIG, 0x7FFFFFFFFFFFFFFF
418                entrypoint:
419                    lddw r1, {expr}
420                    exit
421                "#
422            );
423            let result = assemble(&source);
424            assert!(result.is_err(), "expected error for '{expr}'");
425            assert!(matches!(
426                result.unwrap_err().first(),
427                Some(CompileError::ArithmeticError { .. })
428            ));
429        }
430    }
431
432    #[test]
433    fn test_assemble_const_expr_div_by_zero_errors() {
434        let source = r#"
435        .globl entrypoint
436        entrypoint:
437            lddw r1, 8 / 0
438            exit
439        "#;
440        let result = assemble(source);
441        assert!(result.is_err());
442        let errors = result.unwrap_err();
443        assert!(matches!(
444            errors.first(),
445            Some(CompileError::ArithmeticError { .. })
446        ));
447        // The diagnostic must point at the offending expression, not 0..len.
448        assert!(errors[0].span().start > 0);
449    }
450
451    #[test]
452    fn test_assemble_const_expr_overflow_in_directive_errors() {
453        // The directive-value evaluator must also be overflow-safe.
454        let source = r#"
455        .globl entrypoint
456        .equ BAD, 0x7FFFFFFFFFFFFFFF + 1
457        entrypoint:
458            lddw r1, BAD
459            exit
460        "#;
461        let result = assemble(source);
462        assert!(result.is_err());
463        assert!(matches!(
464            result.unwrap_err().first(),
465            Some(CompileError::ArithmeticError { .. })
466        ));
467    }
468
469    #[test]
470    fn test_assemble_valid_const_expr_still_folds() {
471        // Regression: in-range arithmetic must keep folding correctly.
472        let source = r#"
473        .globl entrypoint
474        .equ A, 0x10
475        .equ B, 0x20
476        entrypoint:
477            lddw r1, A * 4 + B
478            exit
479        "#;
480        assert!(assemble(source).is_ok());
481    }
482
483    #[test]
484    fn test_assemble_duplicate_label_error() {
485        let source = r#"
486        .globl entrypoint
487        entrypoint:
488            mov64 r1, 1
489        entrypoint:
490            exit
491        "#;
492        let result = assemble(source);
493        assert!(result.is_err());
494        let errors = result.unwrap_err();
495        assert!(!errors.is_empty());
496    }
497
498    #[test]
499    fn test_assemble_undefined_label_error() {
500        let instructions = [
501            "mov64 r1, UNDEF",
502            "jeq r1, UNDEF, done",
503            "jeq32 r1, UNDEF, done",
504            "stb [r1+0], UNDEF",
505            "stb [r1+UNDEF], 0x1",
506            "stxb [r1+UNDEF], r2",
507            "ldxb r1, [r2+UNDEF]",
508            "call UNDEF",
509            "lddw r1, UNDEF",
510            "ja UNDEF",
511            "jeq r1, 0x1, UNDEF",
512            "jeq r1, r2, UNDEF",
513            "jeq32 r1, 0x1, UNDEF",
514            "jeq32 r1, r2, UNDEF",
515        ];
516
517        for instruction in instructions {
518            let source = format!(
519                r#"
520                .globl e
521                e:
522                    {instruction}
523                    exit
524                done:
525                    exit
526                "#
527            );
528            let instruction_start = source.find(instruction).unwrap();
529            let instruction_end = instruction_start + instruction.len();
530            let errors = assemble(&source).unwrap_err();
531
532            assert_eq!(errors.len(), 1, "unexpected errors for `{instruction}`");
533            match &errors[0] {
534                CompileError::UndefinedLabel { label, span, .. } => {
535                    assert_eq!(label, "UNDEF", "unexpected label for `{instruction}`");
536                    assert_eq!(
537                        span,
538                        &(instruction_start..instruction_end),
539                        "unexpected span for `{instruction}`"
540                    );
541                }
542                error => {
543                    panic!("unexpected error for `{instruction}`: {error:?}")
544                }
545            }
546        }
547    }
548
549    #[test]
550    fn test_assemble_extern_directive() {
551        let source = r#"
552        .globl entrypoint
553        .extern my_extern_symbol
554        entrypoint:
555            exit
556        "#;
557        let result = assemble(source);
558        assert!(result.is_ok());
559    }
560
561    #[test]
562    fn test_assemble_rodata_section() {
563        let source = r#"
564        .globl entrypoint
565        .rodata
566        my_data: .ascii "hello"
567        .text
568        entrypoint:
569            exit
570        "#;
571        let result = assemble(source);
572        assert!(result.is_ok());
573    }
574
575    #[test]
576    fn test_assemble_rodata_byte() {
577        let source = r#"
578        .globl entrypoint
579        .rodata
580        my_byte: .byte 0x42
581        .text
582        entrypoint:
583            exit
584        "#;
585        let result = assemble(source);
586        assert!(result.is_ok());
587    }
588
589    #[test]
590    fn test_assemble_rodata_multiple_bytes() {
591        let source = r#"
592        .globl entrypoint
593        .rodata
594        my_bytes: .byte 0x01, 0x02, 0x03, 0x04
595        .text
596        entrypoint:
597            exit
598        "#;
599        let result = assemble(source);
600        assert!(result.is_ok());
601    }
602
603    #[test]
604    fn test_assemble_rodata_mixed() {
605        let source = r#"
606        .globl entrypoint
607        .rodata
608        data1: .byte 0x42
609        data2: .ascii "test"
610        .text
611        entrypoint:
612            exit
613        "#;
614        let result = assemble(source);
615        assert!(result.is_ok());
616    }
617
618    #[test]
619    fn test_assemble_jump_operations() {
620        let source = r#"
621        .globl entrypoint
622        entrypoint:
623            jeq r1, 0, +1
624            ja +2
625        target:
626            jne r1, r2, target
627            exit
628        "#;
629        let result = assemble(source);
630        assert!(result.is_ok());
631    }
632
633    #[test]
634    fn test_assemble_jump32_v3() {
635        let source = r#"
636        .globl entrypoint
637        entrypoint:
638            jeq32 r1, 0, +1
639            jset32 r1, r2, +1
640            exit
641        "#;
642        let assembler = Assembler::new(AssemblerOption::default());
643        let result = assembler.assemble(source);
644        assert!(result.is_ok());
645    }
646
647    #[test]
648    fn test_assemble_jump32_v0() {
649        let source = r#"
650        .globl entrypoint
651        entrypoint:
652            jeq32 r1, 0, +1
653            exit
654        "#;
655        let assembler = Assembler::new(AssemblerOption::default().with_arch(SbpfArch::V0));
656        let result = assembler.assemble(source);
657        // jmp32 operations should not work in v0
658        assert!(result.is_err());
659    }
660
661    #[test]
662    fn test_assemble_llvm_jump32_v3() {
663        let source = r#"
664        .globl entrypoint
665        entrypoint:
666            if w1 == 0 goto +1
667            if w1 & w2 goto +1
668            exit
669        "#;
670        let assembler = Assembler::new(AssemblerOption::default());
671        let result = assembler.assemble(source);
672        assert!(result.is_ok(),);
673    }
674
675    #[test]
676    fn test_assemble_llvm_jump32_v0() {
677        let source = r#"
678        .globl entrypoint
679        entrypoint:
680            if w1 == 0 goto +1
681            exit
682        "#;
683        let assembler = Assembler::new(AssemblerOption::default().with_arch(SbpfArch::V0));
684        let result = assembler.assemble(source);
685        // jmp32 operations should not work in v0
686        assert!(result.is_err());
687    }
688
689    #[test]
690    fn test_assemble_offset_expression() {
691        let source = r#"
692        .globl entrypoint
693        .equ BASE, 100
694        entrypoint:
695            mov64 r1, BASE+10
696            exit
697        "#;
698        let result = assemble(source);
699        assert!(result.is_ok());
700    }
701
702    #[test]
703    fn test_assemble_equ_expression() {
704        let source = r#"
705        .globl entrypoint
706        .equ BASE, 100
707        .equ OFFSET, 20
708        .equ COMPUTED, BASE
709        entrypoint:
710            mov64 r1, BASE
711            mov64 r2, OFFSET
712            mov64 r3, COMPUTED
713            exit
714        "#;
715        let result = assemble(source);
716        assert!(result.is_ok());
717    }
718
719    #[test]
720    fn test_assemble_label_arithmetic_rodata_length() {
721        // The primary use case: compute string length via label subtraction
722        let source = r#"
723        .globl entrypoint
724        .rodata
725        msg: .ascii "Hello"
726        msg_end:
727        .text
728        entrypoint:
729            lddw r1, msg
730            mov64 r2, msg_end - msg
731            exit
732        "#;
733        let result = assemble(source);
734        assert!(result.is_ok(), "Failed: {:?}", result.err());
735    }
736
737    #[test]
738    fn test_assemble_label_arithmetic_with_offset() {
739        // Label arithmetic with additional constant offset
740        let source = r#"
741        .globl entrypoint
742        .rodata
743        msg: .ascii "Hello!"
744        msg_end:
745        .text
746        entrypoint:
747            lddw r1, msg
748            mov64 r2, msg_end - msg - 1
749            exit
750        "#;
751        let result = assemble(source);
752        assert!(result.is_ok(), "Failed: {:?}", result.err());
753    }
754
755    #[test]
756    fn test_assemble_label_arithmetic_text_section() {
757        // Label arithmetic works in the text section too
758        let source = r#"
759        .globl entrypoint
760        entrypoint:
761            mov64 r1, 1
762        middle:
763            mov64 r2, 2
764        end:
765            mov64 r3, end - entrypoint
766            exit
767        "#;
768        let result = assemble(source);
769        assert!(result.is_ok(), "Failed: {:?}", result.err());
770    }
771
772    #[test]
773    fn test_assemble_label_arithmetic_forward_reference() {
774        // Text section before rodata — forward references to rodata labels
775        let source = r#"
776        .globl entrypoint
777        entrypoint:
778            lddw r1, message
779            mov64 r2, message_end - message
780            call sol_log_
781            exit
782            lddw r10, 1
783        .rodata
784            message: .ascii "Hello, Solana!"
785            message_end:
786        "#;
787        let result = assemble(source);
788        assert!(
789            result.is_ok(),
790            "Forward reference failed: {:?}",
791            result.err()
792        );
793    }
794
795    #[test]
796    fn test_assemble_label_arithmetic_multiline_rodata() {
797        // Rodata label and directive on separate lines (as from macro expansion)
798        let source = r#"
799        .globl entrypoint
800        entrypoint:
801            lddw r1, message
802            mov64 r2, message_end - message
803            call sol_log_
804            exit
805        .rodata
806        message:
807            .ascii "Hello, Solana!"
808        message_end:
809        "#;
810        let result = assemble(source);
811        assert!(
812            result.is_ok(),
813            "Multi-line rodata failed: {:?}",
814            result.err()
815        );
816    }
817
818    #[test]
819    fn test_assemble_label_arithmetic_macro_e2e() {
820        // Full end-to-end test with macro expansion + label arithmetic
821        let source = r#"
822.macro DEF_STR name, text
823\name:
824    .ascii \text
825\name\()_end:
826.endm
827
828.macro SOL_LOG name
829    lddw r1, \name
830    mov64 r2, \name\()_end - \name
831    call sol_log_
832.endm
833
834.globl entrypoint
835entrypoint:
836    SOL_LOG message
837    exit
838.rodata
839    DEF_STR message, "Hello, Solana!"
840"#;
841        let assembler = Assembler::new(AssemblerOption::default());
842        let result = assembler.assemble_with_preprocess(source, "test.s", None);
843        assert!(result.is_ok(), "Macro e2e failed: {:?}", result.err());
844    }
845
846    #[test]
847    fn test_assemble_label_arithmetic_cross_section_error() {
848        // Cross-section arithmetic should fail
849        let source = r#"
850        .globl entrypoint
851        .rodata
852        msg: .ascii "Hello"
853        .text
854        entrypoint:
855            mov64 r1, msg - entrypoint
856            exit
857        "#;
858        let result = assemble(source);
859        assert!(result.is_err(), "Cross-section arithmetic should fail");
860    }
861
862    #[test]
863    fn test_assemble_label_arithmetic_complex_expression() {
864        // More complex expression with multiple rodata entries
865        let source = r#"
866        .globl entrypoint
867        .rodata
868        str1: .ascii "Hello"
869        str2: .ascii " World"
870        str2_end:
871        .text
872        entrypoint:
873            lddw r1, str2
874            mov64 r2, str2_end - str2
875            exit
876        "#;
877        let result = assemble(source);
878        assert!(result.is_ok(), "Failed: {:?}", result.err());
879    }
880
881    #[test]
882    fn test_parse_error_column_through_preprocess() {
883        // Verify the column offset is correctly computed through the
884        // preprocessing pipeline so errors point at the invalid token,
885        // not at the label before it.
886        let source = ".globl e\ne:\n    exit\n.rodata\n    thing1: .bogus 5\n";
887        let assembler = Assembler::new(AssemblerOption::default());
888        let result = assembler.assemble_with_preprocess(source, "test.s", None);
889        assert!(result.is_err());
890        let errors = result.unwrap_err();
891        let err = &errors.errors[0];
892        let origin = err.origin.as_ref().expect("Expected origin");
893        let registry = &errors.file_registry;
894
895        // Verify the column points at ".bogus", not "thing1"
896        let col = err.column.expect("Expected column info");
897        let line_start = registry.line_byte_offset(origin.file_id, origin.line);
898        let content = registry.content(origin.file_id);
899        let at_col = &content[line_start + col..];
900        assert!(
901            at_col.starts_with(".bogus"),
902            "Expected column to point at '.bogus', but points at: '{}'",
903            &at_col[..at_col.len().min(20)]
904        );
905    }
906
907    #[test]
908    fn test_parse_error_column_with_include_and_macros() {
909        // Simulate the user's actual scenario: .include with macros,
910        // then an invalid directive on a later line.
911        use std::collections::HashMap;
912
913        struct TestResolver {
914            files: HashMap<String, String>,
915        }
916        impl FileResolver for TestResolver {
917            fn resolve(&self, path: &str, _base: &str) -> Result<String, std::io::Error> {
918                self.files.get(path).cloned().ok_or_else(|| {
919                    std::io::Error::new(
920                        std::io::ErrorKind::NotFound,
921                        format!("not found: {}", path),
922                    )
923                })
924            }
925        }
926
927        let mut files = HashMap::new();
928        files.insert(
929            "syscalls/sol_log.s".to_string(),
930            r#".macro SOL_LOG name
931    lddw r1, \name
932    mov64 r2, \name\()_end - \name
933    call sol_log_
934.endm
935
936.macro DEF_STR name, text
937\name:
938    .ascii \text
939\name\()_end:
940.endm
941"#
942            .to_string(),
943        );
944
945        let source = r#".include "syscalls/sol_log.s"
946.globl e
947e:
948    SOL_LOG message
949.rodata
950    DEF_STR message, "TEST"
951    thing2: .int 1
952    thing1: .bogus "text"
953"#;
954        let resolver = TestResolver { files };
955        let assembler = Assembler::new(AssemblerOption::default());
956        let result = assembler.assemble_with_preprocess(source, "main.s", Some(&resolver));
957        assert!(result.is_err());
958        let errors = result.unwrap_err();
959        let err = &errors.errors[0];
960        let origin = err.origin.as_ref().expect("Expected origin");
961        let registry = &errors.file_registry;
962
963        // Verify origin points to main.s, not the included file
964        assert_eq!(registry.path(origin.file_id), "main.s");
965
966        // Verify the column points at ".bogus"
967        let col = err.column.expect("Expected column info");
968        let line_start = registry.line_byte_offset(origin.file_id, origin.line);
969        let content = registry.content(origin.file_id);
970        let at_col = &content[line_start + col..];
971        assert!(
972            at_col.starts_with(".bogus"),
973            "Expected column to point at '.bogus', but at line {} col {}, points at: '{}'",
974            origin.line,
975            col,
976            &at_col[..at_col.len().min(20)]
977        );
978    }
979
980    #[test]
981    fn test_assemble_with_debug_data() {
982        let source = r#".equ MSG_LEN, 14
983
984.globl entrypoint
985entrypoint:
986  lddw r1, message
987  mov64 r2, MSG_LEN
988  call sol_log_
989  exit
990.rodata
991  message: .ascii "Hello, Solana!"
992"#;
993        let result = assemble_with_debug_data(source, "hello_solana.s", "/tmp");
994        assert!(result.is_ok());
995        let bytecode = result.unwrap();
996
997        // Verify the ELF has all debug sections.
998        let bytecode_str = String::from_utf8_lossy(&bytecode);
999        assert!(
1000            bytecode_str.contains(".debug_abbrev"),
1001            "Missing .debug_abbrev section"
1002        );
1003        assert!(
1004            bytecode_str.contains(".debug_info"),
1005            "Missing .debug_info section"
1006        );
1007        assert!(
1008            bytecode_str.contains(".debug_line"),
1009            "Missing .debug_line section"
1010        );
1011        assert!(
1012            bytecode_str.contains(".debug_line_str"),
1013            "Missing .debug_line_str section"
1014        );
1015    }
1016}