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

ptx_parser/parser/
common.rs
