ptx_parser/parser/
mod.rs

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