orql 0.1.0

use std::{fmt::Debug, num::NonZeroU32};

use super::Location;

use crate::scanner::{Comment as ScannerComment, CommentStyle, Token, TokenType};

/// Aids the parser in tracking metadata to be associated with AST nodes; in
/// particular, the tracker is designed to collect location and comments
/// associated the nodes.
///
/// Upon starting to parse a particular AST node, the parser obtains a "node
/// id" `start_node` and stores it in the AST; clients can then use it to
/// refer to metadata corresponding to the particular AST node. Upon parsing
/// the end, the parser calls `end_last_node` to signal the end of the lastly
/// opened node giving the track a chance to associated so far collected
/// information with the processed node.
///
/// Note: this is intended for use by the parser only.
pub(super) trait Tracker<'s> {
    /// the type of node identifiers
    type NodeId: Copy + Ord + Debug;

    /// the type of the final output structure
    type Metadata: Metadata<'s, NodeId = Self::NodeId>;

    /// Allocates a new node id (returning it) in the assumption that a new
    /// AST node is being started.
    fn on_node_start(&mut self, loc: Location) -> Self::NodeId;

    /// Given a chance to see if `t` is a comment and returns `true` if so.
    /// The parser can then skip the processing of the comment transparently.
    fn accept_comment(&mut self, t: &Token<'s>) -> bool {
        matches!(t.ttype, TokenType::Comment(_))
    }

    /// Explicitly ends a node, flushing so far collected metadata.
    fn on_node_end(&mut self);

    /// Like [Self::on_node_end] but restores `node` as the last node.
    fn on_node_end_restore(&mut self, node: Self::NodeId);

    /// Allows temporary access to so far tracked metadata.
    fn metadata(&self) -> &impl Metadata<'s, NodeId = Self::NodeId>;

    /// Finishes the tracking and turns it into the final metadata structure
    /// for querying purposes.
    fn finish(self) -> Self::Metadata;
}

// ----------------------------------------------------------------------------

/// Provides access to an AST node's metadata.
pub trait Metadata<'s> {
    type NodeId: Copy + Ord;

    /// Retrieves the start location (in the originally parsed source) of a
    /// given AST node.
    fn location(&self, node: Self::NodeId) -> Location;

    /// Retrieves leading and trailing comments associated with a node,
    /// i.e. comments appearning before and after a particular node.
    ///
    /// Note: a comment may be trailing to one node and, at the same time, a
    /// leading comment to another node. Example: `a /* hello */ + b` where
    /// `/* hello */` is trailing to `a` and leading to the `+` operator. You
    /// can use the [comment's location](Comment::loc) to disambiguate the
    /// situation.
    fn comments(&self, node: Self::NodeId) -> (&[Comment<'s>], &[Comment<'s>]);
}

/// A source code comment with its location, e.g. `-- ...` or `/* ... */`
#[derive(Debug, PartialEq, Eq)]
pub struct Comment<'s> {
    /// The verbatim content of the comment excluding its comment markers
    ///
    /// Note: single line comments do _not_ contain their terminating newline
    pub text: &'s str,

    /// (single) line (i.e. `-- ...`) or block (multi-line) (i.e. `/* ... */`)
    pub style: CommentStyle,

    /// The start location of the comment (including its marker prefix / beginning)
    pub loc: Location,
}

// ----------------------------------------------------------------------------

/// An opaque node identifier of a node within a single AST.
///
/// The parser may associate additional metadata, like locations and/or
/// comments, with nodes which can be accessed through the nodes' `Id`.
///
/// See [`parse_with_metadata`](crate::parser::parse_with_metadata)
#[derive(PartialEq, Eq, Copy, Clone, PartialOrd, Ord)]
pub struct Id(Location);

impl Debug for Id {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Id({}_{})", self.0.line, self.0.col)
    }
}

/// Collected metadata about nodes within a particular AST.
///
/// Note: the instance and the node [Id]s being queried must originate from
/// the same parse operation. See [`Parser::parse_str_with_metadata`](crate::parser::Parser::parse_with_metadata).
#[derive(Debug)]
pub struct DefaultMetadata<'s>(CommentsAccum<'s>);

impl<'s> Metadata<'s> for DefaultMetadata<'s> {
    type NodeId = Id;

    fn location(&self, node: Self::NodeId) -> Location {
        self.0.location(node)
    }

    fn comments(&self, node: Self::NodeId) -> (&[Comment<'s>], &[Comment<'s>]) {
        self.0.comments(node)
    }
}

#[derive(Default)]
pub(super) struct DefaultTracker<'s> {
    /// the last node - to be associated with encountered comments as its
    /// "trailing comments"
    last_node: Option<Id>,

    /// encountered comments since `last_node`; we collect and emit them upon
    /// encountering the next `node` - at which point that next node becomes
    /// the `last_node` (and the vec of collected comment gets cleared.)
    last_comments: Vec<Comment<'s>>,

    /// all tracked comments accumulated
    final_comments: CommentsAccum<'s>,
}

impl<'s> Tracker<'s> for DefaultTracker<'s> {
    type NodeId = Id;
    type Metadata = DefaultMetadata<'s>;

    fn accept_comment(&mut self, t: &Token<'s>) -> bool {
        match t.ttype {
            TokenType::Comment(ScannerComment(text, style)) => {
                self.last_comments.push(Comment {
                    text,
                    style,
                    loc: t.loc,
                });
                true
            }
            _ => false,
        }
    }

    /// Allocates a new node id (returning it) in the assumption that a new
    /// AST node has been encountered. Emits the last collected comments associating
    /// with the last node (if any) and the new one.
    fn on_node_start(&mut self, loc: Location) -> Self::NodeId {
        // ~ we simply use the location (of a start token) as a unique of the
        // ast node being built
        let new_node = self.final_comments.new_node(loc);

        if !self.last_comments.is_empty() {
            let comments = self.last_comments.drain(..);
            if let Some(last_node) = self.last_node {
                self.final_comments
                    .track_comments_between(comments, last_node, new_node);
            } else {
                self.final_comments
                    .track_comments_before(comments, new_node);
            }
        }
        self.last_node = Some(new_node);
        new_node
    }

    /// Explicitly ends a node, flushing so far collected comments as the last
    /// node's "trailing comments." Further comments encountered after this
    /// method call will _not_ be associated with the last node (only as
    /// "leading comments" of the next node, if any.)
    fn on_node_end(&mut self) {
        if let Some(last_node) = self.last_node.take()
            && !self.last_comments.is_empty()
        {
            let comments = self.last_comments.drain(..);
            self.final_comments
                .track_comments_after(comments, last_node);
        }
    }

    fn on_node_end_restore(&mut self, node: Self::NodeId) {
        self.on_node_end();
        self.last_node = Some(node);
    }

    fn metadata(&self) -> &impl Metadata<'s, NodeId = Self::NodeId> {
        self
    }

    fn finish(mut self) -> Self::Metadata {
        self.on_node_end();
        self.final_comments.finish()
    }
}

#[derive(Default, Debug)]
struct CommentsAccum<'s> {
    /// ~ only ever appended to
    comments: Vec<Comment<'s>>,
    /// ~ keeps only nodes with associated comments;
    /// ~ ordered by [NodeWithComments::node]
    nodes: Vec<NodeWithComments>,
}

#[derive(Debug)]
struct NodeWithComments {
    node: Id,
    leading: WeakRef,
    trailing: WeakRef,
}

#[derive(Debug, Copy, Clone)]
struct WeakRef {
    pos: u32,
    len: u32,
}

impl WeakRef {
    fn sub_slice<T>(self, slice: &[T]) -> &[T] {
        &slice[self.pos as usize..self.pos as usize + self.len as usize]
    }
}

impl<'s> CommentsAccum<'s> {
    fn new_node(&mut self, loc: Location) -> Id {
        Id(loc)
    }

    fn track_comments_after(&mut self, comments: impl Iterator<Item = Comment<'s>>, node: Id) {
        if let Some((num_comments, comments_pos)) = self.push_comments_(comments) {
            self.track_comments_as_trailing_(node, num_comments.get(), comments_pos);
        }
    }

    fn track_comments_before(&mut self, comments: impl Iterator<Item = Comment<'s>>, node: Id) {
        if let Some((num_comments, comments_pos)) = self.push_comments_(comments) {
            self.track_comments_as_leading_(node, num_comments.get(), comments_pos);
        }
    }

    fn track_comments_between(
        &mut self,
        comments: impl Iterator<Item = Comment<'s>>,
        node_before: Id,
        node_after: Id,
    ) {
        let Some((num_comments, comments_pos)) = self.push_comments_(comments) else {
            return;
        };
        self.track_comments_as_trailing_(node_before, num_comments.get(), comments_pos);
        self.track_comments_as_leading_(node_after, num_comments.get(), comments_pos);
    }

    fn finish(self) -> DefaultMetadata<'s> {
        DefaultMetadata(self)
    }

    // impl details -----------------------------------------------------------

    /// Ensures a node to track `num_comments` at `comments_pos` as "trailing",
    /// ie. to be countered _after_ the node location.
    fn track_comments_as_trailing_(&mut self, node: Id, num_comments: u32, comments_pos: u32) {
        debug_assert!(num_comments > 0);
        match self.nodes.binary_search_by(|x| x.node.cmp(&node)) {
            Ok(i) => {
                self.nodes[i].trailing = WeakRef {
                    pos: comments_pos,
                    len: num_comments,
                };
            }
            Err(i) => {
                self.nodes.insert(
                    i,
                    NodeWithComments {
                        node,
                        leading: WeakRef { pos: 0, len: 0 },
                        trailing: WeakRef {
                            pos: comments_pos,
                            len: num_comments,
                        },
                    },
                );
            }
        }
    }

    /// Ensures a node to track `num_comments` at `comment_pos` as "trailing",
    /// ie. to be countered _after_ the node location.
    fn track_comments_as_leading_(&mut self, node: Id, num_comments: u32, comments_pos: u32) {
        debug_assert!(num_comments > 0);
        match self.nodes.binary_search_by(|x| x.node.cmp(&node)) {
            Ok(i) => {
                self.nodes[i].leading = WeakRef {
                    pos: comments_pos,
                    len: num_comments,
                }
            }
            Err(i) => {
                self.nodes.insert(
                    i,
                    NodeWithComments {
                        node,
                        leading: WeakRef {
                            pos: comments_pos,
                            len: num_comments,
                        },
                        trailing: WeakRef { pos: 0, len: 0 },
                    },
                );
            }
        }
    }

    /// Pushes the given comments returning `(push-comments-len, index-of-the-first-pushed-comment)`.
    /// Returns `None` if no comments where pushed.
    fn push_comments_(
        &mut self,
        comments: impl Iterator<Item = Comment<'s>>,
    ) -> Option<(NonZeroU32, u32)> {
        let (mut num_comments, index) = (0, self.comments.len() as u32);
        for comment in comments {
            self.comments.push(comment);
            num_comments += 1;
        }
        if self.comments.len() >= u32::MAX as usize {
            panic!("too many comments");
        }
        if num_comments == 0 {
            None
        } else {
            Some((
                NonZeroU32::new(num_comments).expect("invalid num_comments"),
                index,
            ))
        }
    }
}

impl<'s> Metadata<'s> for CommentsAccum<'s> {
    type NodeId = Id;

    fn location(&self, node: Self::NodeId) -> Location {
        // ~ note: we simply use an AST Node's location is its ID; see `DefaultTracker::new_node()`
        node.0
    }

    fn comments(&self, node: Self::NodeId) -> (&[Comment<'s>], &[Comment<'s>]) {
        match self.nodes.binary_search_by(|x| x.node.cmp(&node)) {
            Ok(i) => {
                let (n, cs) = (&self.nodes[i], self.comments.as_slice());
                (n.leading.sub_slice(cs), n.trailing.sub_slice(cs))
            }
            Err(_) => (&[], &[]),
        }
    }
}

/// `DefaultTracker` can (also) serve not-yet-finalized metadata pretending to
/// associate so far collected, but not yet assigned `last_comments` as
/// trailing with the `last_node`.
impl<'s> Metadata<'s> for DefaultTracker<'s> {
    type NodeId = Id;

    fn location(&self, node: Self::NodeId) -> Location {
        node.0
    }

    fn comments(&self, node: Self::NodeId) -> (&[Comment<'s>], &[Comment<'s>]) {
        let comments = self.final_comments.comments(node);
        if Some(node) == self.last_node && !self.last_comments.is_empty() {
            (comments.0, &self.last_comments)
        } else {
            comments
        }
    }
}

// ----------------------------------------------------------------------------

/// A tracker to expose node locations as node ids.
#[derive(Default)]
pub(super) struct LocationsTracker;

impl<'s> Tracker<'s> for LocationsTracker {
    type NodeId = Location;
    type Metadata = LocationMetadata;
    fn on_node_start(&mut self, loc: Location) -> Self::NodeId {
        loc
    }
    fn on_node_end(&mut self) {}
    fn on_node_end_restore(&mut self, _node: Self::NodeId) {}
    fn metadata(&self) -> &impl Metadata<'s, NodeId = Self::NodeId> {
        &LocationMetadata
    }
    fn finish(self) -> Self::Metadata {
        LocationMetadata
    }
}

pub(super) struct LocationMetadata;

impl<'s> Metadata<'s> for LocationMetadata {
    type NodeId = Location;

    fn location(&self, node: Self::NodeId) -> Location {
        node
    }

    fn comments(&self, _node: Self::NodeId) -> (&[Comment<'s>], &[Comment<'s>]) {
        (&[][..], &[][..])
    }
}

// ----------------------------------------------------------------------------

/// A "void" metadata tracker to ignore all metadata.
#[derive(Default)]
pub(super) struct VoidTracker;

impl<'s> Tracker<'s> for VoidTracker {
    type NodeId = ();
    type Metadata = VoidMetadata;
    fn on_node_start(&mut self, _loc: Location) -> Self::NodeId {}
    fn on_node_end(&mut self) {}
    fn on_node_end_restore(&mut self, _node: Self::NodeId) {}
    fn metadata(&self) -> &impl Metadata<'s, NodeId = Self::NodeId> {
        &VoidMetadata
    }
    fn finish(self) -> Self::Metadata {
        VoidMetadata
    }
}

pub(super) struct VoidMetadata;

impl<'s> Metadata<'s> for VoidMetadata {
    type NodeId = ();

    fn location(&self, _node: Self::NodeId) -> Location {
        Location { line: 0, col: 0 }
    }

    fn comments(&self, _node: Self::NodeId) -> (&[Comment<'s>], &[Comment<'s>]) {
        (&[][..], &[][..])
    }
}
// ----------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use pretty_assertions::assert_eq;

    use super::*;

    impl<'s> DefaultMetadata<'s> {
        fn comments_all(&self) -> &[Comment<'s>] {
            &self.0.comments
        }
    }

    #[test]
    fn test_default_tracker() {
        fn tok_comment<'s>(text: &'s str, loc: Location) -> Token<'s> {
            Token {
                ttype: TokenType::Comment(ScannerComment(text, CommentStyle::Block)),
                loc,
            }
        }
        fn comment<'s>(text: &'s str, loc: Location) -> Comment<'s> {
            Comment {
                text,
                style: CommentStyle::Block,
                loc,
            }
        }

        let mut t = DefaultTracker::default();

        // just like the parser would do, feed the tracker on approx this
        // token sequence: `/* one */ /* two */ 123 + /* three */ 456 /* four */`

        // assert comments get acknoledged correctly
        assert!(t.accept_comment(&tok_comment(" one ", (1, 1).into())));
        assert!(t.accept_comment(&tok_comment(" two ", (1, 5).into())));
        let node_123 = t.on_node_start((1, 15).into()); // ~ value
        // assert non-comments get acknoledged correctly
        assert!(!t.accept_comment(&Token {
            ttype: TokenType::Plus,
            loc: (1, 20).into()
        }));
        let node_plus = t.on_node_start((1, 20).into()); // ~ binary expr
        assert!(t.accept_comment(&tok_comment(" three ", (1, 40).into())));
        let node_456 = t.on_node_start((1, 40).into()); // ~ value
        assert!(t.accept_comment(&tok_comment(" four ", (1, 50).into())));
        let meta = t.finish();

        assert_eq!(
            &[
                comment(" one ", (1, 1).into()),
                comment(" two ", (1, 5).into()),
                comment(" three ", (1, 40).into()),
                comment(" four ", (1, 50).into()),
            ][..],
            meta.comments_all()
        );

        assert_eq!(
            (
                &[
                    comment(" one ", (1, 1).into()),
                    comment(" two ", (1, 5).into())
                ][..],
                &[][..]
            ),
            meta.comments(node_123)
        );

        assert_eq!(
            (&[][..], &[comment(" three ", (1, 40).into())][..]),
            meta.comments(node_plus)
        );

        assert_eq!(
            (
                &[comment(" three ", (1, 40).into())][..],
                &[comment(" four ", (1, 50).into())][..]
            ),
            meta.comments(node_456)
        );
    }
}