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}