lutra-compiler 0.6.0

use chumsky::input::ValueInput;
use chumsky::prelude::*;
use indexmap::IndexMap;

use crate::Span;
use crate::pr::*;

use super::helpers::*;
use super::{PError, PExtra, TokenKind};
use super::{expr, types};

/// The top-level parser
pub fn source<'src, I>() -> impl Parser<'src, I, Source, PExtra<'src>> + Clone + 'src
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    let is_submodule = keyword("submodule").or_not().map(|x| x.is_some());

    let ty = types::type_expr();
    let expr = expr::expr(ty.clone());

    let mod_content = recursive(|mod_content| {
        let module_def = keyword("module")
            .ignore_then(def_name())
            .then(
                mod_content
                    .delimited_by(ctrl('{'), ctrl('}'))
                    .map_with(|m, e| set_content_span(m, e.span())),
            )
            .map(|(name, module_def)| (name, DefKind::Module(module_def)))
            .labelled("module definition");

        let annotations = ctrl('@')
            .ignore_then(expr.clone().map(Box::new))
            .map(|expr| Anno { expr })
            .labelled("annotation")
            .repeated()
            .collect::<Vec<_>>();

        let def_kind = choice((
            module_def.boxed(),
            type_def(ty.clone()).boxed(),
            import_def().boxed(),
            external_func_def(ty.clone()).boxed(),
            func_def(expr.clone(), ty.clone()).boxed(),
            const_def(expr.clone(), ty.clone()).boxed(),
            anno_def(ty).boxed(),
        ))
        .labelled("definition");

        let def = (doc_annotation().or_not())
            .then(annotations.clone())
            .then(def_kind)
            .map_with(|((doc_ann, ann), (name, def_kind)), e| {
                let mut annotations = Vec::with_capacity(ann.len() + 1);
                annotations.extend(doc_ann);
                annotations.extend(ann);
                let def = Def {
                    kind: def_kind,
                    span: Some(e.span()),
                    span_name: Some(name.1),
                    annotations,
                };
                (name.0, def)
            });

        let defs = def
            .repeated()
            .collect::<Vec<_>>()
            .validate(|defs, e, emitter| into_module(defs, e.span(), emitter));

        let self_annotation = choice((
            ctrl('@')
                .then(ctrl('!'))
                .ignore_then(expr.map(Box::new))
                .map(|expr| Anno { expr })
                .labelled("annotation"),
            self_doc_comment(),
        ));

        self_annotation
            .repeated()
            .collect::<Vec<_>>()
            .then(defs)
            .map(|(annotations, mut module_def)| {
                module_def.annotations = annotations;
                module_def
            })
    });

    is_submodule
        .then(mod_content)
        .map_with(|(is_submodule, root), e| Source {
            is_submodule,
            root,
            span: e.span(),
        })
        .then_ignore(end())
}

fn into_module(
    def_vec: Vec<(String, Def)>,
    span: Span,
    emitter: &mut chumsky::input::Emitter<PError<'_>>,
) -> ModuleDef {
    let mut imports = Vec::new();
    let mut defs = IndexMap::with_capacity(def_vec.len());
    for (name, def) in def_vec {
        // hack that allows imports not to have names
        if name.is_empty() {
            assert!(matches!(def.kind, DefKind::Import(_)));
            imports.push(def);
            continue;
        }

        let span = def.span.unwrap();
        let conflict = defs.insert(name, def);
        if let Some(conflict) = conflict {
            emitter.emit(Rich::custom(span, "duplicate name"));
            emitter.emit(Rich::custom(conflict.span.unwrap(), "duplicate name"));
        }
    }

    ModuleDef {
        imports,
        defs,
        annotations: vec![],
        span_content: Some(span),
    }
}

fn set_content_span(mut module: ModuleDef, mut span: Span) -> ModuleDef {
    span.start += 1;
    span.len = span.len.saturating_sub(2);
    module.span_content = Some(span);
    module
}

/// Parses `## ...` doc comment lines into a `@doc("...")` annotation.
fn doc_annotation<'src, I>() -> impl Parser<'src, I, Anno, PExtra<'src>> + Clone + 'src
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    select! {
        TokenKind::DocComment(text) => text,
    }
    .repeated()
    .at_least(1)
    .collect::<Vec<_>>()
    .map_with(|lines, e| Anno::new_std_doc(lines.join("\n"), e.span()))
    .labelled("doc comment")
}

/// Parses `#! ...` self doc comment lines into a `@!doc("...")` annotation.
fn self_doc_comment<'src, I>() -> impl Parser<'src, I, Anno, PExtra<'src>> + Clone + 'src
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    select! {
        TokenKind::DocCommentSelf(text) => text,
    }
    .repeated()
    .at_least(1)
    .collect::<Vec<_>>()
    .map_with(|lines, e| Anno::new_std_doc(lines.join("\n"), e.span()))
    .labelled("self doc comment")
}

struct NameAndSpan(String, Span);

impl NameAndSpan {
    fn new(n: String, s: Span) -> Self {
        Self(n, s)
    }
}

fn def_name<'src, I>() -> impl Parser<'src, I, NameAndSpan, PExtra<'src>> + Clone + 'src
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    ident_part().map_with(|n, e| NameAndSpan::new(n, e.span()))
}

fn const_def<'src, I>(
    expr: impl Parser<'src, I, Expr, PExtra<'src>> + Clone + 'src,
    ty: impl Parser<'src, I, Ty, PExtra<'src>> + Clone + 'src,
) -> impl Parser<'src, I, (NameAndSpan, DefKind), PExtra<'src>> + Clone
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    keyword("const")
        .ignore_then(def_name())
        .then(ctrl(':').ignore_then(ty).or_not())
        .then(ctrl('=').ignore_then(expr.clone()).map(Box::new))
        .map(|((name, ty), value)| {
            (
                name,
                DefKind::Expr(ExprDef {
                    value,
                    ty,
                    constant: true,
                }),
            )
        })
        .labelled("constant definition")
}

fn func_def<'src, I>(
    expr: impl Parser<'src, I, Expr, PExtra<'src>> + Clone + 'src,
    ty: impl Parser<'src, I, Ty, PExtra<'src>> + Clone + 'src,
) -> impl Parser<'src, I, (NameAndSpan, DefKind), PExtra<'src>> + Clone
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    let head = keyword("func").ignore_then(def_name());

    let params = delimited_by_parenthesis(
        super::expr::func_param(ty.clone())
            .separated_by(ctrl(','))
            .allow_trailing()
            .collect(),
        |_| vec![],
    );

    let return_ty = ctrl(':').ignore_then(ty.clone()).or_not();

    let ty_params = keyword("where")
        .ignore_then(types::type_params(ty.clone()))
        .or_not()
        .map(|x| x.unwrap_or_default());

    let body = just(TokenKind::ArrowThin).ignore_then(expr.map(Box::new));

    head.then(params)
        .then(return_ty)
        .then(ty_params)
        .then(body)
        .map_with(|((((name, params), return_ty), ty_params), body), e| {
            let span = e.span();
            let func = Func {
                return_ty,
                body,
                params,
                ty_params,
            };
            let value = Expr::new_with_span(ExprKind::Func(Box::new(func)), span);
            let def = DefKind::Expr(ExprDef {
                ty: None,
                value: Box::new(value),
                constant: false,
            });
            (name, def)
        })
        .labelled("function definition")
}

fn external_func_def<'src, I>(
    ty: impl Parser<'src, I, Ty, PExtra<'src>> + Clone + 'src,
) -> impl Parser<'src, I, (NameAndSpan, DefKind), PExtra<'src>> + Clone
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    let head = keyword("external")
        .ignore_then(keyword("func"))
        .ignore_then(def_name());

    let params = super::types::func_params(ty.clone());

    let return_ty = ctrl(':').ignore_then(ty.clone());

    let ty_params = keyword("where")
        .ignore_then(types::type_params(ty.clone()))
        .or_not()
        .map(|x| x.unwrap_or_default());

    head.then(params)
        .then(return_ty)
        .then(ty_params)
        .map(|(((name, params), return_ty), ty_params)| {
            let ty = Ty::new(TyKind::Func(TyFunc {
                params,
                body: Some(Box::new(return_ty)),
                ty_params,
            }));
            (name, DefKind::External(ty))
        })
        .labelled("external function definition")
}

fn type_def<'src, I>(
    ty: impl Parser<'src, I, Ty, PExtra<'src>> + Clone + 'src,
) -> impl Parser<'src, I, (NameAndSpan, DefKind), PExtra<'src>> + Clone
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    keyword("type")
        .ignore_then(def_name())
        .then(
            choice((
                // type alias
                ctrl(':').ignore_then(ty.clone()).map(|ty| TyDef {
                    ty,
                    is_framed: false,
                    framed_label: None,
                }),
                // framed type (nominal / new type)
                delimited_by_parenthesis(
                    ident_part().then_ignore(ctrl(':')).or_not().then(ty),
                    |p| (None, Ty::new_with_span(TyKind::Tuple(vec![]), p)),
                )
                .map(|(framed_label, ty)| TyDef {
                    ty,
                    is_framed: true,
                    framed_label,
                }),
            ))
            .map(DefKind::Ty),
        )
        .labelled("type definition")
}

fn anno_def<'src, I>(
    ty: impl Parser<'src, I, Ty, PExtra<'src>> + Clone + 'src,
) -> impl Parser<'src, I, (NameAndSpan, DefKind), PExtra<'src>> + Clone
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    let params = super::types::func_params(ty);

    keyword("anno")
        .ignore_then(def_name())
        .then(params)
        .map(|(name, params)| (name, DefKind::Anno(AnnoDef { params })))
        .labelled("annotation definition")
}

fn import_def<'src, I>() -> impl Parser<'src, I, (NameAndSpan, DefKind), PExtra<'src>> + Clone
where
    I: ValueInput<'src, Token = TokenKind, Span = Span>,
{
    let import = recursive(|import_part| {
        expr::path()
            .then(choice((
                just(TokenKind::PathSep)
                    .ignore_then(delimited_by_parenthesis(
                        import_part
                            .separated_by(ctrl(','))
                            .allow_trailing()
                            .collect(),
                        |_| vec![],
                    ))
                    .map(|children| ImportKind::Many(Path::empty(), children)),
                just(TokenKind::PathSep)
                    .ignore_then(ctrl('*'))
                    .to(ImportKind::Star(Path::empty())),
                keyword("as")
                    .ignore_then(ident_part())
                    .or_not()
                    .map(|alias| ImportKind::Single(Path::empty(), alias)),
            )))
            .map_with(|(path, mut kind), e| {
                let span = Some(e.span());
                *kind.path_mut() = path;
                ImportDef { kind, span }
            })
    });

    keyword("import")
        .ignore_then(import)
        .map_with(|import, e| {
            let s = e.span();
            (NameAndSpan::new("".into(), s), DefKind::Import(import))
        })
        .labelled("import statement")
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::parser::test::parse_with;

    #[test]
    fn test_module_annotation_basic() {
        let src = "submodule\n@!my_ann\n\nconst x = 1";
        let parsed = parse_with(src, |_| Box::new(source())).unwrap();
        assert!(parsed.is_submodule);
        assert_eq!(parsed.root.annotations.len(), 1);
        assert!(matches!(
            &parsed.root.annotations[0].expr.kind,
            crate::pr::ExprKind::Ident(p) if p.len() == 1 && p.first() == "my_ann"
        ));
        assert_eq!(parsed.root.defs.len(), 1);
    }

    #[test]
    fn test_module_annotation_multiple() {
        let src = "submodule\n@!ann1\n@!ann2\n\nconst x = 1";
        let parsed = parse_with(src, |_| Box::new(source())).unwrap();
        assert_eq!(parsed.root.annotations.len(), 2);
    }

    #[test]
    fn test_module_annotation_no_submodule() {
        let src = "const x = 1";
        let parsed = parse_with(src, |_| Box::new(source())).unwrap();
        assert!(!parsed.is_submodule);
        assert_eq!(parsed.root.annotations.len(), 0);
    }

    #[test]
    fn test_module_annotation_with_call_expr() {
        let src = "submodule\n@!deprecated(\"use other instead\")\n\nconst x = 1";
        let parsed = parse_with(src, |_| Box::new(source())).unwrap();
        assert_eq!(parsed.root.annotations.len(), 1);
        assert!(matches!(
            &parsed.root.annotations[0].expr.kind,
            crate::pr::ExprKind::Call(_)
        ));
    }
}