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        alt(
199            map(identifier_p(), |val, span| c!(Label { val })),
200            map(directive_p(), |name, span| c!(Label {
201                val = format!(".{name}")
202            })),
203        )
204    }
205}
206
207impl PtxParser for PredicateRegister {
208    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
209        try_map(register_p(), |name, span| {
210            if name.starts_with("%p") {
211                ok!(PredicateRegister { name })
212            } else {
213                err!(ParseErrorKind::InvalidLiteral(
214                    "expected predicate register (%pX)".into(),
215                ))
216            }
217        })
218    }
219}
220
221impl PtxParser for Predicate {
222    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
223        try_map(
224            seq_n!(at_p(), optional(exclamation_p()), Operand::parse()),
225            |(_, negation, operand), span| {
226                let negated = negation.is_some();
227                ok!(Predicate { negated, operand })
228            },
229        )
230    }
231}
232
233impl PtxParser for Instruction {
234    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
235        use crate::r#type::instruction::Inst;
236
237        try_map(
238            seq(optional(Predicate::parse()), Inst::parse()),
239            |(predicate, inst), span| ok!(Instruction { predicate, inst }),
240        )
241    }
242}
243
244impl PtxParser for ParamStateSpace {
245    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
246        alt!(
247            map(directive_exact_p("const"), |_, span| {
248                c!(ParamStateSpace::Const)
249            }),
250            map(directive_exact_p("global"), |_, span| {
251                c!(ParamStateSpace::Global)
252            }),
253            map(directive_exact_p("local"), |_, span| {
254                c!(ParamStateSpace::Local)
255            }),
256            map(directive_exact_p("shared"), |_, span| {
257                c!(ParamStateSpace::Shared)
258            }),
259        )
260    }
261}
262
263impl PtxParser for Operand {
264    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
265        let register = mapc!(RegisterOperand::parse(), Operand::Register { operand });
266        let immediate = map(
267            seq(optional(Sign::parse()), Immediate::parse()),
268            |(sign, mut operand), span| {
269                if matches!(sign, Some(Sign::Negative { .. })) && !operand.value.starts_with('-') {
270                    operand.value = format!("-{}", operand.value);
271                }
272                c!(Operand::Immediate { operand })
273            },
274        );
275        let symbol_offset = map(
276            seq_n!(identifier_p(), plus_p(), Immediate::parse()),
277            |(symbol, _, offset), span| c!(Operand::SymbolOffset { symbol, offset }),
278        );
279        let symbol = mapc!(identifier_p(), Operand::Symbol { name });
280        alt!(register, immediate, symbol_offset, symbol)
281    }
282}
283
284impl PtxParser for VectorOperand {
285    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
286        try_map(
287            between(lbrace_p(), rbrace_p(), sep_by1(Operand::parse(), comma_p())),
288            |operands, span| match operands.len() {
289                1 => {
290                    let mut iter = operands.into_iter();
291                    let operand = iter.next().unwrap();
292                    ok!(VectorOperand::Vector1 { operand })
293                }
294                2 => {
295                    let mut iter = operands.into_iter();
296                    let operands = [iter.next().unwrap(), iter.next().unwrap()];
297                    ok!(VectorOperand::Vector2 { operands })
298                }
299                3 => {
300                    let mut iter = operands.into_iter();
301                    let operands = [
302                        iter.next().unwrap(),
303                        iter.next().unwrap(),
304                        iter.next().unwrap(),
305                    ];
306                    ok!(VectorOperand::Vector3 { operands })
307                }
308                4 => {
309                    let mut iter = operands.into_iter();
310                    let operands = [
311                        iter.next().unwrap(),
312                        iter.next().unwrap(),
313                        iter.next().unwrap(),
314                        iter.next().unwrap(),
315                    ];
316                    ok!(VectorOperand::Vector4 { operands })
317                }
318                8 => {
319                    let mut iter = operands.into_iter();
320                    let operands = [
321                        iter.next().unwrap(),
322                        iter.next().unwrap(),
323                        iter.next().unwrap(),
324                        iter.next().unwrap(),
325                        iter.next().unwrap(),
326                        iter.next().unwrap(),
327                        iter.next().unwrap(),
328                        iter.next().unwrap(),
329                    ];
330                    ok!(VectorOperand::Vector8 { operands })
331                }
332                _ => {
333                    let span = operands.first().map(Spanned::span).unwrap_or(span!(0..0));
334                    Err(PtxParseError {
335                        kind: ParseErrorKind::UnexpectedToken {
336                            expected: vec!["vector with 1,2,3,4, or 8 operands".into()],
337                            found: format!("vector with {} operands", operands.len()),
338                        },
339                        span,
340                    })
341                }
342            },
343        )
344    }
345}
346
347impl PtxParser for GeneralOperand {
348    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
349        alt(
350            mapc!(VectorOperand::parse(), GeneralOperand::Vec { operand }),
351            mapc!(Operand::parse(), GeneralOperand::Single { operand }),
352        )
353    }
354}
355
356fn handler_len_error(ty: &str, expected: &str, found: usize) -> PtxParseError {
357    PtxParseError {
358        kind: ParseErrorKind::InvalidLiteral(format!("{ty} expects {expected}, found {found}")),
359        span: Span::new(0, 0),
360    }
361}
362
363fn tex_operands()
364-> impl Fn(&mut PtxTokenStream) -> Result<(Vec<GeneralOperand>, Span), PtxParseError> {
365    between(
366        lbracket_p(),
367        rbracket_p(),
368        sep_by1(GeneralOperand::parse(), comma_p()),
369    )
370}
371
372impl PtxParser for TexHandler2 {
373    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
374        try_map(tex_operands(), |operands, span| {
375            if operands.len() != 2 {
376                return Err(handler_len_error(
377                    "TexHandler2",
378                    "exactly 2 operands",
379                    operands.len(),
380                ));
381            }
382            let mut iter = operands.into_iter();
383            let first = iter.next().unwrap();
384            let second = iter.next().unwrap();
385            let operands = [first, second];
386            ok!(TexHandler2 { operands })
387        })
388    }
389}
390
391impl PtxParser for TexHandler3 {
392    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
393        try_map(tex_operands(), |operands, span| {
394            if operands.len() != 3 {
395                return Err(handler_len_error(
396                    "TexHandler3",
397                    "exactly 3 operands",
398                    operands.len(),
399                ));
400            }
401            let mut iter = operands.into_iter();
402            let handle = iter.next().unwrap();
403            let sampler = iter.next().unwrap();
404            let coords = iter.next().unwrap();
405            ok!(TexHandler3 {
406                handle,
407                sampler,
408                coords
409            })
410        })
411    }
412}
413
414impl PtxParser for TexHandler3Optional {
415    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
416        try_map(tex_operands(), |operands, span| match operands.len() {
417            2 => {
418                let mut iter = operands.into_iter();
419                let handle = iter.next().unwrap();
420                let coords = iter.next().unwrap();
421                let sampler = None;
422                ok!(TexHandler3Optional {
423                    handle,
424                    sampler,
425                    coords
426                })
427            }
428            3 => {
429                let mut iter = operands.into_iter();
430                let handle = iter.next().unwrap();
431                let sampler = Some(iter.next().unwrap());
432                let coords = iter.next().unwrap();
433                ok!(TexHandler3Optional {
434                    handle,
435                    sampler,
436                    coords
437                })
438            }
439            other => Err(handler_len_error(
440                "TexHandler3Optional",
441                "2 or 3 operands",
442                other,
443            )),
444        })
445    }
446}
447
448impl PtxParser for SpecialRegister {
449    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
450        try_map(register_p(), |name, span| {
451            let raw = name.strip_prefix('%').ok_or_else(|| PtxParseError {
452                kind: ParseErrorKind::InvalidLiteral("expected % prefix".into()),
453                span,
454            })?;
455
456            let (base, axis_suffix) = split_axis_suffix(raw).map_err(|msg| PtxParseError {
457                kind: ParseErrorKind::InvalidLiteral(msg.into()),
458                span,
459            })?;
460            let base_lower = base.to_ascii_lowercase();
461
462            let mut node = match base_lower.as_str() {
463                "aggr_smem_size" => {
464                    ensure_no_axis(axis_suffix, span)?;
465                    c!(SpecialRegister::AggrSmemSize)
466                }
467                "dynamic_smem_size" => {
468                    ensure_no_axis(axis_suffix, span)?;
469                    c!(SpecialRegister::DynamicSmemSize)
470                }
471                "lanemask_gt" => {
472                    ensure_no_axis(axis_suffix, span)?;
473                    c!(SpecialRegister::LanemaskGt)
474                }
475                "reserved_smem_offset_begin" => {
476                    ensure_no_axis(axis_suffix, span)?;
477                    c!(SpecialRegister::ReservedSmemOffsetBegin)
478                }
479                "clock" => {
480                    ensure_no_axis(axis_suffix, span)?;
481                    c!(SpecialRegister::Clock)
482                }
483                "lanemask_le" => {
484                    ensure_no_axis(axis_suffix, span)?;
485                    c!(SpecialRegister::LanemaskLe)
486                }
487                "reserved_smem_offset_cap" => {
488                    ensure_no_axis(axis_suffix, span)?;
489                    c!(SpecialRegister::ReservedSmemOffsetCap)
490                }
491                "clock64" => {
492                    ensure_no_axis(axis_suffix, span)?;
493                    c!(SpecialRegister::Clock64)
494                }
495                "globaltimer" => {
496                    ensure_no_axis(axis_suffix, span)?;
497                    c!(SpecialRegister::Globaltimer)
498                }
499                "lanemask_lt" => {
500                    ensure_no_axis(axis_suffix, span)?;
501                    c!(SpecialRegister::LanemaskLt)
502                }
503                "reserved_smem_offset_end" => {
504                    ensure_no_axis(axis_suffix, span)?;
505                    c!(SpecialRegister::ReservedSmemOffsetEnd)
506                }
507                "globaltimer_hi" => {
508                    ensure_no_axis(axis_suffix, span)?;
509                    c!(SpecialRegister::GlobaltimerHi)
510                }
511                "nclusterid" => {
512                    ensure_no_axis(axis_suffix, span)?;
513                    c!(SpecialRegister::Nclusterid)
514                }
515                "smid" => {
516                    ensure_no_axis(axis_suffix, span)?;
517                    c!(SpecialRegister::Smid)
518                }
519                "globaltimer_lo" => {
520                    ensure_no_axis(axis_suffix, span)?;
521                    c!(SpecialRegister::GlobaltimerLo)
522                }
523                "gridid" => {
524                    ensure_no_axis(axis_suffix, span)?;
525                    c!(SpecialRegister::Gridid)
526                }
527                "nsmid" => {
528                    ensure_no_axis(axis_suffix, span)?;
529                    c!(SpecialRegister::Nsmid)
530                }
531                "total_smem_size" => {
532                    ensure_no_axis(axis_suffix, span)?;
533                    c!(SpecialRegister::TotalSmemSize)
534                }
535                "is_explicit_cluster" => {
536                    ensure_no_axis(axis_suffix, span)?;
537                    c!(SpecialRegister::IsExplicitCluster)
538                }
539                "warpid" => {
540                    ensure_no_axis(axis_suffix, span)?;
541                    c!(SpecialRegister::Warpid)
542                }
543                "clusterid" => {
544                    ensure_no_axis(axis_suffix, span)?;
545                    c!(SpecialRegister::Clusterid)
546                }
547                "laneid" => {
548                    ensure_no_axis(axis_suffix, span)?;
549                    c!(SpecialRegister::Laneid)
550                }
551                "nwarpid" => {
552                    ensure_no_axis(axis_suffix, span)?;
553                    c!(SpecialRegister::Nwarpid)
554                }
555                "warpsz" => {
556                    ensure_no_axis(axis_suffix, span)?;
557                    c!(SpecialRegister::WARPSZ)
558                }
559                "lanemask_eq" => {
560                    ensure_no_axis(axis_suffix, span)?;
561                    c!(SpecialRegister::LanemaskEq)
562                }
563                "current_graph_exec" => {
564                    ensure_no_axis(axis_suffix, span)?;
565                    c!(SpecialRegister::CurrentGraphExec)
566                }
567                "lanemask_ge" => {
568                    ensure_no_axis(axis_suffix, span)?;
569                    c!(SpecialRegister::LanemaskGe)
570                }
571                "cluster_ctaid" => {
572                    let axis = axis_with_span(axis_suffix, &span)?;
573                    c!(SpecialRegister::ClusterCtaid { axis })
574                }
575                "cluster_ctarank" => {
576                    let axis = axis_with_span(axis_suffix, &span)?;
577                    c!(SpecialRegister::ClusterCtarank { axis })
578                }
579                "nctaid" => {
580                    let axis = axis_with_span(axis_suffix, &span)?;
581                    c!(SpecialRegister::Nctaid { axis })
582                }
583                "tid" => {
584                    let axis = axis_with_span(axis_suffix, &span)?;
585                    c!(SpecialRegister::Tid { axis })
586                }
587                "cluster_nctaid" => {
588                    let axis = axis_with_span(axis_suffix, &span)?;
589                    c!(SpecialRegister::ClusterNctaid { axis })
590                }
591                "cluster_nctarank" => {
592                    let axis = axis_with_span(axis_suffix, &span)?;
593                    c!(SpecialRegister::ClusterNctarank { axis })
594                }
595                "ntid" => {
596                    let axis = axis_with_span(axis_suffix, &span)?;
597                    c!(SpecialRegister::Ntid { axis })
598                }
599                "ctaid" => {
600                    let axis = axis_with_span(axis_suffix, &span)?;
601                    c!(SpecialRegister::Ctaid { axis })
602                }
603                base if base.starts_with("envreg") => {
604                    ensure_no_axis(axis_suffix, span)?;
605                    let digits = base.strip_prefix("envreg").unwrap();
606                    let index = parse_index_suffix(digits, 31, "envreg", span)?;
607                    c!(SpecialRegister::Envreg { index })
608                }
609                base if base.starts_with("pm") && base.ends_with("_64") => {
610                    ensure_no_axis(axis_suffix, span)?;
611                    let digits = base
612                        .strip_prefix("pm")
613                        .and_then(|rest| rest.strip_suffix("_64"))
614                        .unwrap();
615                    let index = parse_index_suffix(digits, 7, "pm64", span)?;
616                    c!(SpecialRegister::Pm64 { index })
617                }
618                base if base.starts_with("pm") => {
619                    ensure_no_axis(axis_suffix, span)?;
620                    let digits = base.strip_prefix("pm").unwrap();
621                    let index = parse_index_suffix(digits, 7, "pm", span)?;
622                    c!(SpecialRegister::Pm { index })
623                }
624                base if base.starts_with("reserved_smem_offset_") => {
625                    ensure_no_axis(axis_suffix, span)?;
626                    let digits = base.strip_prefix("reserved_smem_offset_").unwrap();
627                    let index = parse_index_suffix(digits, 1, "reserved_smem_offset", span)?;
628                    c!(SpecialRegister::ReservedSmemOffset { index })
629                }
630                _ => {
631                    return err!(ParseErrorKind::InvalidLiteral(format!(
632                        "unknown special register: {name}"
633                    )));
634                }
635            };
636
637            node.set_span(span);
638            Ok(node)
639        })
640    }
641}
642
643impl PtxParser for Immediate {
644    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
645        mapc!(literal_p(), Immediate { value })
646    }
647}
648
649impl PtxParser for AddressBase {
650    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
651        alt(
652            mapc!(RegisterOperand::parse(), AddressBase::Register { operand }),
653            mapc!(VariableSymbol::parse(), AddressBase::Variable { symbol }),
654        )
655    }
656}
657
658impl PtxParser for AddressOffset {
659    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
660        alt(
661            mapc!(
662                seq(plus_p(), RegisterOperand::parse()),
663                AddressOffset::Register { _, operand }
664            ),
665            mapc!(
666                seq(Sign::parse(), Immediate::parse()),
667                AddressOffset::Immediate { sign, value }
668            ),
669        )
670    }
671}
672
673impl PtxParser for AddressOperand {
674    fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError> {
675        alt!(
676            mapc!(
677                seq(
678                    VariableSymbol::parse(),
679                    between(lbracket_p(), rbracket_p(), Immediate::parse()),
680                ),
681                AddressOperand::Array { base, index }
682            ),
683            mapc!(
684                between(
685                    lbracket_p(),
686                    rbracket_p(),
687                    seq(AddressBase::parse(), optional(AddressOffset::parse())),
688                ),
689                AddressOperand::Offset { base, offset }
690            ),
691            map(
692                between(
693                    lbracket_p(),
694                    rbracket_p(),
695                    seq(optional(Sign::parse()), Immediate::parse()),
696                ),
697                |(sign, mut addr), span| {
698                    if matches!(sign, Some(Sign::Negative { .. })) && !addr.value.starts_with('-') {
699                        addr.value = format!("-{}", addr.value);
700                    }
701                    c!(AddressOperand::ImmediateAddress { addr })
702                },
703            ),
704        )
705    }
706}
707
708fn split_axis_suffix(raw: &str) -> Result<(&str, Option<char>), &'static str> {
709    if let Some(idx) = raw.rfind('.') {
710        let (base, suffix) = raw.split_at(idx);
711        if suffix.len() != 2 {
712            return Err("invalid axis suffix");
713        }
714        let axis_char = suffix.chars().nth(1).unwrap();
715        Ok((base, Some(axis_char)))
716    } else {
717        Ok((raw, None))
718    }
719}
720
721fn axis_from_suffix(axis: Option<char>, span: Span) -> Result<Axis, PtxParseError> {
722    match axis.map(|c| c.to_ascii_lowercase()) {
723        None => Ok(Axis::None { span }),
724        Some('x') => Ok(Axis::X { span }),
725        Some('y') => Ok(Axis::Y { span }),
726        Some('z') => Ok(Axis::Z { span }),
727        Some(other) => Err(PtxParseError {
728            kind: ParseErrorKind::InvalidLiteral(format!(
729                "invalid axis '.{}' for special register",
730                other
731            )),
732            span,
733        }),
734    }
735}
736
737fn axis_with_span(axis: Option<char>, span: &Span) -> Result<Axis, PtxParseError> {
738    axis_from_suffix(axis, *span)
739}
740
741fn ensure_no_axis(axis: Option<char>, span: Span) -> Result<(), PtxParseError> {
742    if let Some(ch) = axis {
743        Err(PtxParseError {
744            kind: ParseErrorKind::InvalidLiteral(format!(
745                "register does not accept axis '.{}'",
746                ch
747            )),
748            span,
749        })
750    } else {
751        Ok(())
752    }
753}
ptx_parser/parser/common.rs

ptx_parser/parser/
common.rs
