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