ptx_parser/parser/
variable.rs

1use crate::{
2    lexer::PtxToken,
3    parser::{
4        PtxParseError, PtxParser, PtxTokenStream, Span, common::parse_u64_literal, invalid_literal,
5        peek_directive, unexpected_value,
6    },
7    r#type::{
8        common::{AddressSpace, AttributeDirective, DataLinkage, DataType},
9        variable::{
10            GlobalInitializer, InitializerValue, ModuleVariableDirective, NumericLiteral,
11            VariableDirective, VariableModifier,
12        },
13    },
14};
15
16const DATA_TYPE_NAMES: &[&str] = &[
17    "u8", "u16", "u32", "u64", "s8", "s16", "s32", "s64", "f16", "f16x2", "f32", "f64", "b8",
18    "b16", "b32", "b64", "b128", "pred",
19];
20
21#[derive(Clone, Copy, Debug, PartialEq, Eq)]
22enum VariableDirectiveKind {
23    Tex,
24    Shared,
25    Global,
26    Const,
27    Other,
28}
29
30fn is_data_type_directive(name: &str) -> bool {
31    DATA_TYPE_NAMES.iter().any(|candidate| candidate == &name)
32}
33
34fn is_vector_modifier(name: &str) -> bool {
35    let mut chars = name.chars();
36    match (chars.next(), chars.next()) {
37        (Some('v'), Some(digit)) if digit.is_ascii_digit() => chars.all(|ch| ch.is_ascii_digit()),
38        _ => false,
39    }
40}
41
42fn parse_alignment_value(stream: &mut PtxTokenStream) -> Result<u32, PtxParseError> {
43    let (value, value_span) = parse_u64_literal(stream)?;
44    if value > u32::MAX as u64 {
45        return Err(invalid_literal(
46            value_span,
47            "alignment value exceeds u32 range",
48        ));
49    }
50    Ok(value as u32)
51}
52
53fn parse_numeric_string(text: &str, span: Span) -> Result<u128, PtxParseError> {
54    text.parse::<u128>()
55        .map_err(|_| invalid_literal(span, "invalid integer literal"))
56}
57
58impl PtxParser for NumericLiteral {
59    fn parse(stream: &mut PtxTokenStream) -> Result<Self, PtxParseError> {
60        let negative = stream
61            .consume_if(|token| matches!(token, PtxToken::Minus))
62            .is_some();
63        let positive = stream
64            .consume_if(|token| matches!(token, PtxToken::Plus))
65            .is_some();
66
67        if negative && positive {
68            let (_, span) = stream.peek()?;
69            return Err(invalid_literal(
70                span.clone(),
71                "cannot have both '+' and '-' signs",
72            ));
73        }
74
75        let (token, span_ref) = stream.consume()?;
76        let span = span_ref.clone();
77        match token {
78            PtxToken::DecimalInteger(text) => {
79                let value = parse_numeric_string(text.as_str(), span.clone())?;
80                if negative {
81                    if value > (i64::MAX as u128) + 1 {
82                        return Err(invalid_literal(span.clone(), "signed integer underflow"));
83                    }
84                    let signed = -(value as i128);
85                    Ok(NumericLiteral::Signed(signed as i64))
86                } else {
87                    if value > u64::MAX as u128 {
88                        return Err(invalid_literal(span.clone(), "unsigned integer overflow"));
89                    }
90                    Ok(NumericLiteral::Unsigned(value as u64))
91                }
92            }
93            PtxToken::HexInteger(text) => {
94                let stripped = text
95                    .strip_prefix("0x")
96                    .or_else(|| text.strip_prefix("0X"))
97                    .unwrap_or(text.as_str());
98                let value = u128::from_str_radix(stripped, 16)
99                    .map_err(|_| invalid_literal(span.clone(), "invalid hex literal"))?;
100                if negative {
101                    if value > (i64::MAX as u128) + 1 {
102                        return Err(invalid_literal(span.clone(), "signed integer underflow"));
103                    }
104                    let signed = -(value as i128);
105                    Ok(NumericLiteral::Signed(signed as i64))
106                } else {
107                    if value > u64::MAX as u128 {
108                        return Err(invalid_literal(span.clone(), "unsigned integer overflow"));
109                    }
110                    Ok(NumericLiteral::Unsigned(value as u64))
111                }
112            }
113            PtxToken::BinaryInteger(text) => {
114                let stripped = text
115                    .strip_prefix("0b")
116                    .or_else(|| text.strip_prefix("0B"))
117                    .unwrap_or(text.as_str());
118                let value = u128::from_str_radix(stripped, 2)
119                    .map_err(|_| invalid_literal(span.clone(), "invalid binary literal"))?;
120                if negative {
121                    if value > (i64::MAX as u128) + 1 {
122                        return Err(invalid_literal(span.clone(), "signed integer underflow"));
123                    }
124                    let signed = -(value as i128);
125                    Ok(NumericLiteral::Signed(signed as i64))
126                } else {
127                    if value > u64::MAX as u128 {
128                        return Err(invalid_literal(span.clone(), "unsigned integer overflow"));
129                    }
130                    Ok(NumericLiteral::Unsigned(value as u64))
131                }
132            }
133            PtxToken::OctalInteger(text) => {
134                let stripped = &text.as_str()[1..];
135                let value = u128::from_str_radix(stripped, 8)
136                    .map_err(|_| invalid_literal(span.clone(), "invalid octal literal"))?;
137                if negative {
138                    if value > (i64::MAX as u128) + 1 {
139                        return Err(invalid_literal(span.clone(), "signed integer underflow"));
140                    }
141                    let signed = -(value as i128);
142                    Ok(NumericLiteral::Signed(signed as i64))
143                } else {
144                    if value > u64::MAX as u128 {
145                        return Err(invalid_literal(span.clone(), "unsigned integer overflow"));
146                    }
147                    Ok(NumericLiteral::Unsigned(value as u64))
148                }
149            }
150            PtxToken::Float(text) | PtxToken::FloatExponent(text) => {
151                let mut value = text
152                    .parse::<f64>()
153                    .map_err(|_| invalid_literal(span.clone(), "invalid floating-point literal"))?;
154                if negative {
155                    value = -value;
156                }
157                Ok(NumericLiteral::Float64(value.to_bits()))
158            }
159            PtxToken::HexFloat(text) => {
160                if text.len() < 3 {
161                    return Err(invalid_literal(
162                        span.clone(),
163                        "invalid hexadecimal float literal",
164                    ));
165                }
166                let (prefix, digits) = text.split_at(2);
167                match prefix.to_ascii_lowercase().as_str() {
168                    "0f" => {
169                        let mut bits = u32::from_str_radix(digits, 16)
170                            .map_err(|_| invalid_literal(span.clone(), "invalid float literal"))?;
171                        if negative {
172                            bits ^= 0x8000_0000;
173                        }
174                        Ok(NumericLiteral::Float32(bits))
175                    }
176                    "0d" => {
177                        let mut bits = u64::from_str_radix(digits, 16)
178                            .map_err(|_| invalid_literal(span.clone(), "invalid float literal"))?;
179                        if negative {
180                            bits ^= 0x8000_0000_0000_0000;
181                        }
182                        Ok(NumericLiteral::Float64(bits))
183                    }
184                    _ => Err(invalid_literal(
185                        span.clone(),
186                        "hexadecimal float must start with 0f or 0d",
187                    )),
188                }
189            }
190            _ => Err(unexpected_value(
191                span.clone(),
192                &["numeric literal"],
193                format!("{token:?}"),
194            )),
195        }
196    }
197}
198
199impl PtxParser for InitializerValue {
200    fn parse(stream: &mut PtxTokenStream) -> Result<Self, PtxParseError> {
201        if let Some((token, span)) = stream.peek().ok() {
202            match token {
203                PtxToken::StringLiteral(value) => {
204                    let value = value.clone();
205                    stream.consume()?;
206                    return Ok(InitializerValue::StringLiteral(value));
207                }
208                PtxToken::Identifier(_) => {
209                    let (symbol, _) = stream.expect_identifier()?;
210                    return Ok(InitializerValue::Symbol(symbol));
211                }
212                PtxToken::Plus | PtxToken::Minus => {
213                    let literal = NumericLiteral::parse(stream)?;
214                    return Ok(InitializerValue::Numeric(literal));
215                }
216                PtxToken::DecimalInteger(_)
217                | PtxToken::HexInteger(_)
218                | PtxToken::BinaryInteger(_)
219                | PtxToken::OctalInteger(_)
220                | PtxToken::Float(_)
221                | PtxToken::FloatExponent(_)
222                | PtxToken::HexFloat(_) => {
223                    let literal = NumericLiteral::parse(stream)?;
224                    return Ok(InitializerValue::Numeric(literal));
225                }
226                _ => {
227                    return Err(unexpected_value(
228                        span.clone(),
229                        &["numeric literal", "symbol", "string literal"],
230                        format!("{token:?}"),
231                    ));
232                }
233            }
234        }
235        let span = stream.peek()?.1.clone();
236        Err(unexpected_value(
237            span,
238            &["numeric literal", "symbol", "string literal"],
239            "end of input".to_string(),
240        ))
241    }
242}
243
244impl PtxParser for GlobalInitializer {
245    fn parse(stream: &mut PtxTokenStream) -> Result<Self, PtxParseError> {
246        if stream
247            .consume_if(|token| matches!(token, PtxToken::LBrace))
248            .is_some()
249        {
250            let mut children = Vec::new();
251            if !stream.check(|token| matches!(token, PtxToken::RBrace)) {
252                loop {
253                    let initializer = GlobalInitializer::parse(stream)?;
254                    children.push(initializer);
255                    if !(stream
256                        .consume_if(|token| matches!(token, PtxToken::Comma))
257                        .is_some())
258                    {
259                        break;
260                    }
261                }
262            }
263            stream.expect(&PtxToken::RBrace)?;
264            Ok(GlobalInitializer::Aggregate(children))
265        } else {
266            let value = InitializerValue::parse(stream)?;
267            Ok(GlobalInitializer::Scalar(value))
268        }
269    }
270}
271
272impl PtxParser for VariableModifier {
273    fn parse(stream: &mut PtxTokenStream) -> Result<Self, PtxParseError> {
274        let (directive, span_ref) = stream.expect_directive()?;
275        let span = span_ref.clone();
276        match directive.as_str() {
277            "align" => {
278                let value = parse_alignment_value(stream)?;
279                Ok(VariableModifier::Alignment(value))
280            }
281            "ptr" => Ok(VariableModifier::Ptr),
282            "visible" => Ok(VariableModifier::Linkage(DataLinkage::Visible)),
283            "extern" => Ok(VariableModifier::Linkage(DataLinkage::Extern)),
284            "weak" => Ok(VariableModifier::Linkage(DataLinkage::Weak)),
285            "common" => Ok(VariableModifier::Linkage(DataLinkage::Common)),
286            other if is_vector_modifier(other) => {
287                let digits = &other[1..];
288                let value = digits
289                    .parse::<u32>()
290                    .map_err(|_| invalid_literal(span.clone(), "invalid vector width"))?;
291                Ok(VariableModifier::Vector(value))
292            }
293            other => Err(unexpected_value(
294                span.clone(),
295                &[
296                    ".align", ".ptr", ".visible", ".extern", ".weak", ".common", ".vN",
297                ],
298                format!(".{other}"),
299            )),
300        }
301    }
302}
303
304impl VariableDirective {
305    fn parse_with_kind(
306        stream: &mut PtxTokenStream,
307    ) -> Result<(VariableDirective, VariableDirectiveKind, Option<Span>), PtxParseError> {
308        let first_span = stream.peek().ok().map(|(_, span)| span.clone());
309
310        let mut address_space: Option<AddressSpace> = None;
311        let mut attributes = Vec::new();
312        let mut modifiers = Vec::new();
313        let mut ty: Option<DataType> = None;
314        let mut array = Vec::new();
315        let mut initializer = None;
316        let mut seen_tex = false;
317        let mut kind = VariableDirectiveKind::Other;
318        let mut kind_span = None;
319
320        loop {
321            let Some((directive, directive_span)) = peek_directive(stream)? else {
322                break;
323            };
324            match directive.as_str() {
325                "tex" => {
326                    stream.expect_directive()?;
327                    if !seen_tex {
328                        seen_tex = true;
329                        kind = VariableDirectiveKind::Tex;
330                        kind_span = Some(directive_span);
331                    }
332                }
333                "global" | "const" | "shared" | "local" | "param" | "reg" => {
334                    if address_space.is_some() {
335                        return Err(unexpected_value(
336                            directive_span.clone(),
337                            &["single address space qualifier"],
338                            format!(".{directive}"),
339                        ));
340                    }
341                    let space = AddressSpace::parse(stream)?;
342                    address_space = Some(space);
343                    match space {
344                        AddressSpace::Global => {
345                            kind = VariableDirectiveKind::Global;
346                            kind_span = Some(directive_span.clone());
347                        }
348                        AddressSpace::Const => {
349                            kind = VariableDirectiveKind::Const;
350                            kind_span = Some(directive_span.clone());
351                        }
352                        AddressSpace::Shared => {
353                            kind = VariableDirectiveKind::Shared;
354                            kind_span = Some(directive_span.clone());
355                        }
356                        _ => {}
357                    }
358                }
359                "managed" | "unified" => {
360                    attributes.push(AttributeDirective::parse(stream)?);
361                }
362                "align" | "ptr" | "visible" | "extern" | "weak" | "common" => {
363                    modifiers.push(VariableModifier::parse(stream)?);
364                }
365                other if is_vector_modifier(other) => {
366                    modifiers.push(VariableModifier::parse(stream)?);
367                }
368                other if is_data_type_directive(other) => {
369                    if ty.is_some() {
370                        return Err(unexpected_value(
371                            directive_span.clone(),
372                            &["single data type qualifier"],
373                            format!(".{other}"),
374                        ));
375                    }
376                    ty = Some(DataType::parse(stream)?);
377                }
378                _ => break,
379            }
380        }
381
382        let (name, _) = stream.expect_identifier()?;
383
384        loop {
385            if stream
386                .consume_if(|token| matches!(token, PtxToken::LBracket))
387                .is_none()
388            {
389                break;
390            }
391
392            if stream
393                .consume_if(|token| matches!(token, PtxToken::RBracket))
394                .is_some()
395            {
396                array.push(None);
397                continue;
398            }
399
400            let size_span = stream.peek()?.1.clone();
401            let literal = NumericLiteral::parse(stream)?;
402            let size = match literal {
403                NumericLiteral::Unsigned(value) => value,
404                NumericLiteral::Signed(value) if value >= 0 => value as u64,
405                _ => {
406                    return Err(invalid_literal(
407                        size_span.clone(),
408                        "array size must be a non-negative integer",
409                    ));
410                }
411            };
412
413            stream.expect(&PtxToken::RBracket)?;
414            array.push(Some(size));
415        }
416
417        if stream
418            .consume_if(|token| matches!(token, PtxToken::Equals))
419            .is_some()
420        {
421            initializer = Some(GlobalInitializer::parse(stream)?);
422        }
423
424        stream.expect(&PtxToken::Semicolon)?;
425
426        let mut final_kind = kind;
427        if seen_tex {
428            final_kind = VariableDirectiveKind::Tex;
429        } else if matches!(final_kind, VariableDirectiveKind::Other) {
430            final_kind = match address_space {
431                Some(AddressSpace::Shared) => VariableDirectiveKind::Shared,
432                Some(AddressSpace::Global) => VariableDirectiveKind::Global,
433                Some(AddressSpace::Const) => VariableDirectiveKind::Const,
434                _ => VariableDirectiveKind::Other,
435            };
436        }
437
438        let directive = VariableDirective {
439            address_space,
440            attributes,
441            ty,
442            modifiers,
443            name,
444            array,
445            initializer,
446            raw: String::new(),
447        };
448
449        Ok((directive, final_kind, kind_span.or(first_span)))
450    }
451}
452
453impl PtxParser for VariableDirective {
454    fn parse(stream: &mut PtxTokenStream) -> Result<Self, PtxParseError> {
455        let (directive, _, _) = VariableDirective::parse_with_kind(stream)?;
456        Ok(directive)
457    }
458}
459
460impl PtxParser for ModuleVariableDirective {
461    fn parse(stream: &mut PtxTokenStream) -> Result<Self, PtxParseError> {
462        let (directive, kind, span) = VariableDirective::parse_with_kind(stream)?;
463        match kind {
464            VariableDirectiveKind::Tex => Ok(ModuleVariableDirective::Tex(directive)),
465            VariableDirectiveKind::Shared => Ok(ModuleVariableDirective::Shared(directive)),
466            VariableDirectiveKind::Global => Ok(ModuleVariableDirective::Global(directive)),
467            VariableDirectiveKind::Const => Ok(ModuleVariableDirective::Const(directive)),
468            VariableDirectiveKind::Other => Err(unexpected_value(
469                span.unwrap_or(0..0),
470                &[".tex", ".shared", ".global", ".const"],
471                "variable directive".to_string(),
472            )),
473        }
474    }
475}
ptx_parser/parser/variable.rs

ptx_parser/parser/
variable.rs
