shibari-macros 0.0.0

use std::collections::{HashMap, HashSet};

use indexmap::IndexMap;
use syn::{Block, Expr, Ident, Lifetime, LitStr, Pat, Type, Visibility};

use super::parse::{
    BranchDefault, BranchYield, ExtendKind, ExtendWith, GrammarDef, GrammarExtend, GrammarNode,
    NodeBranch, NodeKind, NodeLeaf, Output, TokenDef, TokenOp,
};
use crate::prelude::*;

pub struct ParsedGrammar {
    pub vis: Visibility,
    pub output: Type,
    pub root: ParsedBranch,
}

pub struct ParsedBranch {
    pub label: Option<Lifetime>,
    pub out: Option<Expr>,
    pub nodes: IndexMap<Ident, ParsedNode>,
    pub extends: Vec<(Ident, Option<ParsedExtendConv>)>,
    pub default: Option<Output>,
}

impl ParsedBranch {
    pub fn extend_references<'ast>(&'ast self, refs: &mut HashSet<&'ast Ident>) {
        refs.extend(self.extends.iter().map(|(i, _)| i));

        for node in self.nodes.values() {
            match node {
                ParsedNode::Leaf(_) => (),
                ParsedNode::Branch(b) => b.extend_references(refs),
            }
        }
    }
}

pub type ParsedExtendConv = (Pat, Block);

#[expect(clippy::large_enum_variant)]
pub enum ParsedNode {
    Leaf(Output),
    Branch(ParsedBranch),
}

pub fn parse_tokens(
    tokens: Vec<TokenDef>,
    diag: &mut TokenStream,
) -> HashMap<Ident, (Pat, Option<LitStr>)> {
    let mut pats = HashMap::new();

    for TokenDef {
        token: _,
        ident,
        eq_token: _,
        pat,
        op,
        semi_token: _,
    } in tokens
    {
        use std::collections::hash_map::Entry;

        let ident_span = ident.span();
        match pats.entry(ident) {
            Entry::Vacant(v) => {
                v.insert((
                    pat,
                    op.map(
                        |TokenOp {
                             arrow_token: _,
                             lit,
                         }| lit,
                    ),
                ));
            },
            Entry::Occupied(_) => diag.extend(
                ident_span
                    .error("Duplicate token name")
                    .into_compile_error(),
            ),
        }
    }

    pats
}

pub fn parse_grammars(
    grammars: Vec<GrammarDef>,
    diag: &mut TokenStream,
) -> IndexMap<Ident, ParsedGrammar> {
    let mut parsed = IndexMap::new();
    for GrammarDef {
        vis,
        grammar_token: _,
        ident,
        colon_token: _,
        output,
        root,
    } in grammars
    {
        use indexmap::map::Entry;

        let ident_span = ident.span();
        match parsed.entry(ident) {
            Entry::Vacant(v) => {
                v.insert(ParsedGrammar {
                    vis,
                    output,
                    root: parse_branch(None, root, diag),
                });
            },
            Entry::Occupied(_) => diag.extend(
                ident_span
                    .error("Duplicate grammar name")
                    .into_compile_error(),
            ),
        }
    }

    parsed
}

fn parse_branch(
    label: Option<Lifetime>,
    branch: NodeBranch,
    diag: &mut TokenStream,
) -> ParsedBranch {
    let mut node_map = IndexMap::new();

    let NodeBranch {
        brace: _,
        out,
        nodes,
        extends,
        default,
    } = branch;

    for GrammarNode {
        label,
        tok_ident,
        kind,
    } in nodes
    {
        use indexmap::map::Entry;

        let ident_span = tok_ident.span();
        match node_map.entry(tok_ident) {
            Entry::Vacant(v) => {
                v.insert(match kind {
                    NodeKind::Leaf(NodeLeaf {
                        arrow_token: _,
                        out,
                        semi_token: _,
                    }) => ParsedNode::Leaf(out),
                    NodeKind::Branch(b) => {
                        ParsedNode::Branch(parse_branch(label.map(|l| l.name), b, diag))
                    },
                });
            },
            Entry::Occupied(_) => diag.extend(
                ident_span
                    .error("Duplicate token in grammar block")
                    .into_compile_error(),
            ),
        }
    }

    ParsedBranch {
        label,
        out: out.map(
            |BranchYield {
                 yield_token: _,
                 expr,
                 semi_token: _,
             }| expr,
        ),
        nodes: node_map,
        extends: extends
            .into_iter()
            .map(
                |GrammarExtend {
                     extend_token: _,
                     ident,
                     kind,
                 }| {
                    (ident, match kind {
                        ExtendKind::From(_) => None,
                        ExtendKind::With(ExtendWith {
                            paren: _,
                            pat,
                            block,
                        }) => Some((pat, block)),
                    })
                },
            )
            .collect(),
        default: default.map(|BranchDefault { ddot_token: _, out }| out),
    }
}