Skip to main content

singe_ptx/parser/
mod.rs

1mod common;
2mod directive;
3mod function;
4mod instruction;
5mod module;
6mod operand;
7
8use crate::{
9    ast::Module,
10    error::{ParseError, ParseErrorKind},
11};
12
13/// Parses a complete PTX module.
14///
15/// The input must contain the module directives required by the PTX grammar
16/// supported by this crate. Leading and trailing whitespace is accepted, but any
17/// non-whitespace input left after the parsed module is reported as
18/// [`ParseErrorKind::TrailingInput`]. Other parse failures report the byte offset
19/// of the unexpected token.
20///
21/// This parser is intentionally syntax-focused: it builds the AST and reports
22/// parse errors, but it does not validate target-specific instruction semantics.
23///
24/// # Errors
25///
26/// Returns an error if `input` is not a complete PTX module accepted by this
27/// parser or if non-whitespace trailing input remains.
28pub fn parse_module(input: &str) -> Result<Module, ParseError> {
29    match module::module(input) {
30        Ok((remaining, module)) => {
31            let remaining = remaining.trim();
32            if remaining.is_empty() {
33                Ok(module)
34            } else {
35                Err(ParseError {
36                    kind: ParseErrorKind::TrailingInput,
37                    offset: input.len() - remaining.len(),
38                })
39            }
40        }
41        Err(e) => {
42            let offset = match &e {
43                nom::Err::Error(e) | nom::Err::Failure(e) => input.len() - e.input.len(),
44                nom::Err::Incomplete(_) => input.len(),
45            };
46            Err(ParseError {
47                kind: ParseErrorKind::UnexpectedToken(input[offset..].chars().take(20).collect()),
48                offset,
49            })
50        }
51    }
52}
53
54pub(crate) fn offset(source: &str, input: &str) -> usize {
55    source.len() - input.len()
56}
57
58#[cfg(test)]
59mod tests {
60    use super::*;
61
62    use crate::ast::*;
63
64    #[test]
65    fn test_parse_no_whitespace() {
66        let src = ".version 9.1\n.target sm_90\n.address_size 64";
67        let module = parse_module(src).unwrap();
68        assert_eq!(module.version.major, 9);
69        assert_eq!(module.version.minor, 1);
70    }
71
72    #[test]
73    fn test_parse_minimal_module() {
74        let src = r#"
75            .version 9.1
76            .target sm_90
77            .address_size 64
78        "#;
79
80        let module = parse_module(src).unwrap();
81        assert_eq!(module.version.major, 9);
82        assert_eq!(module.version.minor, 1);
83        assert_eq!(module.target.specifiers, vec!["sm_90"]);
84        assert_eq!(module.address_size, AddressSize::Bits64);
85        assert!(module.items.is_empty());
86    }
87
88    #[test]
89    fn test_parse_multiple_targets() {
90        let src = r#"
91            .version 9.1
92            .target sm_90, texmode_unified
93            .address_size 64
94        "#;
95
96        let module = parse_module(src).unwrap();
97        assert_eq!(module.target.specifiers, vec!["sm_90", "texmode_unified"]);
98    }
99
100    #[test]
101    fn test_parse_32bit_address() {
102        let src = r#"
103            .version 8.5
104            .target sm_80
105            .address_size 32
106        "#;
107
108        let module = parse_module(src).unwrap();
109        assert_eq!(module.address_size, AddressSize::Bits32);
110    }
111
112    #[test]
113    fn test_parse_default_32bit_address_when_omitted() {
114        let src = r#"
115            .version 9.1
116            .target sm_90
117        "#;
118
119        let module = parse_module(src).unwrap();
120        assert_eq!(module.address_size, AddressSize::Bits32);
121    }
122
123    #[test]
124    fn test_parse_with_comments() {
125        let src = r#"
126            // PTX module
127            .version 9.1
128            .target sm_90  // target arch
129            .address_size 64
130            /* end of preamble */
131        "#;
132
133        let module = parse_module(src).unwrap();
134        assert_eq!(module.version.major, 9);
135    }
136
137    #[test]
138    fn test_parse_empty_entry_function() {
139        let src = r#"
140            .version 9.1
141            .target sm_90
142            .address_size 64
143
144            .entry kernel() {
145            }
146        "#;
147
148        let module = parse_module(src).unwrap();
149        assert_eq!(module.items.len(), 1);
150        match &module.items[0] {
151            TopLevelItem::Function(f) => {
152                assert!(f.entry);
153                assert_eq!(f.name, "kernel");
154                assert!(f.params.is_empty());
155                assert!(f.body.as_ref().unwrap().is_empty());
156            }
157            other => panic!("expected function, got {:?}", other),
158        }
159    }
160
161    #[test]
162    fn test_parse_entry_with_instructions() {
163        let src = r#"
164            .version 9.1
165            .target sm_90
166            .address_size 64
167
168            .entry kernel() {
169                add.s32 %r1, %r2, %r3;
170                ret;
171            }
172        "#;
173
174        let module = parse_module(src).unwrap();
175        let func = match &module.items[0] {
176            TopLevelItem::Function(f) => f,
177            other => panic!("expected function, got {:?}", other),
178        };
179        let body = func.body.as_ref().unwrap();
180        assert_eq!(body.len(), 2);
181
182        match &body[0] {
183            Statement::Instruction(inst) => {
184                assert_eq!(inst.opcode, "add");
185                assert_eq!(inst.modifiers, vec!["s32"]);
186                assert_eq!(inst.operands.len(), 3);
187            }
188            other => panic!("expected instruction, got {:?}", other),
189        }
190
191        match &body[1] {
192            Statement::Instruction(inst) => {
193                assert_eq!(inst.opcode, "ret");
194                assert!(inst.modifiers.is_empty());
195                assert!(inst.operands.is_empty());
196            }
197            other => panic!("expected instruction, got {:?}", other),
198        }
199    }
200
201    #[test]
202    fn test_parse_predicated_instruction() {
203        let src = r#"
204            .version 9.1
205            .target sm_90
206            .address_size 64
207
208            .entry kernel() {
209                @%p1 add.s32 %r1, %r2, %r3;
210                @!%p2 bra label;
211            }
212        "#;
213
214        let module = parse_module(src).unwrap();
215        let func = match &module.items[0] {
216            TopLevelItem::Function(f) => f,
217            other => panic!("expected function, got {:?}", other),
218        };
219        let body = func.body.as_ref().unwrap();
220
221        match &body[0] {
222            Statement::Instruction(inst) => {
223                let guard = inst.guard.as_ref().unwrap();
224                assert!(!guard.negated);
225                assert_eq!(guard.register, "%p1");
226            }
227            other => panic!("expected instruction, got {:?}", other),
228        }
229
230        match &body[1] {
231            Statement::Instruction(inst) => {
232                let guard = inst.guard.as_ref().unwrap();
233                assert!(guard.negated);
234                assert_eq!(guard.register, "%p2");
235                assert_eq!(inst.opcode, "bra");
236                match &inst.operands[0] {
237                    Operand::Label { name, .. } => assert_eq!(name, "label"),
238                    other => panic!("expected label, got {:?}", other),
239                }
240            }
241            other => panic!("expected instruction, got {:?}", other),
242        }
243    }
244
245    #[test]
246    fn test_parse_function_with_params() {
247        let src = r#"
248            .version 9.1
249            .target sm_90
250            .address_size 64
251
252            .entry foo(.param .b32 N, .param .b8 buffer[64]) {
253                ret;
254            }
255        "#;
256
257        let module = parse_module(src).unwrap();
258        let func = match &module.items[0] {
259            TopLevelItem::Function(f) => f,
260            other => panic!("expected function, got {:?}", other),
261        };
262        assert_eq!(func.params.len(), 2);
263        assert_eq!(func.params[0].name, "N");
264        assert_eq!(func.params[1].name, "buffer");
265        assert_eq!(func.params[1].array_bounds, vec![Some(64)]);
266    }
267
268    #[test]
269    fn test_parse_func_without_return_params() {
270        let src = r#"
271            .version 9.1
272            .target sm_90
273            .address_size 64
274
275            .func foo(.param .b32 x) {
276                ret;
277            }
278        "#;
279
280        let module = parse_module(src).unwrap();
281        let func = match &module.items[0] {
282            TopLevelItem::Function(func) => func,
283            other => panic!("expected function, got {:?}", other),
284        };
285
286        assert!(func.return_params.is_empty());
287        assert_eq!(func.params.len(), 1);
288        assert_eq!(func.params[0].name, "x");
289    }
290
291    #[test]
292    fn test_parse_function_with_return_params_and_directives() {
293        let src = r#"
294            .version 9.1
295            .target sm_90
296            .address_size 64
297
298            .func (.reg .b32 rval) foo(.reg .b32 x) .noreturn .abi_preserve 8 .abi_preserve_control 2;
299        "#;
300
301        let module = parse_module(src).unwrap();
302        let func = match &module.items[0] {
303            TopLevelItem::Function(func) => func,
304            other => panic!("expected function, got {:?}", other),
305        };
306
307        assert_eq!(func.return_params.len(), 1);
308        assert_eq!(func.return_params[0].name, "rval");
309        assert_eq!(func.params.len(), 1);
310        assert_eq!(func.directives.len(), 3);
311        assert!(matches!(
312            func.directives[0],
313            FunctionDirective::NoReturn { .. }
314        ));
315        assert!(matches!(
316            func.directives[1],
317            FunctionDirective::AbiPreserve { value: 8, .. }
318        ));
319        assert!(matches!(
320            func.directives[2],
321            FunctionDirective::AbiPreserveControl { value: 2, .. }
322        ));
323    }
324
325    #[test]
326    fn test_parse_label_and_branch() {
327        let src = r#"
328            .version 9.1
329            .target sm_90
330            .address_size 64
331
332            .entry kernel() {
333            start:
334                mov.b32 %r1, %r2;
335                bra start;
336            }
337        "#;
338
339        let module = parse_module(src).unwrap();
340        let func = match &module.items[0] {
341            TopLevelItem::Function(f) => f,
342            other => panic!("expected function, got {:?}", other),
343        };
344        let body = func.body.as_ref().unwrap();
345
346        match &body[0] {
347            Statement::Label { name, .. } => assert_eq!(name, "start"),
348            other => panic!("expected label, got {:?}", other),
349        }
350    }
351
352    #[test]
353    fn test_parse_variable_initializers_and_multidimensional_arrays() {
354        let src = r#"
355            .version 9.1
356            .target sm_90
357            .address_size 64
358
359            .visible .global .align 8 .u64 g_data[2] = {1, 2};
360            .global .s32 offset[][2] = {{-1, 0}, {0, -1}};
361            .const .u32 p = generic(g_data) + 8;
362        "#;
363
364        let module = parse_module(src).unwrap();
365        assert_eq!(module.items.len(), 3);
366
367        let variable = match &module.items[0] {
368            TopLevelItem::Variable(variable) => variable,
369            other => panic!("expected variable, got {:?}", other),
370        };
371        assert_eq!(variable.array_bounds, vec![Some(2)]);
372        match variable.initializer.as_ref().unwrap() {
373            Initializer::List { values, .. } => assert_eq!(values.len(), 2),
374            other => panic!("expected list initializer, got {:?}", other),
375        }
376
377        let variable = match &module.items[1] {
378            TopLevelItem::Variable(variable) => variable,
379            other => panic!("expected variable, got {:?}", other),
380        };
381        assert_eq!(variable.array_bounds, vec![None, Some(2)]);
382
383        let variable = match &module.items[2] {
384            TopLevelItem::Variable(variable) => variable,
385            other => panic!("expected variable, got {:?}", other),
386        };
387        match variable.initializer.as_ref().unwrap() {
388            Initializer::Binary {
389                op: InitializerBinaryOp::Add,
390                left,
391                right,
392                ..
393            } => {
394                assert!(
395                    matches!(left.as_ref(), Initializer::Generic { expr: inner, .. } if matches!(inner.as_ref(), Initializer::Symbol { name: symbol, .. } if symbol == "g_data"))
396                );
397                assert!(matches!(
398                    right.as_ref(),
399                    Initializer::Integer { value: 8, .. }
400                ));
401            }
402            other => panic!("expected generic add initializer, got {:?}", other),
403        }
404    }
405
406    #[test]
407    fn test_parse_function_body_variables_blocks_and_directives() {
408        let src = r#"
409            .version 8.5
410            .target sm_80
411            .address_size 64
412
413            .visible .func helper(.param .b32 helper_param)
414            {
415                .reg .b32 %r0;
416                mov.u32 %r0, %r0;
417                ret;
418            }
419
420            .visible .entry step64_kernel(
421                .param .u64 param0,
422                .param .u32 param1
423            )
424            .maxnreg 32
425            .reqntid 256, 1, 1
426            {
427                .reg .pred %p1;
428                .reg .b32 %r1;
429                .loc 1 42 3
430                mov.u32 %r1, %tid.x;
431                .pragma "nounroll";
432            $L_loop:
433                entry_br: .branchtargets $L_loop;
434                entry_call: .calltargets helper;
435                entry_proto: .callprototype _ (.param .u32 param_placeholder) .abi_preserve 4 .abi_preserve_control 2;
436                {
437                    .reg .b32 %r3;
438                    add.s32 %r3, %r1, -1;
439                }
440                ret;
441            }
442        "#;
443
444        let module = parse_module(src).unwrap();
445        assert_eq!(module.items.len(), 2);
446
447        let entry = match &module.items[1] {
448            TopLevelItem::Function(func) => func,
449            other => panic!("expected function, got {:?}", other),
450        };
451        assert!(matches!(
452            entry.directives[0],
453            FunctionDirective::MaxNReg { value: 32, .. }
454        ));
455        assert!(
456            matches!(entry.directives[1], FunctionDirective::ReqNTid { ref values, .. } if values == &vec![256, 1, 1])
457        );
458
459        let body = entry.body.as_ref().unwrap();
460        assert!(matches!(body[0], Statement::Variable(_)));
461        assert!(matches!(body[1], Statement::Variable(_)));
462        assert!(matches!(
463            body[2],
464            Statement::Directive(DirectiveStatement::Loc(_))
465        ));
466        assert!(matches!(
467            body[4],
468            Statement::Directive(DirectiveStatement::Pragma { .. })
469        ));
470        assert!(matches!(&body[5], Statement::Label { name, .. } if name == "$L_loop"));
471        assert!(matches!(&body[6], Statement::Label { name, .. } if name == "entry_br"));
472        assert!(matches!(
473            body[7],
474            Statement::Directive(DirectiveStatement::BranchTargets { .. })
475        ));
476        assert!(matches!(&body[8], Statement::Label { name, .. } if name == "entry_call"));
477        assert!(matches!(
478            body[9],
479            Statement::Directive(DirectiveStatement::CallTargets { .. })
480        ));
481        assert!(matches!(&body[10], Statement::Label { name, .. } if name == "entry_proto"));
482        assert!(matches!(
483            body[11],
484            Statement::Directive(DirectiveStatement::CallPrototype { .. })
485        ));
486        assert!(matches!(body[12], Statement::Block { .. }));
487    }
488
489    #[test]
490    fn test_parse_negative_address_offset() {
491        let src = r#"
492            .version 9.1
493            .target sm_90
494            .address_size 64
495
496            .entry kernel(.param .u64 ptr) {
497                ld.global.u32 %r1, [%rd1 - 8];
498                ret;
499            }
500        "#;
501
502        let module = parse_module(src).unwrap();
503        let func = match &module.items[0] {
504            TopLevelItem::Function(func) => func,
505            other => panic!("expected function, got {:?}", other),
506        };
507        let body = func.body.as_ref().unwrap();
508
509        match &body[0] {
510            Statement::Instruction(inst) => match &inst.operands[1] {
511                Operand::Address {
512                    address:
513                        Address::RegisterOffset {
514                            register, offset, ..
515                        },
516                    ..
517                } => {
518                    assert_eq!(register, "%rd1");
519                    assert_eq!(*offset, -8);
520                }
521                other => panic!("expected register offset address, got {:?}", other),
522            },
523            other => panic!("expected instruction, got {:?}", other),
524        }
525    }
526
527    #[test]
528    fn test_parse_binary_and_octal_literals() {
529        let src = r#"
530            .version 9.1
531            .target sm_90
532            .address_size 64
533
534            .entry kernel() {
535                mov.u32 %r1, 0b1010;
536                mov.u32 %r2, 0o17;
537                ret;
538            }
539        "#;
540
541        let module = parse_module(src).unwrap();
542        let func = match &module.items[0] {
543            TopLevelItem::Function(func) => func,
544            other => panic!("expected function, got {:?}", other),
545        };
546        let body = func.body.as_ref().unwrap();
547
548        match &body[0] {
549            Statement::Instruction(inst) => {
550                assert!(matches!(inst.operands[1], Operand::Immediate { .. }));
551            }
552            other => panic!("expected instruction, got {:?}", other),
553        }
554        match &body[1] {
555            Statement::Instruction(inst) => {
556                assert!(matches!(inst.operands[1], Operand::Immediate { .. }));
557            }
558            other => panic!("expected instruction, got {:?}", other),
559        }
560    }
561
562    #[test]
563    fn test_reject_identifier_with_dot() {
564        let src = r#"
565            .version 9.1
566            .target sm_90.something
567            .address_size 64
568        "#;
569
570        assert!(parse_module(src).is_err());
571    }
572
573    #[test]
574    fn test_parse_alias_and_section_and_file_metadata() {
575        let src = r#"
576            .version 9.1
577            .target sm_90, debug
578            .address_size 64
579            .file 1 "kernel.cu", 1339013327, 64118;
580            .visible .func foo();
581            .visible .func bar();
582            .alias bar, foo;
583            .section .debug_info {
584                .b32 11430
585            }
586        "#;
587
588        let module = parse_module(src).unwrap();
589        assert!(matches!(
590            module.items[0],
591            TopLevelItem::File {
592                timestamp: Some(1339013327),
593                file_size: Some(64118),
594                ..
595            }
596        ));
597        assert!(matches!(module.items[3], TopLevelItem::Alias { .. }));
598        assert!(
599            matches!(module.items[4], TopLevelItem::Section(ref section) if section.name == ".debug_info")
600        );
601    }
602
603    #[test]
604    fn test_parse_file_without_semicolon() {
605        let src = r#"
606            .version 9.1
607            .target sm_90
608            .file 1 "kernel.cu"
609        "#;
610
611        let module = parse_module(src).unwrap();
612        assert!(matches!(
613            module.items[0],
614            TopLevelItem::File {
615                index: 1,
616                timestamp: Some(0),
617                file_size: Some(0),
618                ..
619            }
620        ));
621    }
622
623    #[test]
624    fn test_parse_entry_scoped_pragma_before_body() {
625        let src = r#"
626            .version 9.1
627            .target sm_90
628            .address_size 64
629
630            .entry kernel() .pragma "nounroll"; {
631                ret;
632            }
633        "#;
634
635        let module = parse_module(src).unwrap();
636        let func = match &module.items[0] {
637            TopLevelItem::Function(func) => func,
638            other => panic!("expected function, got {:?}", other),
639        };
640
641        assert!(matches!(
642            func.directives[0],
643            FunctionDirective::Pragma { ref value, .. } if value == "nounroll"
644        ));
645        assert!(func.body.is_some());
646    }
647
648    #[test]
649    fn test_parse_leading_dot_float_initializer() {
650        let src = r#"
651            .version 9.1
652            .target sm_90
653            .address_size 64
654
655            .global .f32 blur_kernel[][3] = {{.05, .1, .05}, {.1, .4, .1}, {.05, .1, .05}};
656        "#;
657
658        let module = parse_module(src).unwrap();
659        let variable = match &module.items[0] {
660            TopLevelItem::Variable(variable) => variable,
661            other => panic!("expected variable, got {:?}", other),
662        };
663
664        assert_eq!(variable.array_bounds, vec![None, Some(3)]);
665        assert!(matches!(
666            variable.initializer,
667            Some(Initializer::List { .. })
668        ));
669    }
670
671    #[test]
672    fn test_parse_initializer_expression_precedence_and_casts() {
673        let src = r#"
674            .version 9.1
675            .target sm_90
676            .address_size 64
677
678            .global .u64 expr0 = 1 + 2 * 3;
679            .global .u64 expr1 = (1 + 2) * 3;
680            .global .u64 expr2 = 1 ? 2 : 3;
681            .global .u64 expr3 = (.u64) (1 + 2);
682            .global .u64 expr4 = ~1 & 3;
683        "#;
684
685        let module = parse_module(src).unwrap();
686
687        let expr0 = match &module.items[0] {
688            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
689            other => panic!("expected variable, got {:?}", other),
690        };
691        assert!(matches!(
692            expr0,
693            Initializer::Binary {
694                op: InitializerBinaryOp::Add,
695                ..
696            }
697        ));
698
699        let expr1 = match &module.items[1] {
700            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
701            other => panic!("expected variable, got {:?}", other),
702        };
703        assert!(matches!(
704            expr1,
705            Initializer::Binary {
706                op: InitializerBinaryOp::Mul,
707                ..
708            }
709        ));
710
711        let expr2 = match &module.items[2] {
712            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
713            other => panic!("expected variable, got {:?}", other),
714        };
715        assert!(matches!(expr2, Initializer::Conditional { .. }));
716
717        let expr3 = match &module.items[3] {
718            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
719            other => panic!("expected variable, got {:?}", other),
720        };
721        assert!(matches!(
722            expr3,
723            Initializer::Cast {
724                ty: InitializerCastType::U64,
725                ..
726            }
727        ));
728
729        let expr4 = match &module.items[4] {
730            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
731            other => panic!("expected variable, got {:?}", other),
732        };
733        assert!(matches!(
734            expr4,
735            Initializer::Binary {
736                op: InitializerBinaryOp::BitAnd,
737                left,
738                ..
739            } if matches!(left.as_ref(), Initializer::Unary { op: InitializerUnaryOp::BitwiseNot, .. })
740        ));
741    }
742
743    #[test]
744    fn test_parse_mask_and_generic_expression_initializers() {
745        let src = r#"
746            .version 9.1
747            .target sm_90
748            .address_size 64
749
750            .const .u32 foo = 42;
751            .global .u32 ptr = generic(foo) + 8;
752            .global .u8 addr[] = {0xff(foo + 4), 0xff(generic(foo) + 4), 0xff(1000 + 546)};
753        "#;
754
755        let module = parse_module(src).unwrap();
756
757        let ptr = match &module.items[1] {
758            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
759            other => panic!("expected variable, got {:?}", other),
760        };
761        assert!(matches!(
762            ptr,
763            Initializer::Binary {
764                op: InitializerBinaryOp::Add,
765                left,
766                right,
767                ..
768            } if matches!(left.as_ref(), Initializer::Generic { .. }) && matches!(right.as_ref(), Initializer::Integer { value: 8, .. })
769        ));
770
771        let addr = match &module.items[2] {
772            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
773            other => panic!("expected variable, got {:?}", other),
774        };
775        assert!(
776            matches!(addr, Initializer::List { values, .. } if matches!(&values[0], Initializer::Masked { .. }))
777        );
778    }
779
780    #[test]
781    fn test_parse_setp_dual_destination_and_negated_predicate_input() {
782        let src = r#"
783            .version 9.1
784            .target sm_90
785            .address_size 64
786
787            .entry kernel() {
788                setp.lt.and.s32 %p1|%p2, %r1, %r2, !%p3;
789                ret;
790            }
791        "#;
792
793        let module = parse_module(src).unwrap();
794        let func = match &module.items[0] {
795            TopLevelItem::Function(func) => func,
796            other => panic!("expected function, got {:?}", other),
797        };
798        let body = func.body.as_ref().unwrap();
799
800        match &body[0] {
801            Statement::Instruction(inst) => {
802                assert_eq!(inst.opcode, "setp");
803                assert_eq!(inst.operands.len(), 4);
804                let details = inst.setp_details.as_ref().unwrap();
805                assert!(matches!(
806                    details.secondary_destination,
807                    Some(Operand::Register { ref name, .. }) if name == "%p2"
808                ));
809                assert!(details.predicate_input_negated);
810                assert!(
811                    matches!(inst.operands[0], Operand::Register { ref name, .. } if name == "%p1")
812                );
813                assert!(
814                    matches!(inst.operands[3], Operand::Register { ref name, .. } if name == "%p3")
815                );
816            }
817            other => panic!("expected instruction, got {:?}", other),
818        }
819    }
820
821    #[test]
822    fn test_parse_setp_with_bit_bucket_destination() {
823        let src = r#"
824            .version 9.1
825            .target sm_90
826            .address_size 64
827
828            .entry kernel() {
829                setp.eq.s32 _|%p2, %r1, %r2;
830                ret;
831            }
832        "#;
833
834        let module = parse_module(src).unwrap();
835        let func = match &module.items[0] {
836            TopLevelItem::Function(func) => func,
837            other => panic!("expected function, got {:?}", other),
838        };
839        let body = func.body.as_ref().unwrap();
840
841        match &body[0] {
842            Statement::Instruction(inst) => {
843                assert!(matches!(inst.operands[0], Operand::BitBucket { .. }));
844                assert!(matches!(
845                    inst.setp_details.as_ref().unwrap().secondary_destination,
846                    Some(Operand::Register { ref name, .. }) if name == "%p2"
847                ));
848            }
849            other => panic!("expected instruction, got {:?}", other),
850        }
851    }
852
853    #[test]
854    fn test_parse_variable_and_function_attributes() {
855        let src = r#"
856            .version 9.1
857            .target sm_90
858            .address_size 64
859
860            .global .attribute(.managed) .s32 g;
861            .global .attribute(.unified(19, 95)) .f32 f;
862            .func .attribute(.unified(0xAB, 0xCD)) bar() {
863                ret;
864            }
865        "#;
866
867        let module = parse_module(src).unwrap();
868
869        let g = match &module.items[0] {
870            TopLevelItem::Variable(variable) => variable,
871            other => panic!("expected variable, got {:?}", other),
872        };
873        assert!(matches!(
874            g.attributes.as_slice(),
875            [Attribute::Managed { .. }]
876        ));
877
878        let f = match &module.items[1] {
879            TopLevelItem::Variable(variable) => variable,
880            other => panic!("expected variable, got {:?}", other),
881        };
882        assert!(matches!(
883            f.attributes.as_slice(),
884            [Attribute::Unified {
885                uuid1: 19,
886                uuid2: 95,
887                ..
888            }]
889        ));
890
891        let bar = match &module.items[2] {
892            TopLevelItem::Function(function) => function,
893            other => panic!("expected function, got {:?}", other),
894        };
895        assert!(matches!(
896            bar.attributes.as_slice(),
897            [Attribute::Unified {
898                uuid1: 0xAB,
899                uuid2: 0xCD,
900                ..
901            }]
902        ));
903    }
904
905    #[test]
906    fn test_parse_structured_loc_directive() {
907        let src = r#"
908            .version 9.1
909            .target sm_90
910            .address_size 64
911
912            .entry kernel() {
913                .loc 1 15 3, function_name .debug_str+16, inlined_at 1 10 5
914                ret;
915            }
916        "#;
917
918        let module = parse_module(src).unwrap();
919        let func = match &module.items[0] {
920            TopLevelItem::Function(function) => function,
921            other => panic!("expected function, got {:?}", other),
922        };
923
924        match &func.body.as_ref().unwrap()[0] {
925            Statement::Directive(DirectiveStatement::Loc(loc)) => {
926                assert_eq!(loc.file, 1);
927                assert_eq!(loc.line, 15);
928                assert_eq!(loc.column, 3);
929                assert!(matches!(
930                    loc.function_name,
931                    Some(LocFunctionName::LabelOffset { ref label, offset, .. }) if label == ".debug_str" && offset == 16
932                ));
933                let inline = loc.inlined_at.as_ref().unwrap();
934                assert_eq!(inline.file, 1);
935                assert_eq!(inline.line, 10);
936                assert_eq!(inline.column, 5);
937            }
938            other => panic!("expected loc directive, got {:?}", other),
939        }
940    }
941
942    #[test]
943    fn test_reject_negative_version_components() {
944        let src = ".version -1.-2\n.target sm_90\n.address_size 64\n";
945        assert!(parse_module(src).is_err());
946    }
947
948    #[test]
949    fn test_reject_negative_file_indices() {
950        let src = ".version 9.1\n.target sm_90\n.file -1 \"kernel.cu\"\n";
951        assert!(parse_module(src).is_err());
952    }
953
954    #[test]
955    fn test_reject_unknown_instruction_modifiers() {
956        let src = r#"
957            .version 9.1
958            .target sm_90
959            .address_size 64
960
961            .entry kernel() {
962                add.s32.fake %r1, %r2, %r3;
963            }
964        "#;
965
966        assert!(parse_module(src).is_err());
967    }
968
969    #[test]
970    fn test_parse_setp_with_compare_only_modifier() {
971        let src = r#"
972            .version 9.1
973            .target sm_90
974            .address_size 64
975
976            .entry kernel() {
977                setp.ge.s32 %p1, %r1, %r2;
978                ret;
979            }
980        "#;
981
982        let module = parse_module(src).unwrap();
983        let function = match &module.items[0] {
984            TopLevelItem::Function(function) => function,
985            other => panic!("expected function, got {:?}", other),
986        };
987        let body = function.body.as_ref().unwrap();
988        match &body[0] {
989            Statement::Instruction(inst) => {
990                assert_eq!(inst.opcode, "setp");
991                assert_eq!(inst.modifiers, vec!["ge", "s32"]);
992            }
993            other => panic!("expected instruction, got {:?}", other),
994        }
995    }
996
997    #[test]
998    fn test_parse_mul_with_mode_and_type_modifiers() {
999        let src = r#"
1000            .version 9.1
1001            .target sm_90
1002            .address_size 64
1003
1004            .entry kernel() {
1005                mul.wide.s32 %rd4, %r1, 4;
1006                ret;
1007            }
1008        "#;
1009
1010        let module = parse_module(src).unwrap();
1011        let function = match &module.items[0] {
1012            TopLevelItem::Function(function) => function,
1013            other => panic!("expected function, got {:?}", other),
1014        };
1015        let body = function.body.as_ref().unwrap();
1016        match &body[0] {
1017            Statement::Instruction(inst) => {
1018                assert_eq!(inst.opcode, "mul");
1019                assert_eq!(inst.modifiers, vec!["wide", "s32"]);
1020            }
1021            other => panic!("expected instruction, got {:?}", other),
1022        }
1023    }
1024
1025    #[test]
1026    fn test_keep_directive_spans_relative_to_module_source() {
1027        let src = r#"
1028            .version 9.1
1029            .target sm_90
1030            .address_size 64
1031
1032            .file 1 "kernel.cu";
1033            .alias bar, foo;
1034            .section .debug_info {
1035                start:
1036                .b32 start+12
1037            }
1038
1039            .entry kernel() {
1040                .pragma "nounroll";
1041                .loc 1 15 3, function_name .debug_str+16, inlined_at 1 10 5
1042                ret;
1043            }
1044        "#;
1045
1046        let module = parse_module(src).unwrap();
1047
1048        let file = match &module.items[0] {
1049            TopLevelItem::File { span, .. } => span,
1050            other => panic!("expected file directive, got {:?}", other),
1051        };
1052        assert_eq!(&src[file.start..file.end], ".file 1 \"kernel.cu\";");
1053
1054        let alias = match &module.items[1] {
1055            TopLevelItem::Alias { span, .. } => span,
1056            other => panic!("expected alias directive, got {:?}", other),
1057        };
1058        assert_eq!(&src[alias.start..alias.end], ".alias bar, foo;");
1059
1060        let section = match &module.items[2] {
1061            TopLevelItem::Section(section) => section,
1062            other => panic!("expected section directive, got {:?}", other),
1063        };
1064        assert_eq!(
1065            &src[section.span.start..section.span.end],
1066            ".section .debug_info {\n                start:\n                .b32 start+12\n            }"
1067        );
1068        match &section.lines[0] {
1069            SectionLine::Label { span, name } => {
1070                assert_eq!(name, "start");
1071                assert_eq!(&src[span.start..span.end], "start:");
1072            }
1073            other => panic!("expected section label, got {:?}", other),
1074        }
1075        match &section.lines[1] {
1076            SectionLine::Data { values, .. } => match &values[0] {
1077                SectionValue::LabelOffset {
1078                    span,
1079                    label,
1080                    offset,
1081                } => {
1082                    assert_eq!(label, "start");
1083                    assert_eq!(*offset, 12);
1084                    assert_eq!(&src[span.start..span.end], "start+12");
1085                }
1086                other => panic!("expected label offset, got {:?}", other),
1087            },
1088            other => panic!("expected section data, got {:?}", other),
1089        }
1090
1091        let function = match &module.items[3] {
1092            TopLevelItem::Function(function) => function,
1093            other => panic!("expected function, got {:?}", other),
1094        };
1095        let body = function.body.as_ref().unwrap();
1096        match &body[0] {
1097            Statement::Directive(DirectiveStatement::Pragma { span, value }) => {
1098                assert_eq!(value, "nounroll");
1099                assert_eq!(&src[span.start..span.end], ".pragma \"nounroll\";");
1100            }
1101            other => panic!("expected pragma, got {:?}", other),
1102        }
1103        match &body[1] {
1104            Statement::Directive(DirectiveStatement::Loc(loc)) => {
1105                assert_eq!(
1106                    &src[loc.span.start..loc.span.end],
1107                    ".loc 1 15 3, function_name .debug_str+16, inlined_at 1 10 5"
1108                );
1109                match loc.function_name.as_ref().unwrap() {
1110                    LocFunctionName::LabelOffset {
1111                        span,
1112                        label,
1113                        offset,
1114                    } => {
1115                        assert_eq!(label, ".debug_str");
1116                        assert_eq!(*offset, 16);
1117                        assert_eq!(&src[span.start..span.end], ".debug_str+16");
1118                    }
1119                    other => panic!("expected function name label offset, got {:?}", other),
1120                }
1121                let inline = loc.inlined_at.as_ref().unwrap();
1122                assert_eq!(&src[inline.span.start..inline.span.end], " 1 10 5");
1123            }
1124            other => panic!("expected loc directive, got {:?}", other),
1125        }
1126    }
1127
1128    #[test]
1129    fn test_parse_nvcc_emitted_ptx_control_flow_sequence() {
1130        let src = r#"
1131//
1132// Generated by NVIDIA NVVM Compiler
1133//
1134// Compiler Build ID: UNKNOWN
1135// Cuda compilation tools, release 13.2, V13.2.51
1136// Based on NVVM 7.0.1
1137//
1138
1139.version 9.2
1140.target sm_75
1141.address_size 64
1142
1143	// .globl	scale_add
1144
1145.visible .entry scale_add(
1146	.param .u64 scale_add_param_0,
1147	.param .u64 scale_add_param_1,
1148	.param .f32 scale_add_param_2,
1149	.param .u32 scale_add_param_3
1150)
1151{
1152	.reg .pred 	%p<2>;
1153	.reg .f32 	%f<4>;
1154	.reg .b32 	%r<6>;
1155	.reg .b64 	%rd<8>;
1156
1157
1158	ld.param.u64 	%rd1, [scale_add_param_0];
1159	ld.param.u64 	%rd2, [scale_add_param_1];
1160	ld.param.f32 	%f1, [scale_add_param_2];
1161	ld.param.u32 	%r2, [scale_add_param_3];
1162	mov.u32 	%r3, %ctaid.x;
1163	mov.u32 	%r4, %ntid.x;
1164	mov.u32 	%r5, %tid.x;
1165	mad.lo.s32 	%r1, %r3, %r4, %r5;
1166	setp.ge.s32 	%p1, %r1, %r2;
1167	@%p1 bra 	$L__BB0_2;
1168
1169	cvta.to.global.u64 	%rd3, %rd1;
1170	mul.wide.s32 	%rd4, %r1, 4;
1171	add.s64 	%rd5, %rd3, %rd4;
1172	ld.global.f32 	%f2, [%rd5];
1173	fma.rn.f32 	%f3, %f2, %f1, 0f3F800000;
1174	cvta.to.global.u64 	%rd6, %rd2;
1175	add.s64 	%rd7, %rd6, %rd4;
1176	st.global.f32 	[%rd7], %f3;
1177
1178$L__BB0_2:
1179	ret;
1180
1181}
1182"#;
1183
1184        let parsed = parse_module(src);
1185        assert!(parsed.is_ok(), "{parsed:?}");
1186    }
1187
1188    #[test]
1189    fn test_parse_section_debug_lines() {
1190        let src = r#"
1191            .version 9.1
1192            .target sm_90
1193            .address_size 64
1194
1195            .section .debug_info {
1196                start:
1197                .b32 start+12
1198                .b64 start-start
1199                .b16 -5, -65535
1200                .b8 2, 0
1201            }
1202        "#;
1203
1204        let module = parse_module(src).unwrap();
1205        let section = match &module.items[0] {
1206            TopLevelItem::Section(section) => section,
1207            other => panic!("expected section, got {:?}", other),
1208        };
1209
1210        assert_eq!(section.name, ".debug_info");
1211        assert!(matches!(&section.lines[0], SectionLine::Label { name, .. } if name == "start"));
1212        assert!(matches!(
1213            section.lines[1],
1214            SectionLine::Data {
1215                width: SectionDataWidth::B32,
1216                ref values,
1217                ..
1218            } if matches!(values[0], SectionValue::LabelOffset { ref label, offset, .. } if label == "start" && offset == 12)
1219        ));
1220        assert!(matches!(
1221            section.lines[2],
1222            SectionLine::Data {
1223                width: SectionDataWidth::B64,
1224                ref values,
1225                ..
1226            } if matches!(values[0], SectionValue::LabelDifference { ref left, ref right, .. } if left == "start" && right == "start")
1227        ));
1228        assert!(matches!(
1229            section.lines[3],
1230            SectionLine::Data {
1231                width: SectionDataWidth::B16,
1232                ref values,
1233                ..
1234            } if matches!(values[0], SectionValue::Integer { value: -5, .. })
1235                && matches!(values[1], SectionValue::Integer { value: -65535, .. })
1236        ));
1237    }
1238
1239    #[test]
1240    fn test_track_spans_for_structural_nodes() {
1241        let src = ".version 9.1\n.target sm_90\n.address_size 64\n.entry kernel() {\nlabel0:\n  ret;\n}\n";
1242
1243        let module = parse_module(src).unwrap();
1244        assert_eq!(module.span.start, 0);
1245        assert_eq!(module.span.end, src.len());
1246        assert_eq!(
1247            &src[module.version.span.start..module.version.span.end],
1248            ".version 9.1"
1249        );
1250        assert_eq!(
1251            &src[module.target.span.start..module.target.span.end],
1252            ".target sm_90"
1253        );
1254
1255        let func = match &module.items[0] {
1256            TopLevelItem::Function(function) => function,
1257            other => panic!("expected function, got {:?}", other),
1258        };
1259        assert_eq!(
1260            &src[func.span.start..func.span.end],
1261            ".entry kernel() {\nlabel0:\n  ret;\n}"
1262        );
1263
1264        let body = func.body.as_ref().unwrap();
1265        match &body[0] {
1266            Statement::Label { span, name } => {
1267                assert_eq!(name, "label0");
1268                assert_eq!(&src[span.start..span.end], "label0:");
1269            }
1270            other => panic!("expected label, got {:?}", other),
1271        }
1272
1273        match &body[1] {
1274            Statement::Instruction(inst) => {
1275                assert_eq!(&src[inst.span.start..inst.span.end], "ret;");
1276            }
1277            other => panic!("expected instruction, got {:?}", other),
1278        }
1279
1280        let start = func.span.start_position(src);
1281        let end = func.span.end_position(src);
1282        assert_eq!((start.line, start.column), (4, 1));
1283        assert_eq!((end.line, end.column), (7, 2));
1284    }
1285
1286    #[test]
1287    fn test_track_spans_for_operand_leaves() {
1288        let src = ".version 9.1\n.target sm_90\n.address_size 64\n.entry kernel() {\n  @!%p0 add.s32 %r1, [%rd4 + 16], 0b1010;\n}\n";
1289
1290        let module = parse_module(src).unwrap();
1291        let func = match &module.items[0] {
1292            TopLevelItem::Function(function) => function,
1293            other => panic!("expected function, got {:?}", other),
1294        };
1295
1296        let inst = match &func.body.as_ref().unwrap()[0] {
1297            Statement::Instruction(inst) => inst,
1298            other => panic!("expected instruction, got {:?}", other),
1299        };
1300
1301        let guard = inst.guard.as_ref().unwrap();
1302        assert_eq!(&src[guard.span.start..guard.span.end], "@!%p0");
1303
1304        match &inst.operands[0] {
1305            Operand::Register { span, name } => {
1306                assert_eq!(name, "%r1");
1307                assert_eq!(&src[span.start..span.end], "%r1");
1308            }
1309            other => panic!("expected register, got {:?}", other),
1310        }
1311
1312        match &inst.operands[1] {
1313            Operand::Address { span, address } => {
1314                assert_eq!(&src[span.start..span.end], "[%rd4 + 16]");
1315                match address {
1316                    Address::RegisterOffset {
1317                        span,
1318                        register,
1319                        offset,
1320                    } => {
1321                        assert_eq!(register, "%rd4");
1322                        assert_eq!(*offset, 16);
1323                        assert_eq!(&src[span.start..span.end], "[%rd4 + 16]");
1324                    }
1325                    other => panic!("expected register offset address, got {:?}", other),
1326                }
1327            }
1328            other => panic!("expected address operand, got {:?}", other),
1329        }
1330
1331        match &inst.operands[2] {
1332            Operand::Immediate { span, immediate } => {
1333                assert_eq!(&src[span.start..span.end], "0b1010");
1334                match immediate {
1335                    crate::ast::Immediate::UnsignedInteger { span, value } => {
1336                        assert_eq!(*value, 10);
1337                        assert_eq!(&src[span.start..span.end], "0b1010");
1338                    }
1339                    other => panic!("expected unsigned immediate, got {:?}", other),
1340                }
1341            }
1342            other => panic!("expected immediate operand, got {:?}", other),
1343        }
1344    }
1345
1346    #[test]
1347    fn test_track_spans_for_initializer_leaves() {
1348        let src = ".version 9.1\n.target sm_90\n.address_size 64\n.const .u32 foo = 42;\n.global .u64 expr = generic(foo) + 8;\n.global .u8 addr[] = {0xff(foo + 4)};\n";
1349
1350        let module = parse_module(src).unwrap();
1351
1352        let expr = match &module.items[1] {
1353            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
1354            other => panic!("expected variable, got {:?}", other),
1355        };
1356
1357        match expr {
1358            Initializer::Binary {
1359                span,
1360                op: InitializerBinaryOp::Add,
1361                left,
1362                right,
1363            } => {
1364                assert_eq!(&src[span.start..span.end], "generic(foo) + 8");
1365                match left.as_ref() {
1366                    Initializer::Generic { span, expr } => {
1367                        assert_eq!(&src[span.start..span.end], "generic(foo)");
1368                        match expr.as_ref() {
1369                            Initializer::Symbol { span, name } => {
1370                                assert_eq!(name, "foo");
1371                                assert_eq!(&src[span.start..span.end], "foo");
1372                            }
1373                            other => panic!("expected symbol initializer, got {:?}", other),
1374                        }
1375                    }
1376                    other => panic!("expected generic initializer, got {:?}", other),
1377                }
1378                match right.as_ref() {
1379                    Initializer::Integer { span, value } => {
1380                        assert_eq!(*value, 8);
1381                        assert_eq!(&src[span.start..span.end], "8");
1382                    }
1383                    other => panic!("expected integer initializer, got {:?}", other),
1384                }
1385            }
1386            other => panic!("expected binary initializer, got {:?}", other),
1387        }
1388
1389        let addr = match &module.items[2] {
1390            TopLevelItem::Variable(variable) => variable.initializer.as_ref().unwrap(),
1391            other => panic!("expected variable, got {:?}", other),
1392        };
1393
1394        match addr {
1395            Initializer::List { values, .. } => match &values[0] {
1396                Initializer::Masked { span, mask, value } => {
1397                    assert_eq!(*mask, 0xff);
1398                    assert_eq!(&src[span.start..span.end], "0xff(foo + 4)");
1399                    match value.as_ref() {
1400                        Initializer::Binary {
1401                            span,
1402                            op: InitializerBinaryOp::Add,
1403                            left,
1404                            right,
1405                        } => {
1406                            assert_eq!(&src[span.start..span.end], "foo + 4");
1407                            assert!(
1408                                matches!(left.as_ref(), Initializer::Symbol { name, .. } if name == "foo")
1409                            );
1410                            assert!(matches!(
1411                                right.as_ref(),
1412                                Initializer::Integer { value: 4, .. }
1413                            ));
1414                        }
1415                        other => panic!("expected binary initializer, got {:?}", other),
1416                    }
1417                }
1418                other => panic!("expected masked initializer, got {:?}", other),
1419            },
1420            other => panic!("expected list initializer, got {:?}", other),
1421        }
1422    }
1423}