laburnum 1.17.0

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

//! Generic trivia handling for lexers.
//!
//! Trivia represents whitespace and formatting characters that are preserved
//! for accurate source representation but don't affect parsing semantics.

use {
    chumsky::prelude::*,
    core::fmt,
};

use crate::{Span, SpanCache, chumsky::LaburnumSpanExt};

// -- TriviaLexer Trait --------------------------------------------------------

/// Trait alias for trivia parsers.
///
/// Trivia parsers operate on string input and produce `Trivia` output.
/// Generic over input type `I` (must implement `StrInput`) and error type `E`.
pub trait TriviaLexer<'src, I, E, X>:
    chumsky::Parser<'src, I, Trivia, X>
    + Clone
    + 'src
where
    I: chumsky::input::Input<'src, Token = char, Span = chumsky::span::SimpleSpan>
        + chumsky::input::StrInput<'src>,
    E: chumsky::error::Error<'src, I>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<char>>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<()>>
        + 'src,
    X: chumsky::extra::ParserExtra<'src, I>,
{
}

impl<'src, I, E, X, P> TriviaLexer<'src, I, E, X> for P
where
    I: chumsky::input::Input<'src, Token = char, Span = chumsky::span::SimpleSpan>
        + chumsky::input::StrInput<'src>,
    E: chumsky::error::Error<'src, I>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<char>>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<()>>
        + 'src,
    X: chumsky::extra::ParserExtra<'src, I>,
    P: chumsky::Parser<'src, I, Trivia, X>
        + Clone
        + 'src,
{
}

// -- Trivia -------------------------------------------------------------------

/// Trivia represents whitespace and formatting characters.
///
/// Each variant captures its span for accurate source reconstruction.
/// The `Many` variant allows combining multiple trivia elements into a sequence.
#[derive(Clone, PartialEq, Eq, Hash)]
pub enum Trivia {
    /// One or more newline characters (`\n`, `\r`, or `\r\n`).
    Newline(Span),
    /// One or more space characters.
    Space(Span),
    /// One or more tab characters (`\t`).
    Tab(Span),
    /// A sequence of trivia elements combined together.
    Many(Vec<Trivia>, Span),
}

impl Trivia {
    /// Get the span for this trivia.
    pub fn span(&self) -> Span {
        match self {
            | Trivia::Newline(span) => *span,
            | Trivia::Space(span) => *span,
            | Trivia::Tab(span) => *span,
            | Trivia::Many(_, span) => *span,
        }
    }

    /// Convert this trivia to its string representation.
    pub fn as_str(&self) -> String {
        match self {
            | Trivia::Newline(_span) => "\\n".to_string(),
            | Trivia::Space(_span) => " ".to_string(),
            | Trivia::Tab(_span) => "\\t".to_string(),
            | Trivia::Many(t, _span) => {
                let mut s = String::new();
                for t in t {
                    s.push_str(t.as_str().as_str());
                }
                s
            },
        }
    }

    /// Convert this trivia to a span and identifier pair.
    pub fn as_span_ident(self) -> (Span, crate::Ident) {
        match self {
            | Trivia::Newline(span) => (span, crate::Ident::new("\n")),
            | Trivia::Space(span) => (span, crate::Ident::new(" ")),
            | Trivia::Tab(span) => (span, crate::Ident::new("\t")),
            | Trivia::Many(t, span) => (
                span,
                crate::Ident::new(
                    &t.into_iter()
                        .map(|t| t.as_str())
                        .collect::<Vec<_>>()
                        .join(""),
                ),
            ),
        }
    }

    /// Check if this trivia contains any newline characters.
    ///
    /// For `Many` variants, recursively checks all contained trivia.
    pub fn contains_new_line(&self) -> bool {
        match self {
            | Trivia::Newline(_) => true,
            | Trivia::Space(_) => false,
            | Trivia::Tab(_) => false,
            | Trivia::Many(trivia, _) => trivia.iter().any(|t| t.contains_new_line()),
        }
    }

}

impl fmt::Debug for Trivia {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            | Trivia::Newline(span) => f
                .debug_struct("Trivia::Newline")
                .field("span", span)
                .finish(),
            | Trivia::Space(span) => {
                f.debug_struct("Trivia::Space").field("span", span).finish()
            },
            | Trivia::Tab(span) => {
                f.debug_struct("Trivia::Tab").field("span", span).finish()
            },
            | Trivia::Many(vals, _) => f.debug_list().entries(vals).finish(),
        }
    }
}

impl std::fmt::Display for Trivia {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            | Self::Newline(_) => write!(f, "\\n"),
            | Self::Space(_) => write!(f, "' '"),
            | Self::Tab(_) => write!(f, "\\t"),
            | Self::Many(trivia, _) => {
                for trivia in trivia {
                    write!(f, "{trivia}")?;
                }
                Ok(())
            },
        }
    }
}

impl bluegum::Bluegum for Trivia {
    fn node(&self, b: &mut bluegum::Builder) {
        b.name("Trivia").field("kind", "").alt(format!("{self:?}"));
    }
}

// -- Trivia Parsers -----------------------------------------------------------

/// Parser for the plain space character.
///
/// Matches one or more consecutive space characters and combines them into
/// a single `Trivia::Space` with the combined span.
pub fn space<'src, I, E, X>() -> impl TriviaLexer<'src, I, E, X>
where
    I: chumsky::input::Input<'src, Token = char, Span = chumsky::span::SimpleSpan>
        + chumsky::input::StrInput<'src>,
    E: chumsky::error::Error<'src, I>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<char>>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<()>>
        + 'src,
    X: chumsky::extra::ParserExtra<'src, I, Error = E> + 'src,
    X::State: crate::chumsky::SpanCreator,
{
    just(' ')
        .ignored()
        .repeated()
        .at_least(1)
        .count()
        .map_with(|_, e| Trivia::Space(e.create_span()))
        .boxed()
}

/// Parsing tab characters.
///
/// Matches one or more consecutive tab characters and combines them into
/// a single `Trivia::Tab` with the combined span.
pub fn tab<'src, I, E, X>() -> impl TriviaLexer<'src, I, E, X>
where
    I: chumsky::input::Input<'src, Token = char, Span = chumsky::span::SimpleSpan>
        + chumsky::input::StrInput<'src>,
    E: chumsky::error::Error<'src, I>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<char>>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<()>>
        + 'src,
    X: chumsky::extra::ParserExtra<'src, I, Error = E> + 'src,
    X::State: crate::chumsky::SpanCreator,
{
    just('\t')
        .ignored()
        .repeated()
        .at_least(1)
        .count()
        .map_with(|_, e| Trivia::Tab(e.create_span()))
        .boxed()
}

/// Newlines are either `\n`, `\r`, or `\r\n`.
///
/// Matches one or more consecutive newline sequences and combines them into
/// a single `Trivia::Newline` with the combined span.
pub fn newlines<'src, I, E, X>() -> impl TriviaLexer<'src, I, E, X>
where
    I: chumsky::input::Input<'src, Token = char, Span = chumsky::span::SimpleSpan>
        + chumsky::input::StrInput<'src>,
    E: chumsky::error::Error<'src, I>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<char>>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<()>>
        + 'src,
    X: chumsky::extra::ParserExtra<'src, I, Error = E> + 'src,
    X::State: crate::chumsky::SpanCreator,
{
    chumsky::text::newline()
        .ignored()
        .repeated()
        .at_least(1)
        .map_with(|_, e| Trivia::Newline(e.create_span()))
        .boxed()
}

/// Match a single newline character.
///
/// Matches a single newline sequence (`\r\n`, `\n`, or `\r`).
pub fn newline<'src, I, E, X>() -> impl TriviaLexer<'src, I, E, X>
where
    I: chumsky::input::Input<'src, Token = char, Span = chumsky::span::SimpleSpan>
        + chumsky::input::StrInput<'src>,
    E: chumsky::error::Error<'src, I>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<char>>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<()>>
        + 'src,
    X: chumsky::extra::ParserExtra<'src, I, Error = E> + 'src,
    X::State: crate::chumsky::SpanCreator,
{
    chumsky::text::newline()
        .ignored()
        .map_with(|_, e| Trivia::Newline(e.create_span()))
        .boxed()
}

/// Leading trivia is any number of newlines, spaces, or tabs.
///
/// This is typically used before a token to capture whitespace that precedes it.
pub fn leading<'src, I, E, X>() -> impl TriviaLexer<'src, I, E, X>
where
    I: chumsky::input::Input<'src, Token = char, Span = chumsky::span::SimpleSpan>
        + chumsky::input::StrInput<'src>,
    E: chumsky::error::Error<'src, I>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<char>>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<()>>
        + 'src,
    X: chumsky::extra::ParserExtra<'src, I, Error = E> + 'src,
    X::State: crate::chumsky::SpanCreator,
{
    choice((space(), tab(), newline()))
        .repeated()
        .at_least(1)
        .collect::<Vec<_>>()
        .map_with(|t, e| Trivia::Many(t, e.create_span()))
        .boxed()
}

/// Trailing trivia is any number of spaces or tabs, optionally followed by a newline.
///
/// This is typically used after a token to capture whitespace that follows it.
/// The trailing newline is included if present.
pub fn trailing<'src, I, E, X>() -> impl TriviaLexer<'src, I, E, X>
where
    I: chumsky::input::Input<'src, Token = char, Span = chumsky::span::SimpleSpan>
        + chumsky::input::StrInput<'src>,
    E: chumsky::error::Error<'src, I>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<char>>
        + chumsky::error::LabelError<'src, I, chumsky::text::TextExpected<()>>
        + 'src,
    X: chumsky::extra::ParserExtra<'src, I, Error = E> + 'src,
    X::State: crate::chumsky::SpanCreator,
{
    choice((tab(), space(), newline()))
        .repeated()
        .at_least(1)
        .collect::<Vec<_>>()
        .then(newline().or_not())
        .map_with(|(mut t, ln), e| {
            if let Some(ln) = ln {
                t.push(ln)
            }
            Trivia::Many(t, e.create_span())
        })
        .boxed()
}

// -- Bluegum Printer ----------------------------------------------------------

/// Bluegum printer helper for leading and trailing trivia.
///
/// This function renders trivia nodes with a prefix indicating whether they
/// are leading (`"<- "`) or trailing (`"-> "`).
pub fn print_trivia_bluegum(
    t: &Trivia,
    b: &mut bluegum::Builder,
    prefix: &str, // usually "<- " or "-> "
    span_cache: Option<&SpanCache>,
) {
    // keep the text of `field` under 5 char, to match `:: Value`, so that
    // everything lines up nicely
    match t {
        | Trivia::Newline(span) => {
            let field = b
                .field(&format!("{prefix}NewLn"), "")
                .alt(format!("{t}"))
                .debug("span", format!("{:?}", span));
            if let Some(cache) = span_cache {
                let len = span.data(cache).map(|d| d.len).unwrap_or(0);
                field.debug("count", format!("{}", len));
            }
        },
        | Trivia::Space(span) => {
            let field = b
                .field(&format!("{prefix}Space"), "")
                .alt(format!("{t}"))
                .debug("span", format!("{:?}", span));
            if let Some(cache) = span_cache {
                let len = span.data(cache).map(|d| d.len).unwrap_or(0);
                field.debug("count", format!("{}", len));
            }
        },
        | Trivia::Tab(span) => {
            let field = b
                .field(&format!("{prefix}Tab"), "")
                .alt(format!("{t}"))
                .debug("span", format!("{:?}", span));
            if let Some(cache) = span_cache {
                let len = span.data(cache).map(|d| d.len).unwrap_or(0);
                field.debug("count", format!("{}", len));
            }
        },
        | Trivia::Many(vals, _) => {
            for t in vals.iter() {
                print_trivia_bluegum(t, b, prefix, span_cache);
            }
        },
    };
}

// -- Wrap Macro ---------------------------------------------------------------

/// Wraps a token parser with optional leading/trailing trivia.
///
/// This macro simplifies the common pattern of parsing a token with optional
/// whitespace before and after it. The resulting parser produces
/// `Spanned<Token, SimpleSpan>` suitable for downstream CST parsers.
///
/// # Usage
///
/// ```ignore
/// wrap!(
///     { your_token_parser } -> |((leading, inner), trailing), e| {
///         // Map to your token type
///         YourToken::Variant(leading, inner, trailing)
///     }
/// )
/// ```
///
/// The `leading` and `trailing` values are `Option<Trivia>`.
/// The `inner` is whatever your token parser produces.
/// The mapping function should return your Token type.
#[macro_export]
macro_rules! wrap {
    (
        $token_parser:block -> $map:expr
    ) => {
        // leading
        $crate::chumsky::lexer::trivia::leading().or_not()
            // token
            .then($token_parser)
            // trailing
            .then($crate::chumsky::lexer::trivia::trailing().or_not())
            .map_with($map)
            // Wrap in chumsky::span::Spanned<Token> with SimpleSpan for downstream CST parser
            .spanned()
            .boxed()
    };
}

pub use wrap;