grammar.rs

use crate::lexicon::{Cursor, Intern, Interner, Lexeme, Symbol, TokenAlias, TokenList, TokenRef};
use crate::locator::Locator;
use crate::span::Span;
use generational_indextree::NodeEdge;
use sha2::Digest;
use std::cell::{Cell, Ref, RefCell, RefMut};
use std::collections::HashMap;
use std::fmt::{Debug, Display, Formatter};
use std::hash::Hash;

pub type NodeIdx = generational_indextree::NodeId;

/// The trait describing a language grammar.
// Note: we need those bounds on the trait itself to deal with the
// incorrect bounds generated by derive (https://github.com/rust-lang/rust/issues/26925).
pub trait Grammar: Clone + Default + Debug {
    /// The type of lexemes (aka. tokens).
    type Lex: Lexeme;
    /// The type of syntax node kinds.
    type Kind: Copy + Clone + PartialEq + Debug;
    /// The type of parsing expression tags, for memoization.
    type Tag: Copy + Clone + PartialEq + Eq + Hash + Debug;
}

#[derive(Clone)]
pub enum SyntaxTrunk<G: Grammar> {
    Leaf(TokenAlias<G::Lex>),
    Tree(G::Kind),
    Error,
}

impl<G: Grammar> Debug for SyntaxTrunk<G> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            SyntaxTrunk::Leaf(t) => write!(f, "Leaf[{:?}]", t.kind()),
            SyntaxTrunk::Tree(k) => write!(f, "Tree[{k:?}]"),
            SyntaxTrunk::Error => write!(f, "Error"),
        }
    }
}

impl<G: Grammar> Copy for SyntaxTrunk<G> {}

/// A trait for internally mutable node state.
pub trait Core: Default + Clone + Debug {}

impl<C: Default + Clone + Debug> Core for C {}

/// The wrapper type around the syntax node core.
///
/// It is allocated on write as we assume that the number of nodes with a core
/// is small compared to the total number of nodes in a syntax tree.
///
/// An alternative would be to use an arena instead of heap allocations.
type CoreCell<T> = RefCell<Option<Box<T>>>;

#[derive(Clone, Debug)]
pub struct SyntaxNode<T: Core, G: Grammar>(SyntaxTrunk<G>, CoreCell<T>);

impl<T: Core, G: Grammar> SyntaxNode<T, G> {
    pub fn new(trunk: SyntaxTrunk<G>) -> Self {
        SyntaxNode(trunk, RefCell::new(None))
    }

    pub fn trunk(&self) -> &SyntaxTrunk<G> {
        &self.0
    }

    pub fn core_from(&self, other: &Self) {
        self.1.replace(other.1.borrow().clone());
    }

    pub fn has_core(&self) -> bool {
        self.1.borrow().is_some()
    }

    pub fn core_ref(&self) -> Ref<T> {
        Ref::map(self.1.borrow(), |r| {
            r.as_ref().expect("core should exist").as_ref()
        })
    }

    pub fn core_mut(&self) -> RefMut<T> {
        RefMut::map(self.1.borrow_mut(), |r| {
            r.get_or_insert_with(Box::default).as_mut()
        })
    }
}

type TreeArena<T, G> = generational_indextree::Arena<SyntaxNode<T, G>>;

pub struct SyntaxTree<T: Core, G: Grammar> {
    tokens: TokenList<G::Lex>,
    tree: TreeArena<T, G>,
    root: Option<NodeIdx>,
}

impl<T: Core, G: Grammar> Interner for SyntaxTree<T, G> {
    fn register<S: AsRef<str>>(&mut self, s: S) -> Symbol {
        self.tokens.register(s)
    }

    fn resolve(&self, sym: Symbol) -> &str {
        self.tokens.resolve(sym)
    }
}

impl<T: Core, G: Grammar> Debug for SyntaxTree<T, G> {
    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
        if let Some(root) = self.root {
            NodeRef::from(self, root).fmt(f)
        } else {
            write!(f, "[empty tree]")
        }
    }
}

impl<T: Core, G: Grammar> SyntaxTree<T, G> {
    pub fn new(tokens: TokenList<G::Lex>) -> Self {
        SyntaxTree {
            tokens,
            tree: TreeArena::default(),
            root: Option::default(),
        }
    }

    pub fn locator(&self) -> &Locator {
        self.tokens.locator()
    }

    pub fn finalize(mut self, root: NodeIdx) -> Self {
        self.root = Some(root);
        self
    }

    pub fn root(&self) -> NodeRef<T, G> {
        NodeRef::from(self, self.root.unwrap())
    }

    pub fn detach(&self, from: NodeIdx) -> SyntaxTree<T, G> {
        let tokens = TokenList::new(self.tokens.locator().clone());
        let mut tree = SyntaxTree::new(tokens);
        let mut parents: Vec<NodeIdx> = Vec::new();
        let mut root: Option<NodeIdx> = None;

        self.traverse(from).for_each(|edge| {
            match edge {
                NodeEdge::Start(id) => {
                    let node = self.node(id);

                    let new_trunk = match node.trunk() {
                        SyntaxTrunk::Leaf(t) => {
                            let cursor = t.cursor();
                            let (token, span) = self.tokens.token_span(cursor);
                            let new_value = token.value().copy(self, &mut tree);
                            let new_token =
                                <<G as Grammar>::Lex as Lexeme>::new(token.kind(), new_value);
                            let new_cursor = tree.tokens.push(new_token, span.range());
                            SyntaxTrunk::Leaf(tree.tokens.alias(new_cursor))
                        }
                        t => *t,
                    };

                    let new_syntax = SyntaxNode::new(new_trunk);
                    new_syntax.core_from(node);

                    let new_id = tree.new_node(new_syntax);

                    if let Some(parent) = parents.last() {
                        tree.append(*parent, new_id)
                    }

                    parents.push(new_id);
                }
                NodeEdge::End(_) => {
                    root = parents.pop();
                }
            };
        });

        tree.finalize(root.unwrap())
    }

    pub fn count(&self) -> usize {
        self.tree.count()
    }

    fn node(&self, id: NodeIdx) -> &SyntaxNode<T, G> {
        self.tree.get(id).unwrap().get()
    }

    fn children(&self, id: NodeIdx) -> impl Iterator<Item = NodeIdx> + '_ {
        id.children(&self.tree)
    }

    fn reverse_children(&self, id: NodeIdx) -> impl Iterator<Item = NodeIdx> + '_ {
        id.reverse_children(&self.tree)
    }

    fn ancestors(&self, id: NodeIdx) -> impl Iterator<Item = NodeIdx> + '_ {
        id.ancestors(&self.tree)
    }

    fn descendants(&self, id: NodeIdx) -> impl Iterator<Item = NodeIdx> + '_ {
        id.descendants(&self.tree)
    }

    fn traverse(&self, id: NodeIdx) -> impl Iterator<Item = NodeEdge> + '_ {
        id.traverse(&self.tree)
    }

    fn new_node(&mut self, n: SyntaxNode<T, G>) -> NodeIdx {
        self.tree.new_node(n)
    }

    fn append(&mut self, parent: NodeIdx, child: NodeIdx) {
        parent.append(child, &mut self.tree);
    }
}

pub struct NodeRef<'a, T: Core, G: Grammar> {
    tree: &'a SyntaxTree<T, G>,
    idx: NodeIdx,
}

// Note: for some reason the derive macro is not doing the right thing for Clone/Copy.
impl<'a, T: Core, G: Grammar> Clone for NodeRef<'a, T, G> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<'a, T: Core, G: Grammar> Copy for NodeRef<'a, T, G> {}

#[derive(Debug)]
pub enum NodeCursor<'a, T: Core, G: Grammar> {
    Start(NodeRef<'a, T, G>),
    End(NodeRef<'a, T, G>),
}

impl<'a, T: Core, G: Grammar> NodeRef<'a, T, G> {
    pub fn from(tree: &'a SyntaxTree<T, G>, idx: NodeIdx) -> Self {
        NodeRef { tree, idx }
    }

    pub fn tree(&self) -> &'a SyntaxTree<T, G> {
        self.tree
    }

    pub fn index(&self) -> NodeIdx {
        self.idx
    }

    pub fn syntax(&self) -> &'a SyntaxNode<T, G> {
        self.tree.node(self.idx)
    }

    pub fn token(&self) -> TokenRef<'a, G::Lex> {
        match self.syntax().trunk() {
            SyntaxTrunk::Leaf(t) => self.tree.tokens.reference(t.cursor()),
            _ => panic!("a node must be a leaf to link to a token"),
        }
    }

    pub fn len(&self) -> usize {
        self.tree.children(self.idx).count()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn children(&self) -> impl Iterator<Item = NodeRef<'a, T, G>> + 'a {
        self.tree
            .children(self.idx)
            .map(|id| NodeRef::from(self.tree, id))
    }

    pub fn reverse_children(&self) -> impl Iterator<Item = NodeRef<'a, T, G>> + 'a {
        self.tree
            .reverse_children(self.idx)
            .map(|id| NodeRef::from(self.tree, id))
    }

    pub fn first(&self) -> NodeRef<'a, T, G> {
        self.children().next().unwrap()
    }

    pub fn last(&self) -> NodeRef<'a, T, G> {
        self.reverse_children().next().unwrap()
    }

    pub fn nth(&self, n: usize) -> NodeRef<'a, T, G> {
        let Some(node) = self.children().nth(n) else {
            panic!("expected node at {n}")
        };
        node
    }

    pub fn ancestors(&self) -> impl Iterator<Item = NodeRef<'a, T, G>> + 'a {
        self.tree
            .ancestors(self.idx)
            .map(|id| NodeRef::from(self.tree, id))
    }

    pub fn descendants(&self) -> impl Iterator<Item = NodeRef<'a, T, G>> + 'a {
        self.tree
            .descendants(self.idx)
            .map(|id| NodeRef::from(self.tree, id))
    }

    pub fn traverse(&self) -> impl Iterator<Item = NodeCursor<'a, T, G>> + 'a {
        self.tree.traverse(self.idx).map(|edge| match edge {
            NodeEdge::Start(id) => NodeCursor::Start(NodeRef::from(self.tree, id)),
            NodeEdge::End(id) => NodeCursor::End(NodeRef::from(self.tree, id)),
        })
    }

    /// Returns the span of text from first to last token, if any.
    pub fn span(&self) -> Option<Span> {
        if let (Some(start), Some(end)) = (self.start(), self.end()) {
            let s = start.span();
            let e = end.span();
            Some(Span::new(s.locator().clone(), s.start()..e.end()))
        } else {
            None
        }
    }

    /// Returns a reference to the first token, if any.
    pub fn start(&self) -> Option<TokenRef<'a, G::Lex>> {
        match self.syntax().trunk() {
            SyntaxTrunk::Leaf(t) => Some(self.tree.tokens.reference(t.cursor())),
            _ => self.children().find_map(|c| c.start()),
        }
    }

    /// Returns a reference to the last token, if any.
    pub fn end(&self) -> Option<TokenRef<'a, G::Lex>> {
        match self.syntax().trunk() {
            SyntaxTrunk::Leaf(t) => Some(self.tree.tokens.reference(t.cursor())),
            _ => self.reverse_children().find_map(|c| c.end()),
        }
    }

    pub fn detach(&self) -> SyntaxTree<T, G> {
        self.tree.detach(self.idx)
    }

    pub fn as_str(&self) -> &'a str {
        self.token().value().as_str(self.tree)
    }

    /// Update the given digest with node reference information.
    pub fn digest<D: Digest>(&self, digest: &mut D) {
        let (index, generation) = generational_arena::Index::from(self.idx).into_raw_parts();
        digest.update(self.tree.locator().url().as_str());
        digest.update(index.to_be_bytes());
        digest.update(generation.to_be_bytes());
    }
}

impl<'a, T: Core, G: Grammar> Debug for NodeRef<'a, T, G> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{:?}", self.syntax().trunk())?;
        if self.syntax().has_core() {
            write!(f, " ({:?})", self.syntax().core_ref())?;
        }
        if !self.is_empty() {
            write!(f, " -> ")?;
            f.debug_list().entries(self.children()).finish()?;
        }
        Ok(())
    }
}

/// An abstract type over a concrete syntax node.
pub trait AbstractSyntaxNode<'a, T: Core, G: Grammar>
where
    Self: Sized,
{
    fn cast(node: NodeRef<'a, T, G>) -> Option<Self>;
    fn node(&self) -> NodeRef<'a, T, G>;
}

#[macro_export]
macro_rules! syntax_nodes {
    ( $grammar:ident, $( $node:ident ),+ ) => {
        #[allow(dead_code)]
        #[derive(Copy, Clone, PartialEq, Eq, Debug)]
        pub enum SyntaxKind {
            $( $node ),+
        }

        $(
            #[allow(dead_code)]
            #[derive(Clone, Copy, Debug)]
            pub struct $node<'a, T: Core>(NodeRef<'a, T, $grammar>);

            #[allow(dead_code)]
            impl<'a, T: Core> AbstractSyntaxNode<'a, T, $grammar> for $node<'a, T>
            {
                fn cast(node: NodeRef<'a, T, $grammar>) -> Option<Self> {
                    match node.syntax().trunk() {
                        SyntaxTrunk::Tree(SyntaxKind::$node) => Some($node(node)),
                        _ => None,
                    }
                }

                fn node(&self) -> NodeRef<'a, T, $grammar> {
                   self.0
                }
            }
        )+
    }
}

#[macro_export]
macro_rules! terminal_node {
    ( $grammar:ident, $node:ident, $(|)? $( $pattern:pat_param )|+ $( if $guard:expr )? $(,)?  ) => {
        #[allow(dead_code)]
        #[derive(Debug)]
        pub struct $node<'a, T: Core>(NodeRef<'a, T, $grammar>);

        #[allow(dead_code)]
        impl<'a, T: Core> AbstractSyntaxNode<'a, T, $grammar> for $node<'a, T> {
            fn cast(node: NodeRef<'a, T, $grammar>) -> Option<Self> {
                match node.syntax().trunk() {
                    SyntaxTrunk::Leaf(t) if matches!(t.kind(), $( $pattern )|+ $( if $guard )?) => Some(Self(node)),
                    _ => None,
                }
            }

            fn node(&self) -> NodeRef<'a, T, $grammar> {
                self.0
            }
        }
    };
}

#[derive(Debug, Clone)]
pub struct ParserError(&'static str, Span);

impl ParserError {
    pub fn new(error: &'static str, span: Span) -> Self {
        ParserError(error, span)
    }

    pub fn span(&self) -> Span {
        self.1.clone()
    }
}

impl Display for ParserError {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.write_str(self.0)
    }
}

impl std::error::Error for ParserError {}

/// A successful production from a parser function.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParserMatch<G: Grammar> {
    Token(TokenAlias<G::Lex>),
    Syntax(G::Kind),
    Node(NodeIdx),
}

/// A result from a parser function.
pub type ParserResult<G> = std::result::Result<(Cursor, ParserMatch<G>), ParserError>;

/// A syntax analysis context.
pub struct Context<T: Core, G: Grammar> {
    tree: SyntaxTree<T, G>,
    cache: HashMap<(Cursor, G::Tag), ParserResult<G>>,
    hits: Cell<usize>,
    reads: Cell<usize>,
    no_cache: bool,
}

impl<T: Core, G: Grammar> Context<T, G> {
    pub fn new(tokens: TokenList<G::Lex>) -> Self {
        Context {
            tree: SyntaxTree::new(tokens),
            cache: Default::default(),
            hits: Cell::new(0),
            reads: Cell::new(0),
            no_cache: false,
        }
    }

    pub fn without_cache(mut self) -> Self {
        self.no_cache = true;
        self
    }

    pub fn head(&self) -> Cursor {
        self.skip_trivia(self.tree.tokens.head())
    }

    pub fn span(&self, s: Cursor) -> Span {
        if s.is_valid() {
            self.tree.tokens.token_span(s).1
        } else {
            let end = self.tree.tokens.end();
            Span::new(self.tree.tokens.locator().clone(), end..end + 1)
        }
    }

    pub fn tree(self) -> SyntaxTree<T, G> {
        self.tree
    }

    pub fn cache(&mut self, p: G::Tag, s: Cursor, r: ParserResult<G>) {
        if !self.no_cache {
            self.cache.insert((s, p), r);
        }
    }

    pub fn lookup(&mut self, p: G::Tag, s: Cursor) -> Option<ParserResult<G>> {
        if self.no_cache {
            return None;
        }
        let hit = self.cache.get(&(s, p)).cloned();
        if hit.is_some() {
            self.hits.set(self.hits.get() + 1);
        }
        hit
    }

    pub fn compose(&mut self, kind: G::Kind, children: &[ParserMatch<G>]) -> ParserMatch<G> {
        if children.is_empty() {
            return ParserMatch::Syntax(kind);
        }
        let node = SyntaxNode::new(SyntaxTrunk::Tree(kind));
        let parent = self.tree.new_node(node);
        for child in children {
            let n = match *child {
                ParserMatch::Token(t) => self.tree.new_node(SyntaxNode::new(SyntaxTrunk::Leaf(t))),
                ParserMatch::Syntax(k) => self.tree.new_node(SyntaxNode::new(SyntaxTrunk::Tree(k))),
                ParserMatch::Node(n) => n,
            };
            self.tree.append(parent, n)
        }
        ParserMatch::Node(parent)
    }

    pub fn count(&self) -> usize {
        self.tree.count()
    }

    fn skip_trivia(&self, mut s: Cursor) -> Cursor {
        while s.is_valid() && G::Lex::is_trivia(self.tree.tokens.kind(s)) {
            s = self.tree.tokens.advance(s);
        }
        s
    }

    fn peek(&self, s: Cursor) -> Option<TokenAlias<G::Lex>> {
        if s.is_valid() {
            Some(self.tree.tokens.alias(s))
        } else {
            None
        }
    }

    fn pop(&self, s: Cursor) -> Option<(Cursor, TokenAlias<G::Lex>)> {
        self.reads.set(self.reads.get() + 1);
        self.peek(s)
            .map(|t| (self.skip_trivia(self.tree.tokens.advance(s)), t))
    }
}

impl<T: Core, G: Grammar> Debug for Context<T, G> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Context")
            .field("arena_size", &self.count())
            .field("input_length", &self.tree.tokens.len())
            .field("input_reads", &self.reads.get())
            .field("cache_hits", &self.hits.get())
            .field("cache_size", &self.cache.len())
            .finish()
    }
}

/// The type of parser production functions.
pub type ParserFn<T, G> = for<'a> fn(&'a mut Context<T, G>, Cursor) -> ParserResult<G>;

pub fn memoize<T: Core, G: Grammar>(
    t: G::Tag,
    c: &mut Context<T, G>,
    s: Cursor,
    p: ParserFn<T, G>,
) -> ParserResult<G> {
    if let Some(r) = c.lookup(t, s) {
        r
    } else {
        let r = p(c, s);
        c.cache(t, s, r.clone());
        r
    }
}

pub fn repeat<T: Core, G: Grammar>(
    c: &mut Context<T, G>,
    mut s: Cursor,
    ns: &mut Vec<ParserMatch<G>>,
    ps: &[ParserFn<T, G>],
) -> Cursor {
    'outer: loop {
        for p in ps {
            if let Ok((s1, n)) = p(c, s) {
                ns.push(n);
                s = s1;
            } else {
                break 'outer;
            }
        }
    }
    s
}

pub fn intersperse<T: Core, G: Grammar, P, I>(
    c: &mut Context<T, G>,
    s: Cursor,
    ns: &mut Vec<ParserMatch<G>>,
    parser: P,
    infix: I,
) -> std::result::Result<Cursor, ParserError>
where
    P: Fn(&mut Context<T, G>, Cursor) -> ParserResult<G>,
    I: Fn(&mut Context<T, G>, Cursor) -> ParserResult<G>,
{
    let (s, n) = parser(c, s)?;
    ns.push(n);
    let mut sl = s;
    loop {
        let s = sl;
        if let Ok((s, n0, n1)) =
            infix(c, s).and_then(|(s, n0)| parser(c, s).map(|(s, n1)| (s, n0, n1)))
        {
            ns.push(n0);
            ns.push(n1);
            sl = s;
        } else {
            break Ok(s);
        }
    }
}

pub fn parse_token_with<T: Core, G: Grammar, F>(
    c: &mut Context<T, G>,
    s: Cursor,
    pred: F,
) -> ParserResult<G>
where
    F: Fn(&<G::Lex as Lexeme>::Kind) -> bool,
{
    match c.pop(s) {
        Some((s, t)) if pred(&t.kind()) => Ok((s, ParserMatch::Token(t))),
        _ => Err(ParserError::new(
            "unexpected token or end of input",
            c.span(s),
        )),
    }
}

pub fn parse_token<T: Core, G: Grammar>(
    c: &mut Context<T, G>,
    s: Cursor,
    kind: <G::Lex as Lexeme>::Kind,
) -> ParserResult<G> {
    parse_token_with(c, s, |k| *k == kind)
}