Skip to main content

sbpf_assembler/preprocessor/
expand.rs

1use {
2    super::{
3        SourceLine,
4        macro_def::{MacroDef, scan_macro_definitions},
5        source_map::SourceOrigin,
6    },
7    crate::errors::CompileError,
8    std::{
9        collections::HashMap,
10        sync::atomic::{AtomicU64, Ordering},
11    },
12};
13
14const MAX_EXPANSION_DEPTH: u32 = 100;
15
16/// Global expansion counter for `\@` unique IDs
17static EXPANSION_COUNTER: AtomicU64 = AtomicU64::new(0);
18
19/// An expansion error paired with its source origin
20#[derive(Debug)]
21pub(crate) struct ExpandError {
22    pub error: CompileError,
23    pub origin: Option<SourceOrigin>,
24}
25
26/// Expand all macros, `.rept`, and `.irp` directives in the given lines.
27///
28/// Returns the expanded lines and any errors encountered.
29pub(crate) fn expand_macros(
30    lines: Vec<SourceLine>,
31) -> Result<(Vec<SourceLine>, Vec<ExpandError>), Vec<ExpandError>> {
32    // Scan for macro definitions and separate them from regular lines
33    let scan_result = scan_macro_definitions(lines);
34
35    if !scan_result.errors.is_empty() {
36        return Err(scan_result
37            .errors
38            .into_iter()
39            .map(|(error, origin)| ExpandError {
40                error,
41                origin: Some(origin),
42            })
43            .collect());
44    }
45
46    let macros = scan_result.macros;
47    let mut errors = Vec::new();
48    let mut output = Vec::new();
49
50    // Expand macro invocations in remaining lines
51    for line in scan_result.remaining_lines {
52        expand_line(&line, &macros, &mut output, &mut errors, 0);
53    }
54
55    // Now, handle .rept and .irp on the post-macro output
56    let output = expand_repetitions(output)?;
57
58    Ok((output, errors))
59}
60
61/// Expand a single line, checking if it's a macro invocation.
62fn expand_line(
63    line: &SourceLine,
64    macros: &HashMap<String, MacroDef>,
65    output: &mut Vec<SourceLine>,
66    errors: &mut Vec<ExpandError>,
67    depth: u32,
68) {
69    if depth > MAX_EXPANSION_DEPTH {
70        errors.push(ExpandError {
71            error: CompileError::MacroRecursionLimit {
72                limit: MAX_EXPANSION_DEPTH,
73                span: 0..0,
74                custom_label: None,
75            },
76            origin: Some(line.origin.clone()),
77        });
78        return;
79    }
80
81    let trimmed = line.text.trim();
82
83    // Check if the first token matches a known macro name
84    let first_token = first_token(trimmed);
85    if let Some(macro_def) = first_token.and_then(|name| macros.get(name)) {
86        let args_str = trimmed[macro_def.name.len()..].trim();
87        let args = split_args(args_str);
88
89        match bind_args(macro_def, &args) {
90            Ok(bindings) => {
91                let expansion_id = EXPANSION_COUNTER.fetch_add(1, Ordering::Relaxed);
92
93                // Expand each body line with parameter substitution
94                for body_line in &macro_def.body_lines {
95                    let expanded_text = substitute(body_line, &bindings, expansion_id);
96                    let expanded_line = SourceLine {
97                        text: expanded_text,
98                        origin: SourceOrigin::with_macro_expansion(
99                            macro_def.defined_at.file_id,
100                            macro_def.defined_at.line,
101                            macro_def.name.clone(),
102                            line.origin.clone(),
103                            depth + 1,
104                        ),
105                    };
106
107                    // Rescan for further macro invocations
108                    expand_line(&expanded_line, macros, output, errors, depth + 1);
109                }
110            }
111            Err(e) => errors.push(ExpandError {
112                error: e,
113                origin: Some(line.origin.clone()),
114            }),
115        }
116    } else {
117        // Not a macro invocation, pass through
118        output.push(line.clone());
119    }
120}
121
122/// Extract the first whitespace-delimited token from a line.
123/// Skips labels (tokens ending with ':').
124fn first_token(line: &str) -> Option<&str> {
125    let trimmed = line.trim();
126    if trimmed.is_empty() || trimmed.starts_with('#') || trimmed.starts_with(';') {
127        return None;
128    }
129
130    // If line starts with '.', it's a directive, not a macro
131    if trimmed.starts_with('.') {
132        return None;
133    }
134
135    let end = trimmed
136        .find(|c: char| c.is_whitespace() || c == ',')
137        .unwrap_or(trimmed.len());
138    let token = &trimmed[..end];
139
140    // Skip labels (they end with ':')
141    if token.ends_with(':') {
142        // Check if there's something after the label on the same line
143        let rest = trimmed[end..].trim();
144        if rest.is_empty() {
145            return None;
146        }
147        return first_token(rest);
148    }
149
150    Some(token)
151}
152
153/// Split arguments by top-level commas, respecting quotes and brackets.
154fn split_args(args_str: &str) -> Vec<String> {
155    if args_str.is_empty() {
156        return Vec::new();
157    }
158
159    let mut args = Vec::new();
160    let mut current = String::new();
161    let mut depth = 0i32; // bracket depth
162    let mut in_string = false;
163    let mut escape_next = false;
164
165    for ch in args_str.chars() {
166        if escape_next {
167            current.push(ch);
168            escape_next = false;
169            continue;
170        }
171
172        match ch {
173            '\\' if in_string => {
174                current.push(ch);
175                escape_next = true;
176            }
177            '"' => {
178                in_string = !in_string;
179                current.push(ch);
180            }
181            '(' | '[' if !in_string => {
182                depth += 1;
183                current.push(ch);
184            }
185            ')' | ']' if !in_string => {
186                depth -= 1;
187                current.push(ch);
188            }
189            ',' if !in_string && depth == 0 => {
190                args.push(current.trim().to_string());
191                current = String::new();
192            }
193            _ => {
194                current.push(ch);
195            }
196        }
197    }
198
199    let last = current.trim().to_string();
200    if !last.is_empty() {
201        args.push(last);
202    }
203
204    args
205}
206
207/// Bind arguments to macro parameters, producing a name->value map.
208fn bind_args(
209    macro_def: &MacroDef,
210    args: &[String],
211) -> Result<HashMap<String, String>, CompileError> {
212    let mut bindings = HashMap::new();
213
214    let required_count = macro_def
215        .params
216        .iter()
217        .filter(|p| p.default.is_none() && !p.is_vararg)
218        .count();
219
220    let has_vararg = macro_def.params.iter().any(|p| p.is_vararg);
221
222    // Check argument count
223    if !has_vararg && args.len() > macro_def.params.len() {
224        return Err(CompileError::MacroArgCount {
225            name: macro_def.name.clone(),
226            expected: macro_def.params.len(),
227            got: args.len(),
228            span: 0..0,
229            custom_label: None,
230        });
231    }
232    if args.len() < required_count {
233        return Err(CompileError::MacroArgCount {
234            name: macro_def.name.clone(),
235            expected: required_count,
236            got: args.len(),
237            span: 0..0,
238            custom_label: None,
239        });
240    }
241
242    for (i, param) in macro_def.params.iter().enumerate() {
243        if param.is_vararg {
244            // Collect all remaining arguments
245            let vararg_values: Vec<&str> = args[i..].iter().map(|s| s.as_str()).collect();
246            bindings.insert(param.name.clone(), vararg_values.join(", "));
247        } else if i < args.len() {
248            bindings.insert(param.name.clone(), args[i].clone());
249        } else if let Some(ref default) = param.default {
250            bindings.insert(param.name.clone(), default.clone());
251        }
252        // required params without args already caught above
253    }
254
255    // Also add positional references (\1, \2, etc.)
256    for (i, arg) in args.iter().enumerate() {
257        bindings.insert((i + 1).to_string(), arg.clone());
258    }
259
260    Ok(bindings)
261}
262
263/// Perform parameter substitution on a single body line.
264///
265/// Handles:
266/// - `\name` -> argument value
267/// - `\@` -> unique expansion ID
268/// - `\()` -> zero-width concatenation (disappears after adjacent substitution)
269/// - `\\` -> literal `\`
270fn substitute(line: &str, bindings: &HashMap<String, String>, expansion_id: u64) -> String {
271    let mut result = String::with_capacity(line.len());
272    let chars: Vec<char> = line.chars().collect();
273    let len = chars.len();
274    let mut i = 0;
275
276    while i < len {
277        if chars[i] == '\\' && i + 1 < len {
278            let next = chars[i + 1];
279
280            if next == '\\' {
281                // Escaped backslash
282                result.push('\\');
283                i += 2;
284            } else if next == '@' {
285                // Unique expansion ID
286                result.push_str(&expansion_id.to_string());
287                i += 2;
288            } else if next == '(' && i + 2 < len && chars[i + 2] == ')' {
289                // Token concatenation -- just skip it (zero-width separator)
290                i += 3;
291            } else if next.is_alphanumeric() || next == '_' {
292                // Parameter reference: collect the name
293                let start = i + 1;
294                let mut end = start;
295                while end < len && (chars[end].is_alphanumeric() || chars[end] == '_') {
296                    end += 1;
297                }
298                let param_name: String = chars[start..end].iter().collect();
299
300                if let Some(value) = bindings.get(&param_name) {
301                    result.push_str(value);
302                    // Check if followed by \() for concatenation
303                    if end + 2 < len
304                        && chars[end] == '\\'
305                        && chars[end + 1] == '('
306                        && chars[end + 2] == ')'
307                    {
308                        // Skip the \() separator
309                        end += 3;
310                    }
311                } else {
312                    // Unknown parameter -- keep as-is
313                    result.push('\\');
314                    result.push_str(&param_name);
315                }
316                i = end;
317            } else {
318                // Unknown escape, keep as-is
319                result.push('\\');
320                i += 1;
321            }
322        } else {
323            result.push(chars[i]);
324            i += 1;
325        }
326    }
327
328    result
329}
330
331/// Expand `.rept` and `.irp` blocks (these are processed before macro expansion).
332fn expand_repetitions(lines: Vec<SourceLine>) -> Result<Vec<SourceLine>, Vec<ExpandError>> {
333    let mut output = Vec::new();
334    let mut errors = Vec::new();
335    let mut i = 0;
336
337    while i < lines.len() {
338        let trimmed = lines[i].text.trim();
339
340        if let Some(count_str) = trimmed.strip_prefix(".rept").and_then(|r| {
341            if r.starts_with(char::is_whitespace) {
342                Some(r.trim())
343            } else {
344                None
345            }
346        }) {
347            // Find matching .endr
348            let (body, end_idx) = find_endr_block(&lines, i + 1, &lines[i].origin)?;
349
350            match count_str.parse::<usize>() {
351                Ok(count) => {
352                    // Recursively expand nested .rept/.irp inside the body
353                    let expanded_body = expand_repetitions(body)?;
354                    for _ in 0..count {
355                        output.extend(expanded_body.iter().cloned());
356                    }
357                }
358                Err(_) => {
359                    errors.push(ExpandError {
360                        error: CompileError::InvalidReptCount {
361                            value: count_str.to_string(),
362                            span: 0..0,
363                            custom_label: None,
364                        },
365                        origin: Some(lines[i].origin.clone()),
366                    });
367                }
368            }
369
370            i = end_idx + 1;
371        } else if let Some(rest) = trimmed.strip_prefix(".irp").and_then(|r| {
372            if r.starts_with(char::is_whitespace) {
373                Some(r.trim())
374            } else {
375                None
376            }
377        }) {
378            // Parse: .irp var, val1, val2, val3
379            let (body, end_idx) = find_endr_block(&lines, i + 1, &lines[i].origin)?;
380
381            if let Some((var_name, values_str)) = rest.split_once(',') {
382                let var_name = var_name.trim();
383                let values = split_args(values_str.trim());
384
385                for value in &values {
386                    let mut bindings = HashMap::new();
387                    bindings.insert(var_name.to_string(), value.clone());
388
389                    // Substitute for this iteration first, then recursively
390                    // expand any nested .rept/.irp in the result.
391                    let mut iter_lines = Vec::with_capacity(body.len());
392                    for body_line in &body {
393                        let expansion_id = EXPANSION_COUNTER.fetch_add(1, Ordering::Relaxed);
394                        let expanded_text = substitute(&body_line.text, &bindings, expansion_id);
395                        iter_lines.push(SourceLine {
396                            text: expanded_text,
397                            origin: body_line.origin.clone(),
398                        });
399                    }
400                    output.extend(expand_repetitions(iter_lines)?);
401                }
402            }
403
404            i = end_idx + 1;
405        } else {
406            output.push(lines[i].clone());
407            i += 1;
408        }
409    }
410
411    if errors.is_empty() {
412        Ok(output)
413    } else {
414        Err(errors)
415    }
416}
417
418/// Find the matching `.endr` for a `.rept` or `.irp` block, handling nesting.
419fn find_endr_block(
420    lines: &[SourceLine],
421    start: usize,
422    directive_origin: &SourceOrigin,
423) -> Result<(Vec<SourceLine>, usize), Vec<ExpandError>> {
424    let mut depth = 1u32;
425    let mut i = start;
426
427    while i < lines.len() {
428        let trimmed = lines[i].text.trim();
429        if trimmed == ".endr" {
430            depth -= 1;
431            if depth == 0 {
432                let body = lines[start..i].to_vec();
433                return Ok((body, i));
434            }
435        } else if trimmed.starts_with(".rept") || trimmed.starts_with(".irp") {
436            depth += 1;
437        }
438        i += 1;
439    }
440
441    Err(vec![ExpandError {
442        error: CompileError::UnclosedRept {
443            span: 0..0,
444            custom_label: None,
445        },
446        origin: Some(directive_origin.clone()),
447    }])
448}
449
450#[cfg(test)]
451mod tests {
452    use {super::*, crate::preprocessor::source_map::FileId};
453
454    fn make_line(text: &str, line: u32) -> SourceLine {
455        SourceLine {
456            text: text.to_string(),
457            origin: SourceOrigin::new(FileId(0), line),
458        }
459    }
460
461    fn expand_and_collect(lines: Vec<SourceLine>) -> Vec<String> {
462        let (result, errors) = expand_macros(lines).unwrap();
463        assert!(errors.is_empty(), "Unexpected errors: {:?}", errors);
464        result.iter().map(|l| l.text.clone()).collect()
465    }
466
467    #[test]
468    fn test_first_token() {
469        assert_eq!(first_token("MY_MACRO arg1, arg2"), Some("MY_MACRO"));
470        assert_eq!(first_token("  MY_MACRO  "), Some("MY_MACRO"));
471        assert_eq!(first_token(".globl entry"), None); // directive
472        assert_eq!(first_token("# comment"), None);
473        assert_eq!(first_token(""), None);
474        assert_eq!(first_token("label:"), None);
475        assert_eq!(first_token("label: MY_MACRO"), Some("MY_MACRO"));
476    }
477
478    #[test]
479    fn test_split_args() {
480        assert_eq!(split_args("a, b, c"), vec!["a", "b", "c"]);
481        assert_eq!(split_args("r1, 42"), vec!["r1", "42"]);
482        assert_eq!(
483            split_args("\"hello, world\", 1"),
484            vec!["\"hello, world\"", "1"]
485        );
486        assert_eq!(split_args(""), Vec::<String>::new());
487        assert_eq!(split_args("single"), vec!["single"]);
488    }
489
490    #[test]
491    fn test_substitute_simple() {
492        let mut bindings = HashMap::new();
493        bindings.insert("reg".to_string(), "r1".to_string());
494        bindings.insert("val".to_string(), "42".to_string());
495
496        assert_eq!(
497            substitute("    mov64 \\reg, \\val", &bindings, 0),
498            "    mov64 r1, 42"
499        );
500    }
501
502    #[test]
503    fn test_substitute_unique_id() {
504        let bindings = HashMap::new();
505        assert_eq!(substitute("label_\\@:", &bindings, 5), "label_5:");
506    }
507
508    #[test]
509    fn test_substitute_concatenation() {
510        let mut bindings = HashMap::new();
511        bindings.insert("msg".to_string(), "e1".to_string());
512
513        assert_eq!(substitute("\\msg\\()_end", &bindings, 0), "e1_end");
514    }
515
516    #[test]
517    fn test_substitute_escaped_backslash() {
518        let bindings = HashMap::new();
519        assert_eq!(substitute("\\\\n", &bindings, 0), "\\n");
520    }
521
522    #[test]
523    fn test_simple_macro_expansion() {
524        let lines = vec![
525            make_line(".macro NOP", 1),
526            make_line("    mov64 r0, 0", 2),
527            make_line(".endm", 3),
528            make_line("NOP", 4),
529        ];
530
531        let result = expand_and_collect(lines);
532        assert_eq!(result, vec!["    mov64 r0, 0"]);
533    }
534
535    #[test]
536    fn test_macro_with_args() {
537        let lines = vec![
538            make_line(".macro MOV dst, src", 1),
539            make_line("    mov64 \\dst, \\src", 2),
540            make_line(".endm", 3),
541            make_line("MOV r1, r2", 4),
542        ];
543
544        let result = expand_and_collect(lines);
545        assert_eq!(result, vec!["    mov64 r1, r2"]);
546    }
547
548    #[test]
549    fn test_macro_with_default() {
550        let lines = vec![
551            make_line(".macro SET reg, val=0", 1),
552            make_line("    mov64 \\reg, \\val", 2),
553            make_line(".endm", 3),
554            make_line("SET r1", 4),
555            make_line("SET r2, 42", 5),
556        ];
557
558        let result = expand_and_collect(lines);
559        assert_eq!(result, vec!["    mov64 r1, 0", "    mov64 r2, 42"]);
560    }
561
562    #[test]
563    fn test_macro_unique_id() {
564        let lines = vec![
565            make_line(".macro LOOP", 1),
566            make_line("loop_\\@:", 2),
567            make_line(".endm", 3),
568            make_line("LOOP", 4),
569            make_line("LOOP", 5),
570        ];
571
572        let result = expand_and_collect(lines);
573        assert_eq!(result.len(), 2);
574        assert!(result.iter().all(|s| {
575            s.starts_with("loop_")
576                && s.ends_with(':')
577                && s["loop_".len()..s.len() - 1].parse::<u64>().is_ok()
578        }));
579        assert_ne!(result[0], result[1]);
580    }
581
582    #[test]
583    fn test_macro_concatenation() {
584        let lines = vec![
585            make_line(".macro DEF_STR name, text", 1),
586            make_line("\\name:", 2),
587            make_line("    .ascii \\text", 3),
588            make_line("\\name\\()_end:", 4),
589            make_line(".endm", 5),
590            make_line("DEF_STR e1, \"error\"", 6),
591        ];
592
593        let result = expand_and_collect(lines);
594        assert_eq!(result, vec!["e1:", "    .ascii \"error\"", "e1_end:"]);
595    }
596
597    #[test]
598    fn test_nested_macro_expansion() {
599        let lines = vec![
600            make_line(".macro INNER val", 1),
601            make_line("    mov64 r0, \\val", 2),
602            make_line(".endm", 3),
603            make_line(".macro OUTER val", 4),
604            make_line("INNER \\val", 5),
605            make_line(".endm", 6),
606            make_line("OUTER 42", 7),
607        ];
608
609        let result = expand_and_collect(lines);
610        assert_eq!(result, vec!["    mov64 r0, 42"]);
611    }
612
613    #[test]
614    fn test_recursion_limit() {
615        let lines = vec![
616            make_line(".macro FOREVER", 1),
617            make_line("FOREVER", 2),
618            make_line(".endm", 3),
619            make_line("FOREVER", 4),
620        ];
621
622        let (_, errors) = expand_macros(lines).unwrap();
623        assert!(!errors.is_empty());
624        assert!(matches!(
625            &errors[0].error,
626            CompileError::MacroRecursionLimit { .. }
627        ));
628        // Should have origin pointing to the invocation
629        assert!(errors[0].origin.is_some());
630    }
631
632    #[test]
633    fn test_wrong_arg_count() {
634        let lines = vec![
635            make_line(".macro NEED_TWO a, b", 1),
636            make_line("    mov64 \\a, \\b", 2),
637            make_line(".endm", 3),
638            make_line("NEED_TWO r1", 4),
639        ];
640
641        let (_, errors) = expand_macros(lines).unwrap();
642        assert!(!errors.is_empty());
643        assert!(matches!(
644            &errors[0].error,
645            CompileError::MacroArgCount {
646                expected: 2,
647                got: 1,
648                ..
649            }
650        ));
651        assert!(errors[0].origin.is_some());
652        assert_eq!(errors[0].origin.as_ref().unwrap().line, 4);
653    }
654
655    #[test]
656    fn test_vararg() {
657        let lines = vec![
658            make_line(".macro LOG fmt, args:vararg", 1),
659            make_line("    .ascii \\fmt", 2),
660            make_line("    .ascii \\args", 3),
661            make_line(".endm", 4),
662            make_line("LOG \"hello\", a, b, c", 5),
663        ];
664
665        let result = expand_and_collect(lines);
666        assert_eq!(result, vec!["    .ascii \"hello\"", "    .ascii a, b, c"]);
667    }
668
669    #[test]
670    fn test_rept() {
671        let lines = vec![
672            make_line(".rept 3", 1),
673            make_line("    nop", 2),
674            make_line(".endr", 3),
675        ];
676
677        let result = expand_and_collect(lines);
678        assert_eq!(result, vec!["    nop", "    nop", "    nop"]);
679    }
680
681    #[test]
682    fn test_rept_zero() {
683        let lines = vec![
684            make_line(".rept 0", 1),
685            make_line("    nop", 2),
686            make_line(".endr", 3),
687        ];
688
689        let result = expand_and_collect(lines);
690        assert!(result.is_empty());
691    }
692
693    #[test]
694    fn test_irp() {
695        let lines = vec![
696            make_line(".irp reg, r1, r2, r3", 1),
697            make_line("    mov64 \\reg, 0", 2),
698            make_line(".endr", 3),
699        ];
700
701        let result = expand_and_collect(lines);
702        assert_eq!(
703            result,
704            vec!["    mov64 r1, 0", "    mov64 r2, 0", "    mov64 r3, 0"]
705        );
706    }
707
708    #[test]
709    fn test_nested_rept() {
710        let lines = vec![
711            make_line(".rept 2", 1),
712            make_line("    .rept 3", 2),
713            make_line("        mov64 r1, 0x1", 3),
714            make_line("    .endr", 4),
715            make_line(".endr", 5),
716        ];
717
718        let result = expand_and_collect(lines);
719        assert_eq!(result.len(), 6);
720        assert!(result.iter().all(|s| s == "        mov64 r1, 0x1"));
721    }
722
723    #[test]
724    fn test_nested_irp() {
725        let lines = vec![
726            make_line(".irp reg, r1, r2", 1),
727            make_line("    .irp val, 0x1, 0x2", 2),
728            make_line("        mov64 \\reg, \\val", 3),
729            make_line("    .endr", 4),
730            make_line(".endr", 5),
731        ];
732
733        let result = expand_and_collect(lines);
734        assert_eq!(
735            result,
736            vec![
737                "        mov64 r1, 0x1",
738                "        mov64 r1, 0x2",
739                "        mov64 r2, 0x1",
740                "        mov64 r2, 0x2",
741            ]
742        );
743    }
744
745    #[test]
746    fn test_rept_inside_irp() {
747        let lines = vec![
748            make_line(".irp r, r1, r2", 1),
749            make_line("    .rept 2", 2),
750            make_line("        mov64 \\r, 0x1", 3),
751            make_line("    .endr", 4),
752            make_line(".endr", 5),
753        ];
754
755        let result = expand_and_collect(lines);
756        assert_eq!(
757            result,
758            vec![
759                "        mov64 r1, 0x1",
760                "        mov64 r1, 0x1",
761                "        mov64 r2, 0x1",
762                "        mov64 r2, 0x1",
763            ]
764        );
765    }
766
767    #[test]
768    fn test_irp_inside_rept() {
769        let lines = vec![
770            make_line(".rept 2", 1),
771            make_line("    .irp r, r1, r2", 2),
772            make_line("        mov64 \\r, 0x1", 3),
773            make_line("    .endr", 4),
774            make_line(".endr", 5),
775        ];
776
777        let result = expand_and_collect(lines);
778        assert_eq!(
779            result,
780            vec![
781                "        mov64 r1, 0x1",
782                "        mov64 r2, 0x1",
783                "        mov64 r1, 0x1",
784                "        mov64 r2, 0x1",
785            ]
786        );
787    }
788
789    #[test]
790    fn test_rept_count_from_macro_param() {
791        let lines = vec![
792            make_line(".macro TEST num", 1),
793            make_line("    .rept \\num", 2),
794            make_line("        mov64 r1, 0x123", 3),
795            make_line("    .endr", 4),
796            make_line(".endm", 5),
797            make_line("TEST 3", 7),
798        ];
799
800        let result = expand_and_collect(lines);
801        assert_eq!(
802            result,
803            vec![
804                "        mov64 r1, 0x123",
805                "        mov64 r1, 0x123",
806                "        mov64 r1, 0x123",
807            ]
808        );
809    }
810
811    #[test]
812    fn test_irp_values_from_macro_vararg() {
813        let lines = vec![
814            make_line(".macro LOAD regs:vararg", 1),
815            make_line("    .irp r, \\regs", 2),
816            make_line("        mov64 \\r, 0x1", 3),
817            make_line("    .endr", 4),
818            make_line(".endm", 5),
819            make_line("LOAD r1, r2, r3", 7),
820        ];
821
822        let result = expand_and_collect(lines);
823        assert_eq!(
824            result,
825            vec![
826                "        mov64 r1, 0x1",
827                "        mov64 r2, 0x1",
828                "        mov64 r3, 0x1",
829            ]
830        );
831    }
832
833    #[test]
834    fn test_unclosed_rept() {
835        let lines = vec![make_line(".rept 3", 1), make_line("    nop", 2)];
836
837        let result = expand_macros(lines);
838        assert!(result.is_err());
839        let errors = result.unwrap_err();
840        assert!(matches!(
841            &errors[0].error,
842            CompileError::UnclosedRept { .. }
843        ));
844        assert!(errors[0].origin.is_some());
845    }
846
847    #[test]
848    fn test_full_example_from_spec() {
849        let lines = vec![
850            make_line(".macro DEF_STR name, text", 1),
851            make_line("\\name:", 2),
852            make_line("    .ascii \\text", 3),
853            make_line("\\name\\()_end:", 4),
854            make_line(".endm", 5),
855            make_line("", 6),
856            make_line(".macro RETURN_ERR code, msg", 7),
857            make_line("    lddw r0, \\code", 8),
858            make_line("    lddw r1, \\msg", 9),
859            make_line("    lddw r2, \\msg\\()_end - \\msg", 10),
860            make_line("    call sol_log_", 11),
861            make_line("    exit", 12),
862            make_line(".endm", 13),
863            make_line("", 14),
864            make_line("DEF_STR e1, \"error\"", 15),
865            make_line("", 16),
866            make_line("RETURN_ERR 1, e1", 17),
867        ];
868
869        let result = expand_and_collect(lines);
870        assert_eq!(
871            result,
872            vec![
873                "", // blank line between definitions
874                "", // blank line after definitions
875                "e1:",
876                "    .ascii \"error\"",
877                "e1_end:",
878                "", // blank line
879                "    lddw r0, 1",
880                "    lddw r1, e1",
881                "    lddw r2, e1_end - e1",
882                "    call sol_log_",
883                "    exit",
884            ]
885        );
886    }
887
888    #[test]
889    fn test_passthrough_no_macros() {
890        let lines = vec![
891            make_line(".globl entrypoint", 1),
892            make_line("entrypoint:", 2),
893            make_line("    mov64 r1, 42", 3),
894            make_line("    exit", 4),
895        ];
896
897        let result = expand_and_collect(lines);
898        assert_eq!(
899            result,
900            vec![
901                ".globl entrypoint",
902                "entrypoint:",
903                "    mov64 r1, 42",
904                "    exit",
905            ]
906        );
907    }
908}