ptx_parser/parser/
common.rs

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

ptx_parser/parser/
common.rs
