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_extern_directive() {
500        let source = r#"
501        .globl entrypoint
502        .extern my_extern_symbol
503        entrypoint:
504            exit
505        "#;
506        let result = assemble(source);
507        assert!(result.is_ok());
508    }
509
510    #[test]
511    fn test_assemble_rodata_section() {
512        let source = r#"
513        .globl entrypoint
514        .rodata
515        my_data: .ascii "hello"
516        .text
517        entrypoint:
518            exit
519        "#;
520        let result = assemble(source);
521        assert!(result.is_ok());
522    }
523
524    #[test]
525    fn test_assemble_rodata_byte() {
526        let source = r#"
527        .globl entrypoint
528        .rodata
529        my_byte: .byte 0x42
530        .text
531        entrypoint:
532            exit
533        "#;
534        let result = assemble(source);
535        assert!(result.is_ok());
536    }
537
538    #[test]
539    fn test_assemble_rodata_multiple_bytes() {
540        let source = r#"
541        .globl entrypoint
542        .rodata
543        my_bytes: .byte 0x01, 0x02, 0x03, 0x04
544        .text
545        entrypoint:
546            exit
547        "#;
548        let result = assemble(source);
549        assert!(result.is_ok());
550    }
551
552    #[test]
553    fn test_assemble_rodata_mixed() {
554        let source = r#"
555        .globl entrypoint
556        .rodata
557        data1: .byte 0x42
558        data2: .ascii "test"
559        .text
560        entrypoint:
561            exit
562        "#;
563        let result = assemble(source);
564        assert!(result.is_ok());
565    }
566
567    #[test]
568    fn test_assemble_jump_operations() {
569        let source = r#"
570        .globl entrypoint
571        entrypoint:
572            jeq r1, 0, +1
573            ja +2
574        target:
575            jne r1, r2, target
576            exit
577        "#;
578        let result = assemble(source);
579        assert!(result.is_ok());
580    }
581
582    #[test]
583    fn test_assemble_jump32_v3() {
584        let source = r#"
585        .globl entrypoint
586        entrypoint:
587            jeq32 r1, 0, +1
588            jset32 r1, r2, +1
589            exit
590        "#;
591        let assembler = Assembler::new(AssemblerOption::default());
592        let result = assembler.assemble(source);
593        assert!(result.is_ok());
594    }
595
596    #[test]
597    fn test_assemble_jump32_v0() {
598        let source = r#"
599        .globl entrypoint
600        entrypoint:
601            jeq32 r1, 0, +1
602            exit
603        "#;
604        let assembler = Assembler::new(AssemblerOption::default().with_arch(SbpfArch::V0));
605        let result = assembler.assemble(source);
606        // jmp32 operations should not work in v0
607        assert!(result.is_err());
608    }
609
610    #[test]
611    fn test_assemble_llvm_jump32_v3() {
612        let source = r#"
613        .globl entrypoint
614        entrypoint:
615            if w1 == 0 goto +1
616            if w1 & w2 goto +1
617            exit
618        "#;
619        let assembler = Assembler::new(AssemblerOption::default());
620        let result = assembler.assemble(source);
621        assert!(result.is_ok(),);
622    }
623
624    #[test]
625    fn test_assemble_llvm_jump32_v0() {
626        let source = r#"
627        .globl entrypoint
628        entrypoint:
629            if w1 == 0 goto +1
630            exit
631        "#;
632        let assembler = Assembler::new(AssemblerOption::default().with_arch(SbpfArch::V0));
633        let result = assembler.assemble(source);
634        // jmp32 operations should not work in v0
635        assert!(result.is_err());
636    }
637
638    #[test]
639    fn test_assemble_offset_expression() {
640        let source = r#"
641        .globl entrypoint
642        .equ BASE, 100
643        entrypoint:
644            mov64 r1, BASE+10
645            exit
646        "#;
647        let result = assemble(source);
648        assert!(result.is_ok());
649    }
650
651    #[test]
652    fn test_assemble_equ_expression() {
653        let source = r#"
654        .globl entrypoint
655        .equ BASE, 100
656        .equ OFFSET, 20
657        .equ COMPUTED, BASE
658        entrypoint:
659            mov64 r1, BASE
660            mov64 r2, OFFSET
661            mov64 r3, COMPUTED
662            exit
663        "#;
664        let result = assemble(source);
665        assert!(result.is_ok());
666    }
667
668    #[test]
669    fn test_assemble_label_arithmetic_rodata_length() {
670        // The primary use case: compute string length via label subtraction
671        let source = r#"
672        .globl entrypoint
673        .rodata
674        msg: .ascii "Hello"
675        msg_end:
676        .text
677        entrypoint:
678            lddw r1, msg
679            mov64 r2, msg_end - msg
680            exit
681        "#;
682        let result = assemble(source);
683        assert!(result.is_ok(), "Failed: {:?}", result.err());
684    }
685
686    #[test]
687    fn test_assemble_label_arithmetic_with_offset() {
688        // Label arithmetic with additional constant offset
689        let source = r#"
690        .globl entrypoint
691        .rodata
692        msg: .ascii "Hello!"
693        msg_end:
694        .text
695        entrypoint:
696            lddw r1, msg
697            mov64 r2, msg_end - msg - 1
698            exit
699        "#;
700        let result = assemble(source);
701        assert!(result.is_ok(), "Failed: {:?}", result.err());
702    }
703
704    #[test]
705    fn test_assemble_label_arithmetic_text_section() {
706        // Label arithmetic works in the text section too
707        let source = r#"
708        .globl entrypoint
709        entrypoint:
710            mov64 r1, 1
711        middle:
712            mov64 r2, 2
713        end:
714            mov64 r3, end - entrypoint
715            exit
716        "#;
717        let result = assemble(source);
718        assert!(result.is_ok(), "Failed: {:?}", result.err());
719    }
720
721    #[test]
722    fn test_assemble_label_arithmetic_forward_reference() {
723        // Text section before rodata — forward references to rodata labels
724        let source = r#"
725        .globl entrypoint
726        entrypoint:
727            lddw r1, message
728            mov64 r2, message_end - message
729            call sol_log_
730            exit
731            lddw r10, 1
732        .rodata
733            message: .ascii "Hello, Solana!"
734            message_end:
735        "#;
736        let result = assemble(source);
737        assert!(
738            result.is_ok(),
739            "Forward reference failed: {:?}",
740            result.err()
741        );
742    }
743
744    #[test]
745    fn test_assemble_label_arithmetic_multiline_rodata() {
746        // Rodata label and directive on separate lines (as from macro expansion)
747        let source = r#"
748        .globl entrypoint
749        entrypoint:
750            lddw r1, message
751            mov64 r2, message_end - message
752            call sol_log_
753            exit
754        .rodata
755        message:
756            .ascii "Hello, Solana!"
757        message_end:
758        "#;
759        let result = assemble(source);
760        assert!(
761            result.is_ok(),
762            "Multi-line rodata failed: {:?}",
763            result.err()
764        );
765    }
766
767    #[test]
768    fn test_assemble_label_arithmetic_macro_e2e() {
769        // Full end-to-end test with macro expansion + label arithmetic
770        let source = r#"
771.macro DEF_STR name, text
772\name:
773    .ascii \text
774\name\()_end:
775.endm
776
777.macro SOL_LOG name
778    lddw r1, \name
779    mov64 r2, \name\()_end - \name
780    call sol_log_
781.endm
782
783.globl entrypoint
784entrypoint:
785    SOL_LOG message
786    exit
787.rodata
788    DEF_STR message, "Hello, Solana!"
789"#;
790        let assembler = Assembler::new(AssemblerOption::default());
791        let result = assembler.assemble_with_preprocess(source, "test.s", None);
792        assert!(result.is_ok(), "Macro e2e failed: {:?}", result.err());
793    }
794
795    #[test]
796    fn test_assemble_label_arithmetic_cross_section_error() {
797        // Cross-section arithmetic should fail
798        let source = r#"
799        .globl entrypoint
800        .rodata
801        msg: .ascii "Hello"
802        .text
803        entrypoint:
804            mov64 r1, msg - entrypoint
805            exit
806        "#;
807        let result = assemble(source);
808        assert!(result.is_err(), "Cross-section arithmetic should fail");
809    }
810
811    #[test]
812    fn test_assemble_label_arithmetic_complex_expression() {
813        // More complex expression with multiple rodata entries
814        let source = r#"
815        .globl entrypoint
816        .rodata
817        str1: .ascii "Hello"
818        str2: .ascii " World"
819        str2_end:
820        .text
821        entrypoint:
822            lddw r1, str2
823            mov64 r2, str2_end - str2
824            exit
825        "#;
826        let result = assemble(source);
827        assert!(result.is_ok(), "Failed: {:?}", result.err());
828    }
829
830    #[test]
831    fn test_parse_error_column_through_preprocess() {
832        // Verify the column offset is correctly computed through the
833        // preprocessing pipeline so errors point at the invalid token,
834        // not at the label before it.
835        let source = ".globl e\ne:\n    exit\n.rodata\n    thing1: .bogus 5\n";
836        let assembler = Assembler::new(AssemblerOption::default());
837        let result = assembler.assemble_with_preprocess(source, "test.s", None);
838        assert!(result.is_err());
839        let errors = result.unwrap_err();
840        let err = &errors.errors[0];
841        let origin = err.origin.as_ref().expect("Expected origin");
842        let registry = &errors.file_registry;
843
844        // Verify the column points at ".bogus", not "thing1"
845        let col = err.column.expect("Expected column info");
846        let line_start = registry.line_byte_offset(origin.file_id, origin.line);
847        let content = registry.content(origin.file_id);
848        let at_col = &content[line_start + col..];
849        assert!(
850            at_col.starts_with(".bogus"),
851            "Expected column to point at '.bogus', but points at: '{}'",
852            &at_col[..at_col.len().min(20)]
853        );
854    }
855
856    #[test]
857    fn test_parse_error_column_with_include_and_macros() {
858        // Simulate the user's actual scenario: .include with macros,
859        // then an invalid directive on a later line.
860        use std::collections::HashMap;
861
862        struct TestResolver {
863            files: HashMap<String, String>,
864        }
865        impl FileResolver for TestResolver {
866            fn resolve(&self, path: &str, _base: &str) -> Result<String, std::io::Error> {
867                self.files.get(path).cloned().ok_or_else(|| {
868                    std::io::Error::new(
869                        std::io::ErrorKind::NotFound,
870                        format!("not found: {}", path),
871                    )
872                })
873            }
874        }
875
876        let mut files = HashMap::new();
877        files.insert(
878            "syscalls/sol_log.s".to_string(),
879            r#".macro SOL_LOG name
880    lddw r1, \name
881    mov64 r2, \name\()_end - \name
882    call sol_log_
883.endm
884
885.macro DEF_STR name, text
886\name:
887    .ascii \text
888\name\()_end:
889.endm
890"#
891            .to_string(),
892        );
893
894        let source = r#".include "syscalls/sol_log.s"
895.globl e
896e:
897    SOL_LOG message
898.rodata
899    DEF_STR message, "TEST"
900    thing2: .int 1
901    thing1: .bogus "text"
902"#;
903        let resolver = TestResolver { files };
904        let assembler = Assembler::new(AssemblerOption::default());
905        let result = assembler.assemble_with_preprocess(source, "main.s", Some(&resolver));
906        assert!(result.is_err());
907        let errors = result.unwrap_err();
908        let err = &errors.errors[0];
909        let origin = err.origin.as_ref().expect("Expected origin");
910        let registry = &errors.file_registry;
911
912        // Verify origin points to main.s, not the included file
913        assert_eq!(registry.path(origin.file_id), "main.s");
914
915        // Verify the column points at ".bogus"
916        let col = err.column.expect("Expected column info");
917        let line_start = registry.line_byte_offset(origin.file_id, origin.line);
918        let content = registry.content(origin.file_id);
919        let at_col = &content[line_start + col..];
920        assert!(
921            at_col.starts_with(".bogus"),
922            "Expected column to point at '.bogus', but at line {} col {}, points at: '{}'",
923            origin.line,
924            col,
925            &at_col[..at_col.len().min(20)]
926        );
927    }
928
929    #[test]
930    fn test_assemble_with_debug_data() {
931        let source = r#".equ MSG_LEN, 14
932
933.globl entrypoint
934entrypoint:
935  lddw r1, message
936  mov64 r2, MSG_LEN
937  call sol_log_
938  exit
939.rodata
940  message: .ascii "Hello, Solana!"
941"#;
942        let result = assemble_with_debug_data(source, "hello_solana.s", "/tmp");
943        assert!(result.is_ok());
944        let bytecode = result.unwrap();
945
946        // Verify the ELF has all debug sections.
947        let bytecode_str = String::from_utf8_lossy(&bytecode);
948        assert!(
949            bytecode_str.contains(".debug_abbrev"),
950            "Missing .debug_abbrev section"
951        );
952        assert!(
953            bytecode_str.contains(".debug_info"),
954            "Missing .debug_info section"
955        );
956        assert!(
957            bytecode_str.contains(".debug_line"),
958            "Missing .debug_line section"
959        );
960        assert!(
961            bytecode_str.contains(".debug_line_str"),
962            "Missing .debug_line_str section"
963        );
964    }
965}