ptx_parser/parser/
common.rs

1use crate::{
2    alt, c, err, mapc, ok,
3    parser::{
4        ParseErrorKind, PtxParseError, PtxParser, PtxTokenStream, Span,
5        util::{
6            alt, at_p, between, comma_p, directive_exact_p, directive_p, exclamation_p,
7            identifier_p, lbrace_p, lbracket_p, literal_p, lparen_p, map, minus_p, optional,
8            parse_index_suffix, plus_p, rbrace_p, rbracket_p, register_p, rparen_p, sep_by1, seq,
9            skip_first, try_map, u64_p,
10        },
11    },
12    seq_n, span,
13    span::Spanned,
14    r#type::{
15        AddressBase, AddressOffset, AddressOperand, AttributeDirective, Axis, CodeLinkage,
16        DataLinkage, DataType, FunctionSymbol, GeneralOperand, Immediate, Instruction, Label,
17        Operand, ParamStateSpace, Predicate, PredicateRegister, RegisterOperand, Sign,
18        SpecialRegister, TexHandler2, TexHandler3, TexHandler3Optional, VariableSymbol,
19        VectorOperand,
20    },
21};
22
23const CODE_LINKAGE_EXPECTED: [&str; 3] = [".visible", ".extern", ".weak"];
24
25impl PtxParser for CodeLinkage {
26    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
27        try_map(directive_p(), |name, span| {
28            let node = match name.as_str() {
29                "visible" => c!(CodeLinkage::Visible),
30                "extern" => c!(CodeLinkage::Extern),
31                "weak" => c!(CodeLinkage::Weak),
32                _ => {
33                    return err!(ParseErrorKind::UnexpectedToken {
34                        expected: CODE_LINKAGE_EXPECTED
35                            .iter()
36                            .map(|s| s.to_string())
37                            .collect(),
38                        found: format!(".{name}"),
39                    });
40                }
41            };
42            Ok(node)
43        })
44    }
45}
46
47const DATA_LINKAGE_EXPECTED: [&str; 4] = [".visible", ".extern", ".weak", ".common"];
48
49impl PtxParser for DataLinkage {
50    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
51        try_map(directive_p(), |name, span| {
52            let node = match name.as_str() {
53                "visible" => c!(DataLinkage::Visible),
54                "extern" => c!(DataLinkage::Extern),
55                "weak" => c!(DataLinkage::Weak),
56                "common" => c!(DataLinkage::Common),
57                _ => {
58                    return err!(ParseErrorKind::UnexpectedToken {
59                        expected: DATA_LINKAGE_EXPECTED
60                            .iter()
61                            .map(|s| s.to_string())
62                            .collect(),
63                        found: format!(".{name}"),
64                    });
65                }
66            };
67            Ok(node)
68        })
69    }
70}
71
72const DATA_TYPE_EXPECTED: [&str; 21] = [
73    ".u8",
74    ".u16",
75    ".u32",
76    ".u64",
77    ".s8",
78    ".s16",
79    ".s32",
80    ".s64",
81    ".f16",
82    ".f16x2",
83    ".f32",
84    ".f64",
85    ".b8",
86    ".b16",
87    ".b32",
88    ".b64",
89    ".b128",
90    ".pred",
91    ".texref",
92    ".samplerref",
93    ".surfref",
94];
95
96impl PtxParser for DataType {
97    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
98        try_map(directive_p(), |name, span| {
99            let node = match name.as_str() {
100                "u8" => c!(DataType::U8),
101                "u16" => c!(DataType::U16),
102                "u32" => c!(DataType::U32),
103                "u64" => c!(DataType::U64),
104                "s8" => c!(DataType::S8),
105                "s16" => c!(DataType::S16),
106                "s32" => c!(DataType::S32),
107                "s64" => c!(DataType::S64),
108                "f16" => c!(DataType::F16),
109                "f16x2" => c!(DataType::F16x2),
110                "f32" => c!(DataType::F32),
111                "f64" => c!(DataType::F64),
112                "b8" => c!(DataType::B8),
113                "b16" => c!(DataType::B16),
114                "b32" => c!(DataType::B32),
115                "b64" => c!(DataType::B64),
116                "b128" => c!(DataType::B128),
117                "pred" => c!(DataType::Pred),
118                "texref" => c!(DataType::TexRef),
119                "samplerref" => c!(DataType::SamplerRef),
120                "surfref" => c!(DataType::SurfRef),
121                _ => {
122                    return err!(ParseErrorKind::UnexpectedToken {
123                        expected: DATA_TYPE_EXPECTED.iter().map(|s| s.to_string()).collect(),
124                        found: format!(".{name}"),
125                    });
126                }
127            };
128            Ok(node)
129        })
130    }
131}
132
133impl PtxParser for AttributeDirective {
134    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
135        alt(
136            map(directive_exact_p("managed"), |_, span| {
137                c!(AttributeDirective::Managed)
138            }),
139            mapc!(
140                skip_first(
141                    directive_exact_p("unified"),
142                    between(
143                        lparen_p(),
144                        rparen_p(),
145                        seq_n!(u64_p(), skip_first(comma_p(), u64_p()))
146                    ),
147                ),
148                AttributeDirective::Unified { uuid1, uuid2 }
149            ),
150        )
151    }
152}
153
154impl PtxParser for Sign {
155    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
156        alt(
157            map(plus_p(), |_, span| c!(Sign::Positive)),
158            map(minus_p(), |_, span| c!(Sign::Negative)),
159        )
160    }
161}
162
163impl PtxParser for RegisterOperand {
164    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
165        try_map(register_p(), |mut raw, span| {
166            let component = if let Some(idx) = raw.rfind('.') {
167                if idx + 1 >= raw.len() {
168                    return err!(ParseErrorKind::InvalidLiteral(
169                        "register component missing after '.'".into(),
170                    ));
171                }
172                let suffix = raw[idx + 1..].to_string();
173                raw.truncate(idx);
174                Some(suffix)
175            } else {
176                None
177            };
178            let name = raw;
179            ok!(RegisterOperand { name, component })
180        })
181    }
182}
183
184impl PtxParser for VariableSymbol {
185    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
186        map(identifier_p(), |val, span| c!(VariableSymbol { val }))
187    }
188}
189
190impl PtxParser for FunctionSymbol {
191    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
192        map(identifier_p(), |val, span| c!(FunctionSymbol { val }))
193    }
194}
195
196impl PtxParser for Label {
197    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
198        map(identifier_p(), |val, span| c!(Label { val }))
199    }
200}
201
202impl PtxParser for PredicateRegister {
203    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
204        try_map(register_p(), |name, span| {
205            if name.starts_with("%p") {
206                ok!(PredicateRegister { name })
207            } else {
208                err!(ParseErrorKind::InvalidLiteral(
209                    "expected predicate register (%pX)".into(),
210                ))
211            }
212        })
213    }
214}
215
216impl PtxParser for Predicate {
217    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
218        try_map(
219            seq_n!(at_p(), optional(exclamation_p()), Operand::parse()),
220            |(_, negation, operand), span| {
221                let negated = negation.is_some();
222                ok!(Predicate { negated, operand })
223            },
224        )
225    }
226}
227
228impl PtxParser for Instruction {
229    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
230        use crate::r#type::instruction::Inst;
231
232        try_map(
233            seq(optional(Predicate::parse()), Inst::parse()),
234            |(predicate, inst), span| ok!(Instruction { predicate, inst }),
235        )
236    }
237}
238
239impl PtxParser for ParamStateSpace {
240    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
241        alt!(
242            map(directive_exact_p("const"), |_, span| {
243                c!(ParamStateSpace::Const)
244            }),
245            map(directive_exact_p("global"), |_, span| {
246                c!(ParamStateSpace::Global)
247            }),
248            map(directive_exact_p("local"), |_, span| {
249                c!(ParamStateSpace::Local)
250            }),
251            map(directive_exact_p("shared"), |_, span| {
252                c!(ParamStateSpace::Shared)
253            }),
254        )
255    }
256}
257
258impl PtxParser for Operand {
259    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
260        let register = mapc!(RegisterOperand::parse(), Operand::Register { operand });
261        let immediate = map(
262            seq(optional(Sign::parse()), Immediate::parse()),
263            |(sign, mut operand), span| {
264                if matches!(sign, Some(Sign::Negative { .. })) && !operand.value.starts_with('-') {
265                    operand.value = format!("-{}", operand.value);
266                }
267                c!(Operand::Immediate { operand })
268            },
269        );
270        let symbol_offset = map(
271            seq_n!(identifier_p(), plus_p(), Immediate::parse()),
272            |(symbol, _, offset), span| c!(Operand::SymbolOffset { symbol, offset }),
273        );
274        let symbol = mapc!(identifier_p(), Operand::Symbol { name });
275        alt!(register, immediate, symbol_offset, symbol)
276    }
277}
278
279impl PtxParser for VectorOperand {
280    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
281        try_map(
282            between(lbrace_p(), rbrace_p(), sep_by1(Operand::parse(), comma_p())),
283            |operands, span| match operands.len() {
284                1 => {
285                    let mut iter = operands.into_iter();
286                    let operand = iter.next().unwrap();
287                    ok!(VectorOperand::Vector1 { operand })
288                }
289                2 => {
290                    let mut iter = operands.into_iter();
291                    let operands = [iter.next().unwrap(), iter.next().unwrap()];
292                    ok!(VectorOperand::Vector2 { operands })
293                }
294                3 => {
295                    let mut iter = operands.into_iter();
296                    let operands = [
297                        iter.next().unwrap(),
298                        iter.next().unwrap(),
299                        iter.next().unwrap(),
300                    ];
301                    ok!(VectorOperand::Vector3 { operands })
302                }
303                4 => {
304                    let mut iter = operands.into_iter();
305                    let operands = [
306                        iter.next().unwrap(),
307                        iter.next().unwrap(),
308                        iter.next().unwrap(),
309                        iter.next().unwrap(),
310                    ];
311                    ok!(VectorOperand::Vector4 { operands })
312                }
313                8 => {
314                    let mut iter = operands.into_iter();
315                    let operands = [
316                        iter.next().unwrap(),
317                        iter.next().unwrap(),
318                        iter.next().unwrap(),
319                        iter.next().unwrap(),
320                        iter.next().unwrap(),
321                        iter.next().unwrap(),
322                        iter.next().unwrap(),
323                        iter.next().unwrap(),
324                    ];
325                    ok!(VectorOperand::Vector8 { operands })
326                }
327                _ => {
328                    let span = operands.first().map(Spanned::span).unwrap_or(span!(0..0));
329                    Err(PtxParseError {
330                        kind: ParseErrorKind::UnexpectedToken {
331                            expected: vec!["vector with 1,2,3,4, or 8 operands".into()],
332                            found: format!("vector with {} operands", operands.len()),
333                        },
334                        span,
335                    })
336                }
337            },
338        )
339    }
340}
341
342impl PtxParser for GeneralOperand {
343    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
344        alt(
345            mapc!(VectorOperand::parse(), GeneralOperand::Vec { operand }),
346            mapc!(Operand::parse(), GeneralOperand::Single { operand }),
347        )
348    }
349}
350
351fn handler_len_error(ty: &str, expected: &str, found: usize) -> PtxParseError {
352    PtxParseError {
353        kind: ParseErrorKind::InvalidLiteral(format!("{ty} expects {expected}, found {found}")),
354        span: Span::new(0, 0),
355    }
356}
357
358fn tex_operands()
359-> impl Fn(&mut PtxTokenStream) -> Result<(Vec<GeneralOperand>, Span), PtxParseError> {
360    between(
361        lbracket_p(),
362        rbracket_p(),
363        sep_by1(GeneralOperand::parse(), comma_p()),
364    )
365}
366
367impl PtxParser for TexHandler2 {
368    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
369        try_map(tex_operands(), |operands, span| {
370            if operands.len() != 2 {
371                return Err(handler_len_error(
372                    "TexHandler2",
373                    "exactly 2 operands",
374                    operands.len(),
375                ));
376            }
377            let mut iter = operands.into_iter();
378            let first = iter.next().unwrap();
379            let second = iter.next().unwrap();
380            let operands = [first, second];
381            ok!(TexHandler2 { operands })
382        })
383    }
384}
385
386impl PtxParser for TexHandler3 {
387    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
388        try_map(tex_operands(), |operands, span| {
389            if operands.len() != 3 {
390                return Err(handler_len_error(
391                    "TexHandler3",
392                    "exactly 3 operands",
393                    operands.len(),
394                ));
395            }
396            let mut iter = operands.into_iter();
397            let handle = iter.next().unwrap();
398            let sampler = iter.next().unwrap();
399            let coords = iter.next().unwrap();
400            ok!(TexHandler3 {
401                handle,
402                sampler,
403                coords
404            })
405        })
406    }
407}
408
409impl PtxParser for TexHandler3Optional {
410    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
411        try_map(tex_operands(), |operands, span| match operands.len() {
412            2 => {
413                let mut iter = operands.into_iter();
414                let handle = iter.next().unwrap();
415                let coords = iter.next().unwrap();
416                let sampler = None;
417                ok!(TexHandler3Optional {
418                    handle,
419                    sampler,
420                    coords
421                })
422            }
423            3 => {
424                let mut iter = operands.into_iter();
425                let handle = iter.next().unwrap();
426                let sampler = Some(iter.next().unwrap());
427                let coords = iter.next().unwrap();
428                ok!(TexHandler3Optional {
429                    handle,
430                    sampler,
431                    coords
432                })
433            }
434            other => Err(handler_len_error(
435                "TexHandler3Optional",
436                "2 or 3 operands",
437                other,
438            )),
439        })
440    }
441}
442
443impl PtxParser for SpecialRegister {
444    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
445        try_map(register_p(), |name, span| {
446            let raw = name.strip_prefix('%').ok_or_else(|| PtxParseError {
447                kind: ParseErrorKind::InvalidLiteral("expected % prefix".into()),
448                span,
449            })?;
450
451            let (base, axis_suffix) = split_axis_suffix(raw).map_err(|msg| PtxParseError {
452                kind: ParseErrorKind::InvalidLiteral(msg.into()),
453                span,
454            })?;
455            let base_lower = base.to_ascii_lowercase();
456
457            let mut node = match base_lower.as_str() {
458                "aggr_smem_size" => {
459                    ensure_no_axis(axis_suffix, span)?;
460                    c!(SpecialRegister::AggrSmemSize)
461                }
462                "dynamic_smem_size" => {
463                    ensure_no_axis(axis_suffix, span)?;
464                    c!(SpecialRegister::DynamicSmemSize)
465                }
466                "lanemask_gt" => {
467                    ensure_no_axis(axis_suffix, span)?;
468                    c!(SpecialRegister::LanemaskGt)
469                }
470                "reserved_smem_offset_begin" => {
471                    ensure_no_axis(axis_suffix, span)?;
472                    c!(SpecialRegister::ReservedSmemOffsetBegin)
473                }
474                "clock" => {
475                    ensure_no_axis(axis_suffix, span)?;
476                    c!(SpecialRegister::Clock)
477                }
478                "lanemask_le" => {
479                    ensure_no_axis(axis_suffix, span)?;
480                    c!(SpecialRegister::LanemaskLe)
481                }
482                "reserved_smem_offset_cap" => {
483                    ensure_no_axis(axis_suffix, span)?;
484                    c!(SpecialRegister::ReservedSmemOffsetCap)
485                }
486                "clock64" => {
487                    ensure_no_axis(axis_suffix, span)?;
488                    c!(SpecialRegister::Clock64)
489                }
490                "globaltimer" => {
491                    ensure_no_axis(axis_suffix, span)?;
492                    c!(SpecialRegister::Globaltimer)
493                }
494                "lanemask_lt" => {
495                    ensure_no_axis(axis_suffix, span)?;
496                    c!(SpecialRegister::LanemaskLt)
497                }
498                "reserved_smem_offset_end" => {
499                    ensure_no_axis(axis_suffix, span)?;
500                    c!(SpecialRegister::ReservedSmemOffsetEnd)
501                }
502                "globaltimer_hi" => {
503                    ensure_no_axis(axis_suffix, span)?;
504                    c!(SpecialRegister::GlobaltimerHi)
505                }
506                "nclusterid" => {
507                    ensure_no_axis(axis_suffix, span)?;
508                    c!(SpecialRegister::Nclusterid)
509                }
510                "smid" => {
511                    ensure_no_axis(axis_suffix, span)?;
512                    c!(SpecialRegister::Smid)
513                }
514                "globaltimer_lo" => {
515                    ensure_no_axis(axis_suffix, span)?;
516                    c!(SpecialRegister::GlobaltimerLo)
517                }
518                "gridid" => {
519                    ensure_no_axis(axis_suffix, span)?;
520                    c!(SpecialRegister::Gridid)
521                }
522                "nsmid" => {
523                    ensure_no_axis(axis_suffix, span)?;
524                    c!(SpecialRegister::Nsmid)
525                }
526                "total_smem_size" => {
527                    ensure_no_axis(axis_suffix, span)?;
528                    c!(SpecialRegister::TotalSmemSize)
529                }
530                "is_explicit_cluster" => {
531                    ensure_no_axis(axis_suffix, span)?;
532                    c!(SpecialRegister::IsExplicitCluster)
533                }
534                "warpid" => {
535                    ensure_no_axis(axis_suffix, span)?;
536                    c!(SpecialRegister::Warpid)
537                }
538                "clusterid" => {
539                    ensure_no_axis(axis_suffix, span)?;
540                    c!(SpecialRegister::Clusterid)
541                }
542                "laneid" => {
543                    ensure_no_axis(axis_suffix, span)?;
544                    c!(SpecialRegister::Laneid)
545                }
546                "nwarpid" => {
547                    ensure_no_axis(axis_suffix, span)?;
548                    c!(SpecialRegister::Nwarpid)
549                }
550                "warpsz" => {
551                    ensure_no_axis(axis_suffix, span)?;
552                    c!(SpecialRegister::WARPSZ)
553                }
554                "lanemask_eq" => {
555                    ensure_no_axis(axis_suffix, span)?;
556                    c!(SpecialRegister::LanemaskEq)
557                }
558                "current_graph_exec" => {
559                    ensure_no_axis(axis_suffix, span)?;
560                    c!(SpecialRegister::CurrentGraphExec)
561                }
562                "lanemask_ge" => {
563                    ensure_no_axis(axis_suffix, span)?;
564                    c!(SpecialRegister::LanemaskGe)
565                }
566                "cluster_ctaid" => {
567                    let axis = axis_with_span(axis_suffix, &span)?;
568                    c!(SpecialRegister::ClusterCtaid { axis })
569                }
570                "cluster_ctarank" => {
571                    let axis = axis_with_span(axis_suffix, &span)?;
572                    c!(SpecialRegister::ClusterCtarank { axis })
573                }
574                "nctaid" => {
575                    let axis = axis_with_span(axis_suffix, &span)?;
576                    c!(SpecialRegister::Nctaid { axis })
577                }
578                "tid" => {
579                    let axis = axis_with_span(axis_suffix, &span)?;
580                    c!(SpecialRegister::Tid { axis })
581                }
582                "cluster_nctaid" => {
583                    let axis = axis_with_span(axis_suffix, &span)?;
584                    c!(SpecialRegister::ClusterNctaid { axis })
585                }
586                "cluster_nctarank" => {
587                    let axis = axis_with_span(axis_suffix, &span)?;
588                    c!(SpecialRegister::ClusterNctarank { axis })
589                }
590                "ntid" => {
591                    let axis = axis_with_span(axis_suffix, &span)?;
592                    c!(SpecialRegister::Ntid { axis })
593                }
594                "ctaid" => {
595                    let axis = axis_with_span(axis_suffix, &span)?;
596                    c!(SpecialRegister::Ctaid { axis })
597                }
598                base if base.starts_with("envreg") => {
599                    ensure_no_axis(axis_suffix, span)?;
600                    let digits = base.strip_prefix("envreg").unwrap();
601                    let index = parse_index_suffix(digits, 31, "envreg", span)?;
602                    c!(SpecialRegister::Envreg { index })
603                }
604                base if base.starts_with("pm") && base.ends_with("_64") => {
605                    ensure_no_axis(axis_suffix, span)?;
606                    let digits = base
607                        .strip_prefix("pm")
608                        .and_then(|rest| rest.strip_suffix("_64"))
609                        .unwrap();
610                    let index = parse_index_suffix(digits, 7, "pm64", span)?;
611                    c!(SpecialRegister::Pm64 { index })
612                }
613                base if base.starts_with("pm") => {
614                    ensure_no_axis(axis_suffix, span)?;
615                    let digits = base.strip_prefix("pm").unwrap();
616                    let index = parse_index_suffix(digits, 7, "pm", span)?;
617                    c!(SpecialRegister::Pm { index })
618                }
619                base if base.starts_with("reserved_smem_offset_") => {
620                    ensure_no_axis(axis_suffix, span)?;
621                    let digits = base.strip_prefix("reserved_smem_offset_").unwrap();
622                    let index = parse_index_suffix(digits, 1, "reserved_smem_offset", span)?;
623                    c!(SpecialRegister::ReservedSmemOffset { index })
624                }
625                _ => {
626                    return err!(ParseErrorKind::InvalidLiteral(format!(
627                        "unknown special register: {name}"
628                    )));
629                }
630            };
631
632            node.set_span(span);
633            Ok(node)
634        })
635    }
636}
637
638impl PtxParser for Immediate {
639    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
640        mapc!(literal_p(), Immediate { value })
641    }
642}
643
644impl PtxParser for AddressBase {
645    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
646        alt(
647            mapc!(RegisterOperand::parse(), AddressBase::Register { operand }),
648            mapc!(VariableSymbol::parse(), AddressBase::Variable { symbol }),
649        )
650    }
651}
652
653impl PtxParser for AddressOffset {
654    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
655        alt(
656            mapc!(
657                seq(plus_p(), RegisterOperand::parse()),
658                AddressOffset::Register { _, operand }
659            ),
660            mapc!(
661                seq(Sign::parse(), Immediate::parse()),
662                AddressOffset::Immediate { sign, value }
663            ),
664        )
665    }
666}
667
668impl PtxParser for AddressOperand {
669    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
670        alt!(
671            mapc!(
672                seq(
673                    VariableSymbol::parse(),
674                    between(lbracket_p(), rbracket_p(), Immediate::parse()),
675                ),
676                AddressOperand::Array { base, index }
677            ),
678            mapc!(
679                between(
680                    lbracket_p(),
681                    rbracket_p(),
682                    seq(AddressBase::parse(), optional(AddressOffset::parse())),
683                ),
684                AddressOperand::Offset { base, offset }
685            ),
686            map(
687                between(
688                    lbracket_p(),
689                    rbracket_p(),
690                    seq(optional(Sign::parse()), Immediate::parse()),
691                ),
692                |(sign, mut addr), span| {
693                    if matches!(sign, Some(Sign::Negative { .. })) && !addr.value.starts_with('-') {
694                        addr.value = format!("-{}", addr.value);
695                    }
696                    c!(AddressOperand::ImmediateAddress { addr })
697                },
698            ),
699        )
700    }
701}
702
703fn split_axis_suffix(raw: &str) -> Result<(&str, Option<char>), &'static str> {
704    if let Some(idx) = raw.rfind('.') {
705        let (base, suffix) = raw.split_at(idx);
706        if suffix.len() != 2 {
707            return Err("invalid axis suffix");
708        }
709        let axis_char = suffix.chars().nth(1).unwrap();
710        Ok((base, Some(axis_char)))
711    } else {
712        Ok((raw, None))
713    }
714}
715
716fn axis_from_suffix(axis: Option<char>, span: Span) -> Result<Axis, PtxParseError> {
717    match axis.map(|c| c.to_ascii_lowercase()) {
718        None => Ok(Axis::None { span }),
719        Some('x') => Ok(Axis::X { span }),
720        Some('y') => Ok(Axis::Y { span }),
721        Some('z') => Ok(Axis::Z { span }),
722        Some(other) => Err(PtxParseError {
723            kind: ParseErrorKind::InvalidLiteral(format!(
724                "invalid axis '.{}' for special register",
725                other
726            )),
727            span,
728        }),
729    }
730}
731
732fn axis_with_span(axis: Option<char>, span: &Span) -> Result<Axis, PtxParseError> {
733    axis_from_suffix(axis, *span)
734}
735
736fn ensure_no_axis(axis: Option<char>, span: Span) -> Result<(), PtxParseError> {
737    if let Some(ch) = axis {
738        Err(PtxParseError {
739            kind: ParseErrorKind::InvalidLiteral(format!(
740                "register does not accept axis '.{}'",
741                ch
742            )),
743            span,
744        })
745    } else {
746        Ok(())
747    }
748}
ptx_parser/parser/common.rs

ptx_parser/parser/
common.rs
