sipha_parse/state/

mod.rs

//! Parser state and node builder utilities.

mod builder;
mod recovery;

pub use builder::NodeBuilder;
pub use recovery::RecoveryStrategy;

use sipha_core::{
    span::Span,
    token::Token,
    token::TokenTrivia,
    traits::{GrammarContext, NodeId, RuleId, TokenKind},
};
use sipha_error::{ErrorContext, Expected, ParseError};
use sipha_memo::{MemoEntry, MemoTable};
use sipha_tree::{CstKind, CstNode, NodeArena, NodeChildren};

/// Stateful helper passed around by grammar routines.
pub struct ParserState<'tokens, 'arena, K, R, N, Ctx = ()>
where
    K: TokenKind,
    R: RuleId,
    N: NodeId,
    Ctx: GrammarContext,
{
    tokens: &'tokens [Token<K>],
    pos: usize,
    pub(crate) arena: &'arena mut NodeArena<K, R, N>,
    include_trivia: bool,
    pub context: Ctx,
    /// Stack of rule IDs currently being parsed, for error context.
    rule_stack: Vec<R>,
    /// Memoization table for packrat parsing.
    pub(crate) memo: Option<&'arena mut MemoTable<K, R, N>>,
    /// Precomputed index of non-trivia token positions (for optimization).
    /// This is `None` if `include_trivia` is true or if not yet computed.
    non_trivia_index: Option<Vec<usize>>,
}

impl<'tokens, 'arena, K, R, N, Ctx> ParserState<'tokens, 'arena, K, R, N, Ctx>
where
    K: TokenKind,
    R: RuleId,
    N: NodeId,
    Ctx: GrammarContext,
{
    /// Construct a new parser state.
    pub fn new(
        tokens: &'tokens [Token<K>],
        arena: &'arena mut NodeArena<K, R, N>,
        include_trivia: bool,
        context: Ctx,
    ) -> Self {
        Self::with_memo(tokens, arena, include_trivia, context, None)
    }

    /// Construct a new parser state with memoization support.
    pub fn with_memo(
        tokens: &'tokens [Token<K>],
        arena: &'arena mut NodeArena<K, R, N>,
        include_trivia: bool,
        context: Ctx,
        memo: Option<&'arena mut MemoTable<K, R, N>>,
    ) -> Self {
        let mut state = Self {
            tokens,
            pos: 0,
            arena,
            include_trivia,
            context,
            rule_stack: Vec::new(),
            memo,
            non_trivia_index: None,
        };
        if !include_trivia {
            // Precompute non-trivia index for faster skipping
            state.non_trivia_index = Some(state.build_non_trivia_index());
            state.skip_trivia();
        }
        state
    }

    /// Skip trivia tokens at the current cursor.
    ///
    /// This method uses the precomputed non-trivia index when available
    /// for better performance on large token streams.
    pub fn skip_trivia(&mut self) {
        if let Some(ref index) = self.non_trivia_index {
            // Use binary search to find next non-trivia position
            match index.binary_search(&self.pos) {
                Ok(_) => {
                    // Exact match - already at non-trivia position
                    return;
                }
                Err(insert_pos) => {
                    if insert_pos < index.len() {
                        // Found next non-trivia position
                        self.pos = index[insert_pos];
                        return;
                    }
                    // Past end of index, set to end
                    self.pos = self.tokens.len();
                    return;
                }
            }
        }
        // Fallback to linear search (when include_trivia is true or index not built)
        while self.pos < self.tokens.len() && self.tokens[self.pos].kind.is_trivia() {
            self.pos += 1;
        }
    }

    /// Build an index of non-trivia token positions.
    ///
    /// This is used to optimize trivia skipping for large token streams.
    fn build_non_trivia_index(&self) -> Vec<usize> {
        self.tokens
            .iter()
            .enumerate()
            .filter(|(_, token)| !token.kind.is_trivia())
            .map(|(idx, _)| idx)
            .collect()
    }

    /// Return a reference to a token `offset` positions away.
    pub fn peek(&self, offset: usize) -> Option<&Token<K>> {
        self.tokens.get(self.pos + offset)
    }

    /// Advance the cursor by one token and return the new token.
    pub fn advance(&mut self) -> Option<&Token<K>> {
        if self.pos < self.tokens.len() {
            self.pos += 1;
        }
        if !self.include_trivia {
            self.skip_trivia();
        }
        self.peek(0)
    }

    /// Consume the current token.
    pub fn consume(&mut self) -> Option<&'tokens Token<K>> {
        let pos = self.pos;
        if pos >= self.tokens.len() {
            return None;
        }
        self.advance();
        self.tokens.get(pos)
    }

    /// Whether the cursor points at the provided kind.
    pub fn at(&self, kind: K) -> bool {
        self.peek(0).map(|t| t.kind == kind).unwrap_or(false)
    }

    /// Expect a specific token kind.
    pub fn expect(&mut self, kind: K) -> Result<&'tokens Token<K>, ParseError<K, R>> {
        if self.at(kind) {
            Ok(self.consume().expect("token should be present"))
        } else {
            Err(ParseError::Expected {
                expected: Expected::Single(kind),
                found: self.peek(0).map(|t| t.kind),
                span: self.current_span(),
                rule_context: self.current_rule(),
                context_stack: self.build_context_stack(),
            })
        }
    }

    /// Expect a token and wrap it inside a CST node.
    pub fn expect_node(&mut self, kind: K) -> Result<N, ParseError<K, R>> {
        let token = self.expect(kind)?;
        Ok(self.alloc_token_node(token))
    }

    /// Current source span or a default span if the stream is empty.
    pub fn current_span(&self) -> Span {
        self.peek(0)
            .map(|token| token.span.clone())
            .or_else(|| self.tokens.last().map(|token| token.span.clone()))
            .unwrap_or_default()
    }

    /// Reset the cursor back to the provided position.
    pub fn reset(&mut self, pos: usize) {
        self.pos = pos.min(self.tokens.len());
    }

    /// Current cursor position.
    pub fn position(&self) -> usize {
        self.pos
    }

    /// Get the current token (if any).
    ///
    /// Returns `None` if at end of input.
    pub fn current_token(&self) -> Option<&Token<K>> {
        self.peek(0)
    }

    /// Get the number of remaining tokens.
    ///
    /// This counts all tokens (including trivia if `include_trivia` is true).
    pub fn remaining_tokens(&self) -> usize {
        self.tokens.len().saturating_sub(self.pos)
    }

    /// Check if we're at the end of input.
    pub fn is_eof(&self) -> bool {
        self.pos >= self.tokens.len()
    }

    /// Get the rule stack for debugging/inspection.
    ///
    /// Returns a slice of rule IDs currently being parsed.
    pub fn rule_stack(&self) -> &[R] {
        &self.rule_stack
    }

    /// Get the number of nodes in the arena.
    pub fn arena_size(&self) -> usize {
        self.arena.len()
    }

    /// Start building a new CST node for the provided rule.
    pub fn start_node(&self, rule_id: R) -> NodeBuilder<R, N> {
        NodeBuilder::new(rule_id, self.current_span())
    }

    /// Allocate a node inside the arena.
    pub fn alloc_node(&mut self, kind: CstKind<K, R>, span: Span, children: NodeChildren<N>) -> N {
        self.arena.alloc(CstNode::create(kind, span, children))
    }

    fn alloc_rule_node(&mut self, rule_id: R, span: Span, children: NodeChildren<N>) -> N {
        self.alloc_node(CstKind::Rule(rule_id), span, children)
    }

    fn alloc_trivia_node(&mut self, trivia: &TokenTrivia<K>) -> N {
        self.alloc_node(
            CstKind::Trivia(trivia.kind),
            trivia.span.clone(),
            NodeChildren::new(),
        )
    }

    fn alloc_token_node(&mut self, token: &Token<K>) -> N {
        let mut children = NodeChildren::new();
        for trivia in &token.leading_trivia {
            children.push(self.alloc_trivia_node(trivia));
        }

        for trivia in &token.trailing_trivia {
            children.push(self.alloc_trivia_node(trivia));
        }

        self.alloc_node(CstKind::Token(token.kind), token.span.clone(), children)
    }

    /// Push a rule ID onto the rule stack for error context tracking.
    pub(crate) fn push_rule(&mut self, rule_id: R) {
        self.rule_stack.push(rule_id);
    }

    /// Pop a rule ID from the rule stack.
    pub(crate) fn pop_rule(&mut self) {
        self.rule_stack.pop();
    }

    /// Get the current rule context (top of the stack).
    pub fn current_rule(&self) -> Option<R> {
        self.rule_stack.last().copied()
    }

    /// Build an error context stack from the current rule stack.
    pub fn build_context_stack(&self) -> Vec<ErrorContext<R>> {
        self.rule_stack
            .iter()
            .map(|&rule_id| ErrorContext {
                rule_id,
                span: self.current_span(),
                message: format!("While parsing {:?}", rule_id),
            })
            .collect()
    }

    /// Check the memo table for a cached result at the current position.
    pub(crate) fn check_memo(&self, rule_id: R) -> Option<&MemoEntry<K, R, N>> {
        self.memo
            .as_ref()
            .and_then(|memo| memo.get(rule_id, self.pos))
    }

    /// Store a successful parse result in the memo table.
    pub(crate) fn store_memo_success(
        &mut self,
        rule_id: R,
        position: usize,
        node: N,
        consumed: usize,
    ) {
        if let Some(memo) = self.memo.as_mut() {
            memo.store_success(rule_id, position, node, consumed);
        }
    }

    /// Store a failed parse result in the memo table.
    pub(crate) fn store_memo_failure(
        &mut self,
        rule_id: R,
        position: usize,
        error: ParseError<K, R>,
        consumed: usize,
    ) {
        if let Some(memo) = self.memo.as_mut() {
            memo.store_failure(rule_id, position, error, consumed);
        }
    }
}