ptx_parser/parser/mod.rs
1use crate::{LexError, lexer::PtxToken, span};
2use serde::Serialize;
3#[cfg(debug_assertions)]
4use stacker;
5use thiserror::Error;
6
7pub(crate) mod common;
8pub(crate) mod function;
9pub(crate) mod instruction;
10pub(crate) mod module;
11pub(crate) mod util;
12pub(crate) mod variable;
13
14#[derive(Debug, Default, PartialEq, Eq, Copy, Clone, Serialize)]
15pub struct Span {
16 pub start: usize,
17 pub end: usize,
18}
19
20impl Span {
21 pub const fn new(start: usize, end: usize) -> Self {
22 Self { start, end }
23 }
24}
25
26impl From<std::ops::Range<usize>> for Span {
27 fn from(range: std::ops::Range<usize>) -> Self {
28 Span::new(range.start, range.end)
29 }
30}
31
32impl From<Span> for std::ops::Range<usize> {
33 fn from(span: Span) -> Self {
34 span.start..span.end
35 }
36}
37
38/// Macro to create an UnexpectedToken error with expected and found values.
39///
40/// # Usage
41/// ```ignore
42/// unexpected_token!(span, ["expected1", "expected2"], "found_value")
43/// unexpected_token!(span, vec!["expected1".to_string()], format!("{:?}", token))
44/// ```
45#[macro_export]
46macro_rules! unexpected_token {
47 ($span:expr, $expected:expr, $found:expr) => {
48 $crate::parser::PtxParseError {
49 kind: $crate::parser::ParseErrorKind::UnexpectedToken {
50 expected: $expected.iter().map(|s| s.to_string()).collect(),
51 found: $found,
52 },
53 span: $span,
54 }
55 };
56}
57
58/// Macro to check if in partial mode and return error if so.
59/// Use this in token-based methods that should only work in complete mode.
60///
61/// # Usage
62/// ```ignore
63/// reject_partial_mode!(self);
64/// ```
65macro_rules! reject_partial_mode {
66 ($self:expr) => {
67 if $self.index.1.is_some() {
68 let span = $self
69 .tokens
70 .get($self.index.0)
71 .map_or(span!(0..0), |(_, s)| *s);
72 return Err($crate::parser::PtxParseError {
73 kind: $crate::parser::ParseErrorKind::InvalidModeForTokenMethod,
74 span,
75 });
76 }
77 };
78}
79
80/// Macro to create an UnexpectedToken error when no candidates match.
81///
82/// # Usage
83/// ```ignore
84/// no_candidate_match!(self, candidates)
85/// ```
86macro_rules! no_candidate_match {
87 ($self:expr, $candidates:expr) => {{
88 let span = $self
89 .tokens
90 .get($self.index.0)
91 .map_or(span!(0..0), |(_, s)| *s);
92 $crate::parser::PtxParseError {
93 kind: $crate::parser::ParseErrorKind::UnexpectedToken {
94 expected: $candidates.iter().map(|s| s.to_string()).collect(),
95 found: "no match".to_string(),
96 },
97 span,
98 }
99 }};
100}
101
102/// Macro to build a standard unexpected-value parse error.
103#[macro_export]
104macro_rules! unexpected_value {
105 ($span:expr, $expected:expr, $found:expr) => {
106 $crate::parser::PtxParseError {
107 kind: $crate::parser::ParseErrorKind::UnexpectedToken {
108 expected: $expected.iter().map(|s| s.to_string()).collect(),
109 found: $found.into(),
110 },
111 span: $span,
112 }
113 };
114}
115
116/// Kinds of parse errors that can occur during PTX parsing.
117#[derive(Debug, Clone, PartialEq, Eq, Error)]
118pub enum ParseErrorKind {
119 #[error("unexpected token: expected one of {expected:?}, found {found}")]
120 UnexpectedToken {
121 expected: Vec<String>,
122 found: String,
123 },
124 #[error("unexpected end of input")]
125 UnexpectedEof,
126 #[error("invalid literal: {0}")]
127 InvalidLiteral(String),
128 #[error("cannot use token-based methods in partial mode")]
129 InvalidModeForTokenMethod,
130}
131
132/// PTX parsing error with location information.
133#[derive(Debug, Clone, PartialEq, Eq, Error)]
134#[error("parsing error at {span:?}: {kind}")]
135pub struct PtxParseError {
136 pub kind: ParseErrorKind,
137 pub span: Span,
138}
139
140impl From<LexError> for PtxParseError {
141 fn from(err: LexError) -> Self {
142 PtxParseError {
143 kind: ParseErrorKind::InvalidLiteral("lexing failed".into()),
144 span: err.span,
145 }
146 }
147}
148
149/// Represents a position in the token stream,
150/// index of the token and optional char offset within the token.
151pub type StreamPosition = (usize, Option<usize>);
152
153/// Token stream wrapper for parsing PTX tokens.
154///
155/// This struct provides methods for consuming and inspecting tokens during parsing.
156pub struct PtxTokenStream<'a> {
157 tokens: &'a [(PtxToken, Span)],
158 /// Current position (index) in the tokens list
159 index: StreamPosition,
160}
161
162impl<'a> PtxTokenStream<'a> {
163 pub fn new(tokens: &'a [(PtxToken, Span)]) -> Self {
164 Self {
165 tokens,
166 index: (0, None),
167 }
168 }
169
170 /// Peek at the next token without consuming it.
171 ///
172 /// # Behavior for complete mode
173 ///
174 /// Returns the token at the current stream position without advancing the position.
175 /// This is a simple array lookup at `index.0`.
176 ///
177 /// # Behavior for partial mode
178 ///
179 /// Returns an error (InvalidModeForTokenMethod). This method only operates on whole tokens
180 /// and cannot be used during partial (character-by-character) matching mode.
181 ///
182 /// # Returns
183 ///
184 /// - `Ok(&(PtxToken, Span))` - The token and its span
185 /// - `Err(PtxParseError)` - If at end of stream (UnexpectedEof) or in partial mode (InvalidModeForTokenMethod)
186 pub fn peek(&self) -> Result<&'a (PtxToken, Span), PtxParseError> {
187 reject_partial_mode!(self);
188 self.tokens.get(self.index.0).ok_or_else(|| {
189 // If the stream is empty, return an EOF error
190 let span = self.tokens.last().map_or(span!(0..0), |(_, s)| *s);
191 PtxParseError {
192 kind: ParseErrorKind::UnexpectedEof,
193 span,
194 }
195 })
196 }
197
198 /// Peek at the token `offset` positions ahead without consuming it.
199 ///
200 /// Behaves like `peek()` but allows inspecting future tokens in complete mode.
201 pub fn peek_n(&self, offset: usize) -> Result<&'a (PtxToken, Span), PtxParseError> {
202 reject_partial_mode!(self);
203 self.tokens.get(self.index.0 + offset).ok_or_else(|| {
204 let span = self.tokens.last().map_or(span!(0..0), |(_, s)| *s);
205 PtxParseError {
206 kind: ParseErrorKind::UnexpectedEof,
207 span,
208 }
209 })
210 }
211
212 /// Consume and return the next token.
213 ///
214 /// # Behavior for complete mode
215 ///
216 /// Advances the stream position by one token (increments `index.0`).
217 /// Returns the token that was at the current position before advancing.
218 ///
219 /// # Behavior for partial mode
220 ///
221 /// Returns an error (InvalidModeForTokenMethod). This method only operates on whole tokens
222 /// and cannot be used during partial (character-by-character) matching mode.
223 ///
224 /// # Returns
225 ///
226 /// - `Ok(&(PtxToken, Span))` - The consumed token and its span
227 /// - `Err(PtxParseError)` - If at end of stream (UnexpectedEof) or in partial mode (InvalidModeForTokenMethod)
228 pub fn consume(&mut self) -> Result<&'a (PtxToken, Span), PtxParseError> {
229 reject_partial_mode!(self);
230 let token = self.peek()?;
231 self.index.0 += 1;
232 Ok(token)
233 }
234
235 /// Conditionally consume the next token if it matches the predicate.
236 ///
237 /// # Returns
238 ///
239 /// - `Some(&(PtxToken, Span))` - If the predicate returns true, consumes and returns the token
240 /// - `None` - If the predicate returns false or if at end of stream
241 pub fn consume_if<F>(&mut self, predicate: F) -> Option<&'a (PtxToken, Span)>
242 where
243 F: FnOnce(&PtxToken) -> bool,
244 {
245 if self.index.1.is_some() {
246 return None; // In partial mode
247 }
248 if let Ok((token, _)) = self.peek() {
249 if predicate(token) {
250 self.index.0 += 1;
251 return self.tokens.get(self.index.0 - 1);
252 }
253 }
254 None
255 }
256
257 /// Check if the next token is the expected type, and if so, consume it.
258 /// Otherwise, return an error and do NOT consume the token.
259 ///
260 /// # Behavior for complete mode
261 ///
262 /// Peeks at the current token and checks if its discriminant (variant type) matches
263 /// the expected token discriminant. If it matches, advances the stream by one token
264 /// and returns the token. If it doesn't match, returns an UnexpectedToken error
265 /// without consuming anything.
266 ///
267 /// # Behavior for partial mode
268 ///
269 /// Returns an error (InvalidModeForTokenMethod). This method only operates on whole tokens
270 /// and cannot be used during partial (character-by-character) matching mode.
271 ///
272 /// # Returns
273 ///
274 /// - `Ok(&(PtxToken, Span))` - The matched and consumed token
275 /// - `Err(PtxParseError)` - If token doesn't match (UnexpectedToken) or in partial mode (InvalidModeForTokenMethod)
276 pub fn expect(&mut self, expected: &PtxToken) -> Result<&'a (PtxToken, Span), PtxParseError> {
277 reject_partial_mode!(self);
278 let token_pair = self.peek()?;
279 let (token, span) = token_pair;
280 if std::mem::discriminant(token) == std::mem::discriminant(expected) {
281 self.index.0 += 1;
282 Ok(token_pair)
283 } else {
284 Err(unexpected_token!(
285 *span,
286 &[format!("{:?}", expected)],
287 format!("{:?}", token)
288 ))
289 }
290 }
291
292 /// Generic helper to extract a String value from a token variant.
293 /// Returns the extracted string and span if the pattern matches, otherwise returns an error.
294 ///
295 /// # Behavior for complete mode
296 ///
297 /// Peeks at the current token and attempts to extract a string value using the provided
298 /// extractor function. If extraction succeeds, advances the stream by one token and returns
299 /// the extracted string with its span. If extraction fails, returns an UnexpectedToken error.
300 ///
301 /// # Behavior for partial mode
302 ///
303 /// Returns an error (InvalidModeForTokenMethod). This method only operates on whole tokens
304 /// and cannot be used during partial (character-by-character) matching mode.
305 ///
306 /// # Returns
307 ///
308 /// - `Ok((String, Span))` - The extracted string value and its span
309 /// - `Err(PtxParseError)` - If extraction fails (UnexpectedToken) or in partial mode (InvalidModeForTokenMethod)
310 fn expect_token_with_string<F>(
311 &mut self,
312 expected_name: &str,
313 extractor: F,
314 ) -> Result<(String, Span), PtxParseError>
315 where
316 F: FnOnce(&PtxToken) -> Option<String>,
317 {
318 reject_partial_mode!(self);
319 let (token, span_ref) = self.peek()?;
320 if let Some(value) = extractor(token) {
321 let span = *span_ref;
322 self.index.0 += 1;
323 Ok((value, span))
324 } else {
325 Err(unexpected_token!(
326 *span_ref,
327 &[expected_name.to_string()],
328 format!("{:?}", token)
329 ))
330 }
331 }
332
333 /// Check if the next token is an identifier, and if so, consume it and return the String.
334 ///
335 /// # Behavior for complete mode
336 ///
337 /// Expects the current token to be an Identifier, consumes it, and returns its string value.
338 ///
339 /// # Behavior for partial mode
340 ///
341 /// Returns an error (InvalidModeForTokenMethod). This is a token-based method.
342 pub fn expect_identifier(&mut self) -> Result<(String, Span), PtxParseError> {
343 self.expect_token_with_string("Identifier", |token| {
344 if let PtxToken::Identifier(name) = token {
345 Some(name.clone())
346 } else {
347 None
348 }
349 })
350 }
351
352 /// Check if the next token is a register, and if so, consume it and return the String.
353 ///
354 /// # Behavior for complete mode
355 ///
356 /// Expects the current token to be a Register, consumes it, and returns its string value.
357 ///
358 /// # Behavior for partial mode
359 ///
360 /// Returns an error (InvalidModeForTokenMethod). This is a token-based method.
361 pub fn expect_register(&mut self) -> Result<(String, Span), PtxParseError> {
362 self.expect_token_with_string("Register", |token| {
363 if let PtxToken::Register(name) = token {
364 Some(name.clone())
365 } else {
366 None
367 }
368 })
369 }
370
371 /// Check if the next token is a directive (Dot + Identifier), and if so, consume them and return the String.
372 ///
373 /// # Behavior for complete mode
374 ///
375 /// Expects a Dot token followed by an Identifier token, consumes both, and returns the
376 /// identifier string with a combined span covering both tokens.
377 ///
378 /// # Behavior for partial mode
379 ///
380 /// Returns an error (InvalidModeForTokenMethod). This is a token-based method.
381 pub fn expect_directive(&mut self) -> Result<(String, Span), PtxParseError> {
382 let (_, dot_span) = self.expect(&PtxToken::Dot)?;
383 let (name, id_span) = self.expect_identifier()?;
384 let span = Span::new(dot_span.start, id_span.end);
385 Ok((name, span))
386 }
387
388 /// Internal helper to match a string pattern against the token stream.
389 /// Returns true if the entire pattern matches and consumes the matched portion.
390 /// Returns false if matching fails (does not modify stream state on failure).
391 ///
392 /// Supports both complete mode (whole token matching) and partial mode (char-by-char).
393 fn match_string_internal(&mut self, pattern: &str) -> bool {
394 let start_pos = self.position();
395 let mut pattern_chars = pattern.chars().peekable();
396
397 loop {
398 // Check if we've consumed the entire pattern
399 if pattern_chars.peek().is_none() {
400 return true; // Successfully matched
401 }
402
403 // Check if we've run out of tokens
404 if self.index.0 >= self.tokens.len() {
405 self.set_position(start_pos);
406 return false;
407 }
408
409 let (token, _span) = &self.tokens[self.index.0];
410 let token_str = token.as_str();
411
412 if let Some(char_offset) = self.index.1 {
413 // Partial mode: match character-by-character
414 let token_chars: Vec<char> = token_str.chars().collect();
415
416 if char_offset >= token_chars.len() {
417 // Consumed entire token, advance to next
418 self.index.0 += 1;
419 self.index.1 = Some(0);
420 continue;
421 }
422
423 // Try to match remaining pattern chars against remaining token chars
424 let mut offset = char_offset;
425 while offset < token_chars.len() && pattern_chars.peek().is_some() {
426 if Some(&token_chars[offset]) == pattern_chars.peek() {
427 pattern_chars.next();
428 offset += 1;
429 } else {
430 // Mismatch
431 self.set_position(start_pos);
432 return false;
433 }
434 }
435 self.index.1 = Some(offset);
436 } else {
437 // Complete mode: match whole token string representation
438 let token_chars: Vec<char> = token_str.chars().collect();
439 let mut token_idx = 0;
440
441 while token_idx < token_chars.len() && pattern_chars.peek().is_some() {
442 if Some(&token_chars[token_idx]) == pattern_chars.peek() {
443 pattern_chars.next();
444 token_idx += 1;
445 } else {
446 // Mismatch
447 self.set_position(start_pos);
448 return false;
449 }
450 }
451
452 // Check if we consumed the entire token
453 if token_idx == token_chars.len() {
454 self.index.0 += 1;
455 } else if pattern_chars.peek().is_none() {
456 // Pattern matched but didn't consume entire token - this is an error in complete mode
457 self.set_position(start_pos);
458 return false;
459 }
460 }
461 }
462 }
463
464 /// Try to match and consume a sequence of tokens that matches one of the candidate strings.
465 /// Returns the index of the matched candidate.
466 ///
467 /// This is used for parsing modifiers that may contain :: sequences like ".to::cluster"
468 /// The candidates should include the leading dot (e.g., [".to::cluster", ".to::cta"])
469 ///
470 /// # Behavior for complete mode
471 ///
472 /// Tries to match each candidate string against the token stream by consuming whole tokens.
473 /// Returns the index of the first candidate that matches. Uses backtracking (position/set_position)
474 /// to try each candidate without consuming tokens on failed attempts.
475 ///
476 /// # Behavior for partial mode
477 ///
478 /// Supports character-by-character matching within tokens using the char offset.
479 /// This allows matching patterns that span across token boundaries or within tokens.
480 /// Uses backtracking to restore position when a candidate fails to match.
481 pub fn expect_strings(&mut self, candidates: &[&str]) -> Result<usize, PtxParseError> {
482 let start_pos = self.position();
483
484 for (idx, candidate) in candidates.iter().enumerate() {
485 self.set_position(start_pos);
486
487 // Try to match this candidate
488 if self.match_string_internal(candidate) {
489 return Ok(idx);
490 }
491 }
492
493 // None matched, restore position and create error
494 self.set_position(start_pos);
495 Err(no_candidate_match!(self, candidates))
496 }
497
498 /// Expect that the next sequence of tokens matches the given string pattern.
499 ///
500 /// # Behavior for complete mode
501 ///
502 /// Matches the pattern against the token stream by consuming whole tokens.
503 /// Each token's string representation must match consecutive characters in the pattern.
504 /// The match succeeds only if the entire pattern is consumed and tokens are fully consumed.
505 ///
506 /// # Behavior for partial mode
507 ///
508 /// Matches the pattern character-by-character against the token stream using the
509 /// character offset for partial token matching. This allows matching patterns that
510 /// don't align with token boundaries. If all characters match, the stream advances.
511 /// If any character fails to match, the stream position is restored.
512 ///
513 /// # Returns
514 ///
515 /// - `Ok(())` if the entire pattern was successfully matched (consumed)
516 /// - `Err(PtxParseError)` if matching failed (UnexpectedToken)
517 pub fn expect_string(&mut self, expected: &str) -> Result<(), PtxParseError> {
518 let start_pos = self.position();
519 if self.match_string_internal(expected) {
520 Ok(())
521 } else {
522 self.set_position(start_pos);
523 Err(no_candidate_match!(self, &[expected]))
524 }
525 }
526
527 /// Ensure we're in complete mode (not in partial token mode).
528 /// This is a no-op in complete mode, and succeeds as long as we're not mid-token.
529 /// Used by generated parsers to enforce token boundaries.
530 pub fn expect_complete(&mut self) -> Result<(), PtxParseError> {
531 if self.index.1.is_some() {
532 let span = self
533 .tokens
534 .get(self.index.0)
535 .map_or(span!(0..0), |(_, s)| *s);
536 return Err(PtxParseError {
537 kind: ParseErrorKind::InvalidModeForTokenMethod,
538 span,
539 });
540 }
541 Ok(())
542 }
543
544 /// Execute a function in partial token mode, enabling character-by-character matching.
545 ///
546 /// # Behavior
547 ///
548 /// This method switches the stream from complete mode to partial mode by setting the
549 /// character offset to `Some(0)`. While in partial mode, string-based methods like
550 /// `expect_string()` can match patterns character-by-character within tokens.
551 ///
552 /// After the closure completes:
553 /// - If the char offset is non-zero, validates that the current token was fully consumed
554 /// - If not fully consumed, reverts to the starting position and returns an error
555 /// - Always resets the mode back to complete mode (sets `index.1` to `None`)
556 ///
557 /// # Errors
558 ///
559 /// Returns an error if:
560 /// - The closure returns an error
561 /// - The token was partially consumed but not completely consumed (incomplete match)
562 ///
563 /// # Panics
564 ///
565 /// Panics if already in partial mode (char offset is already `Some`).
566 pub fn with_partial_token_mode<F, R>(&mut self, f: F) -> Result<R, PtxParseError>
567 where
568 F: FnOnce(&mut PtxTokenStream) -> Result<R, PtxParseError>,
569 {
570 let start_index = self.index;
571 assert!(self.index.1.is_none(), "Already in partial mode");
572 self.index.1 = Some(0);
573 let result = f(self);
574
575 // Check if char offset has consumed the entire token
576 if let Some(char_offset) = self.index.1 {
577 if char_offset != 0 {
578 // if consumed entire token, ok; else, reset position and error
579 if let Some((token, span)) = self.tokens.get(self.index.0) {
580 if token.len() != char_offset {
581 self.index = start_index;
582 return Err(unexpected_token!(
583 *span,
584 &["fully consumed token".to_string()],
585 format!("partially consumed {:?}", token)
586 ));
587 } else {
588 // Token was fully consumed, advance to next token
589 self.index.0 += 1;
590 }
591 }
592 }
593 }
594 self.index.1 = None;
595 result
596 }
597
598 /// Execute a closure with automatic backtracking and span tracking.
599 ///
600 /// Saves the current stream position before running `f`. If `f` returns an
601 /// error, the stream position (including partial-mode offsets) is restored.
602 /// When `f` succeeds, this returns the closure result together with the span
603 /// covering the consumed source range.
604 pub fn try_with_span<F, R>(&mut self, f: F) -> Result<(R, Span), PtxParseError>
605 where
606 F: FnOnce(&mut PtxTokenStream) -> Result<R, PtxParseError>,
607 {
608 let start_pos = self.position();
609 match f(self) {
610 Ok(value) => {
611 let end_pos = self.position();
612 let span_start = self.offset_from_start(start_pos);
613 let span_end = self.offset_from_end(start_pos, end_pos).max(span_start);
614 Ok((value, Span::new(span_start, span_end)))
615 }
616 Err(err) => {
617 self.set_position(start_pos);
618 Err(err)
619 }
620 }
621 }
622
623 /// Get the current position in the stream, for backtracking.
624 ///
625 /// # Behavior for complete mode
626 ///
627 /// Returns a StreamPosition containing the token index (index.0).
628 /// The char offset (index.1) will be `None`.
629 ///
630 /// # Behavior for partial mode
631 ///
632 /// Returns a StreamPosition containing both the token index (index.0) and
633 /// the character offset within that token (index.1 = Some(offset)).
634 ///
635 /// This position can be used with `set_position()` to restore the exact state,
636 /// including the parsing mode and character offset.
637 pub fn position(&self) -> StreamPosition {
638 self.index
639 }
640
641 /// Reset the stream to a previously saved position, for backtracking.
642 ///
643 /// # Behavior for complete mode
644 ///
645 /// Restores the token index to the saved position. If the saved position
646 /// was in complete mode (char offset = None), stays in complete mode.
647 ///
648 /// # Behavior for partial mode
649 ///
650 /// Can restore to either complete or partial mode depending on the saved position.
651 /// If the saved position was in partial mode (char offset = Some(n)), switches
652 /// to partial mode at that exact character offset. This allows proper backtracking
653 /// during character-by-character matching attempts.
654 pub fn set_position(&mut self, pos: StreamPosition) {
655 self.index = pos;
656 }
657
658 /// Check if we've reached the end of the token stream.
659 ///
660 /// # Behavior for complete mode
661 ///
662 /// Returns `true` if the token index is at or past the end of the tokens array
663 /// and we're in complete mode (char offset is `None`).
664 ///
665 /// # Behavior for partial mode
666 ///
667 /// Always returns `false` while in partial mode (char offset is `Some`), even if
668 /// positioned at the last token. This is because partial mode implies we're still
669 /// potentially consuming characters from the current token.
670 pub fn is_at_end(&self) -> bool {
671 self.index.0 >= self.tokens.len() && self.index.1.is_none()
672 }
673
674 /// Create a zero-length span at the current stream position.
675 pub fn current_span(&self) -> Span {
676 let offset = self.offset_from_start(self.index);
677 Span::new(offset, offset)
678 }
679
680 /// Convert a `StreamPosition` into an absolute start offset in source bytes.
681 ///
682 /// Uses the lexer-supplied span of the token at `pos.0` and the character
683 /// offset stored in `pos.1` (if any) to compute the precise byte position,
684 /// preserving partial-mode progress within the token.
685 fn offset_from_start(&self, pos: StreamPosition) -> usize {
686 if let Some((_, span)) = self.tokens.get(pos.0) {
687 let token_offset = pos.1.unwrap_or(0);
688 return (span.start + token_offset).min(span.end);
689 }
690 self.tokens.last().map(|(_, span)| span.end).unwrap_or(0)
691 }
692
693 /// Convert a pair of positions into the absolute end offset of the parsed span.
694 ///
695 /// Handles both complete mode (token-level) and partial mode (character-level)
696 /// states and gracefully falls back to the closest known span when the stream
697 /// is at the very beginning or end.
698 fn offset_from_end(&self, start: StreamPosition, end: StreamPosition) -> usize {
699 if start == end {
700 return self.offset_from_start(start);
701 }
702
703 if let Some(char_offset) = end.1 {
704 if let Some((_, span)) = self.tokens.get(end.0) {
705 return (span.start + char_offset).min(span.end);
706 }
707 } else if end.0 == 0 {
708 if let Some((_, span)) = self.tokens.get(0) {
709 return span.start;
710 }
711 } else if let Some((_, span)) = self.tokens.get(end.0 - 1) {
712 return span.end;
713 }
714
715 self.tokens
716 .last()
717 .map(|(_, span)| span.end)
718 .unwrap_or_else(|| self.offset_from_start(start))
719 }
720}
721
722/// Trait for types that can be parsed from a PTX token stream.
723///
724/// This trait is implemented for all PTX AST node types to enable
725/// recursive descent parsing.
726///
727/// Following the combinator architecture, parse() returns a parser function
728/// rather than directly taking a stream parameter.
729pub trait PtxParser
730where
731 Self: Sized,
732{
733 /// Returns a parser function that can parse an instance of `Self`.
734 fn parse() -> impl Fn(&mut PtxTokenStream) -> Result<(Self, Span), PtxParseError>;
735}
736
737// Parse PTX source code into a structured Module representation.
738//
739// This is the main entry point for parsing PTX code. It performs lexical
740// analysis followed by syntactic parsing.
741//
742// # Arguments
743//
744// * `source` - The PTX source code as a string slice
745//
746// # Returns
747//
748// Returns a parsed `Module` AST node, or a `PtxParseError` if parsing fails.
749//
750// # Example
751//
752// ```no_run
753// use ptx_parser::parse_ptx;
754//
755// let source = r#"
756// .version 8.5
757// .target sm_90
758// .address_size 64
759//
760// .entry kernel() {
761// ret;
762// }
763// "#;
764//
765// let module = parse_ptx(source).expect("Failed to parse PTX");
766// println!("Parsed {} directives", module.directives.len());
767// ```
768pub fn parse_ptx(source: &str) -> Result<crate::r#type::module::Module, PtxParseError> {
769 #[cfg(debug_assertions)]
770 {
771 // Debug builds can have very deep combinator stacks; force a large stack for parsing.
772 return stacker::grow(256 * 1024 * 1024, || parse_ptx_inner(source));
773 }
774
775 #[cfg(not(debug_assertions))]
776 {
777 parse_ptx_inner(source)
778 }
779}
780
781fn parse_ptx_inner(source: &str) -> Result<crate::r#type::module::Module, PtxParseError> {
782 use crate::{PtxTokenStream, tokenize, r#type::Module};
783
784 let tokens = tokenize(source)?;
785 let mut stream = PtxTokenStream::new(&tokens);
786 let (module, _) = Module::parse()(&mut stream)?;
787 if !stream.is_at_end() {
788 let pos = stream.position();
789 let remaining = tokens
790 .get(pos.0)
791 .map(|(tok, _)| format!("{:?}", tok))
792 .unwrap_or_else(|| "EOF".into());
793 return Err(PtxParseError {
794 kind: ParseErrorKind::UnexpectedToken {
795 expected: vec!["end of file".into()],
796 found: remaining,
797 },
798 span: stream.current_span(),
799 });
800 }
801 Ok(module)
802}