tek/proc/src/lib.rs

use proc_macro::{TokenStream, Literal};
use proc_macro2::{TokenStream as TokenStream2, Span};
use quote::{quote, ToTokens, TokenStreamExt};
use std::collections::{HashMap, BTreeMap};
use syn::{
    Error, Path, Ident, Variant, Fields, BinOp,
    Expr, ExprPath, ExprBinary, ExprAssign, ExprLit, Lit, LitStr,
    Item, ItemEnum, ItemImpl, ImplItem, ImplItemFn, ItemTrait, TraitItem, TraitItemFn,
    Signature, MetaList, Type, TypePath, FnArg, PatType, Attribute,
    parse::{ParseStream, Parse, Result},
    spanned::Spanned,
    punctuated::Punctuated,
    token::Comma
};

macro_rules! attribute {
    ($name:ident { $($body:tt)* }) => {
        #[proc_macro_attribute] pub fn $name (meta: TokenStream, item: TokenStream) -> TokenStream {
            write(self::$name::CustomAttribute(
                syn::parse_macro_input!(meta as self::$name::CustomAttributeMeta),
                syn::parse_macro_input!(item as self::$name::CustomAttributeItem),
            ))
        }
        mod $name {
            use crate::*;
            $($body)*
        }
    }
}

attribute!(commands {
    #[derive(Debug, Clone)] pub struct CustomAttribute(
        pub CustomAttributeMeta,
        pub CustomAttributeItem,
    );
    #[derive(Debug, Clone)] pub struct CustomAttributeMeta(
        pub Path,
        pub Option<LitStr>
    );
    #[derive(Debug, Clone)] pub struct CustomAttributeItem(
        pub Path,
        pub CustomAttributeItemDispatch,
        pub Item
    );

    pub type CustomAttributeItemDispatch = HashMap<Ident, (Ident, Vec<FnArg>, LitStr)>;

    impl Parse for CustomAttributeMeta {
        /// Parse contents of `#[command(...)]` attribute tag.
        fn parse (input: ParseStream) -> Result<Self> {
            let meta: Expr = input.parse()?;
            Ok(match meta {
                // Struct name only
                Expr::Path(ExprPath { path, .. }) => Self(path, None),

                // Struct name with namespace prefix
                Expr::Assign(ExprAssign { ref left, ref right, .. })
                    if let Expr::Path(ExprPath { path, .. }) = &**left
                    && let Expr::Lit(ExprLit { lit: Lit::Str(lit), .. }) = &**right
                => Self(path.clone(), Some(lit.clone())),

                // All other variants invalid
                _ => return Err(Error::new(meta.span(), format!(
                    "must be: #[commands(Struct)] or #[commands(Struct = \"struct\")], got: {meta:?}"
                )))
            })
        }
    }

    impl Parse for CustomAttributeItem {
        /// Parse contents of `trait` or `impl` block annotated with `#[command(...)]`
        fn parse (input: ParseStream) -> Result<Self> {
            let mut item: Item = input.parse()?;
            Ok(Self(
                parse_custom_attribute_item_path(&item)?,
                parse_custom_attribute_item_dispatch(&mut item)?,
                item,
            ))
        }
    }

    fn parse_custom_attribute_item_path (item: &Item) -> Result<Path> {
        Ok(match item {
            Item::Trait(ItemTrait { ident, .. }) => Path {
                leading_colon: None,
                segments: syn::punctuated::Punctuated::from_iter([
                    syn::PathSegment {
                        ident: ident.clone(),
                        arguments: syn::PathArguments::None // TODO support generics
                    }
                ])
            },

            Item::Impl(ItemImpl { self_ty, .. }) if let Type::Path(TypePath {
                path, ..
            }) = &**self_ty => path.clone(),

            _ => return Err(
                Error::new(item.span(), format!("#[commands] works on trait or inherent impl"))
            )
        })
    }

    /// Pick out annotated functions from the `trait` or `impl` block,
    /// adding them to the [CustomAttributeItemDispatch] collection.
    fn parse_custom_attribute_item_dispatch (item: &mut Item)
        -> Result<CustomAttributeItemDispatch>
    {
        let mut dispatch: CustomAttributeItemDispatch = Default::default();
        let mut dispatch_attrs = |
            attrs: &Vec<Attribute>, ident: &Ident, inputs: &Punctuated<FnArg, Comma>
        | {
            attrs.iter().filter(|attr|{
                if let syn::Meta::List(MetaList { ref path, ref tokens, .. }) = attr.meta
                && path == &Path::from(Ident::new("command", Span::call_site()))
                && let Ok(handler) = syn::parse2::<Expr>(tokens.clone())
                && let Expr::Assign(ExprAssign { ref left, ref right, .. }) = handler
                && let Expr::Lit(ExprLit { lit: Lit::Str(lit), .. }) = &**right
                && let Expr::Path(ExprPath { path, .. }) = &**left
                && path.segments.len() == 1 {
                    dispatch.insert(ident.clone(), (
                        path.segments.first().unwrap().ident.clone(),
                        inputs.iter().cloned().collect(),
                        lit.clone()
                    ));
                    false
                } else {
                    true
                }
            }).cloned().collect()
        };

        match item {
            Item::Trait(ItemTrait { items, .. }) => {
                for item in items.iter_mut() {
                    if let TraitItem::Fn(TraitItemFn {
                        attrs, sig: Signature { ident, inputs, .. }, ..
                    }) = item {
                        *attrs = dispatch_attrs(attrs, ident, inputs);
                    }
                }
            },
            Item::Impl(ItemImpl { items, .. }) => {
                for item in items.iter_mut() {
                    if let ImplItem::Fn(ImplItemFn {
                        attrs, sig: Signature { ident, inputs, .. }, ..
                    }) = item {
                        *attrs = dispatch_attrs(attrs, ident, inputs);
                    }
                }
            },
            _ => return Err(
                Error::new(item.span(), format!("#[commands] works on trait or inherent impl"))
            )
        }
        Ok(dispatch)
    }

    impl ToTokens for CustomAttribute {
        fn to_tokens (&self, out: &mut TokenStream2) {
            let Self(
                CustomAttributeMeta(command, namespace),
                CustomAttributeItem(state, items, item),
            ) = self;
            let mut variants = quote! {};
            let mut dispatch = quote! {};
            let mut keywords = quote! {};
            let mut expressions = quote! {};
            for (ident, (variant, inputs, keyword)) in items.iter() {
                write_variant(
                    &mut expressions, &mut keywords, &mut variants, &mut dispatch,
                    ident, command, variant, inputs, namespace, keyword,
                )
            }
            let impls = match item {
                Item::Impl(ItemImpl { generics, .. }) => write_for_struct(
                    keywords, expressions, dispatch, state, command, generics,
                ),
                Item::Trait { .. } => write_for_trait(
                    keywords, expressions, dispatch, state, command,
                ),
                _ => panic!("trait or inherent impl needed for #[commands]")
            };
            append(out, quote! {
                /// Command variants
                #[derive(Debug, Clone)] pub enum #command { #variants }
                #impls
                #item
            })
        }
    }

    fn write_variant (
        exps:      &mut TokenStream2,
        syms:      &mut TokenStream2,
        vars:      &mut TokenStream2,
        disp:      &mut TokenStream2,
        ident:     &Ident,
        command:   &Path,
        variant:   &Ident,
        inputs:    &[FnArg],
        namespace: &Option<LitStr>,
        keyword:   &LitStr,
    ) {
        let keyword = if let Some(namespace) = namespace {
            format!("{}/{}", namespace.value(), keyword.value())
        } else {
            keyword.value()
        };
        let mut typed = quote! {};
        let mut params = quote! {};
        let mut obtain = quote! {};
        let mut has_args = false;
        for arg in inputs.iter() {
            if let FnArg::Typed(PatType { pat, ty, .. }) = arg {
                has_args = true;
                append(&mut typed,  quote! { #pat: #ty, });
                append(&mut params, quote! { #pat, });
                append(&mut obtain, quote! { let #pat: #ty = {
                    let head = tail.head()?.unwrap_or_default();
                    let tail = tail.tail()?.unwrap_or_default();
                    match dizzle::Namespace::<#ty>::namespace(state, &head)? {
                        Some(arg) => arg,
                        None => return Err(format!("{}: arg \"{}\" ({}) got: {head} {tail}",
                            #keyword, stringify!(#pat), stringify!(#ty),
                        ).into())
                    }
                };});
            }
        }
        if has_args {
            append(vars, quote! { #variant { #typed }, });
            append(disp, quote! { #command::#variant { #params } => state.#ident(#params), });
            append(exps, quote! {
                let tail_base = tail;
                if head.src()? == Some(#keyword) {
                    let tail = tail_base;
                    #obtain
                    return Ok(Some(#command::#variant { #params }))
                }
            });
        } else {
            append(vars, quote! { #variant, });
            append(disp, quote! { #command::#variant => state.#ident(), });
            append(syms, quote! {
                if word == #keyword { return Ok(Some(#command::#variant)); }
            });
        }
    }

    fn write_for_struct (
        syms:     TokenStream2,
        exps:     TokenStream2,
        disp:     TokenStream2,
        state:    &Path,
        command:  &Path,
        generics: &syn::Generics,
    ) -> TokenStream2 {
        //let lts = Punctuated::<_, Comma>::from_iter(generics.lifetimes());
        //let cns = Punctuated::<_, Comma>::from_iter(generics.const_params());
        let tys = Punctuated::<_, Comma>::from_iter(generics.type_params());
        quote! {
            impl<#tys> dizzle::Namespace<#command> for #state {
                //def_namespace_symbols!('n |state: Self| -> #command {
                    //#syms
                //});
                //def_namespace_exps!('n |state: Self| -> #command {
                    //#exps
                //});
                fn namespace_symbol (&self, word: impl Symbol) -> Perhaps<#command> {
                    if let Some(word) = word.word()? {
                        #syms
                    }
                    Ok(None)
                }
                fn namespace_expression (&self, expr: impl Expression) -> Perhaps<#command> {
                    let state = self;
                    if let Some(expr) = expr.expr()? {
                        let head = expr.head()?;
                        let tail = expr.tail()?;
                        #exps
                    }
                    Ok(None)
                }
            }
            impl #generics dizzle::Dispatch<#state> for #command {
                fn dispatch (self, state: &mut #state) -> Perhaps<Self> {
                    match self {
                        #disp
                        _ => unreachable!()
                    }
                }
            }
        }
    }

    fn write_for_trait (
        syms: TokenStream2,
        exps: TokenStream2,
        disp: TokenStream2,
        state:   &Path,
        command: &Path,
    ) -> TokenStream2 {
        quote! {
            impl<T: #state> dizzle::Namespaced<T> for #command {
                fn namespaced_symbol <L: Symbol> (state: &T, word: L) -> Perhaps<#command> {
                    if let Some(word) = word.word()? {
                        #syms
                    }
                    Ok(None)
                }
                fn namespaced_expression <L: Expression> (state: &T, expr: L) -> Perhaps<#command> {
                    if let Some(expr) = expr.expr()? {
                        let head = expr.head()?;
                        let tail = expr.tail()?;
                        #exps
                    }
                    Ok(None)
                }
            }
            impl<T: #state> dizzle::Dispatch<T> for #command {
                fn dispatch (self, state: &mut T) -> Perhaps<Self> {
                    match self {
                        #disp
                        _ => unreachable!()
                    }
                }
            }
        }
    }
});

attribute!(command {
    #[derive(Debug, Clone)] pub struct CustomAttribute(
        pub CustomAttributeMeta,
        pub CustomAttributeItem
    );
    #[derive(Debug, Clone)] pub struct CustomAttributeMeta(
        pub Path
    );
    #[derive(Debug, Clone)] pub struct CustomAttributeItem(
        pub ItemEnum, pub HashMap<Ident, (Fields, Expr)>
    );
    impl Parse for CustomAttributeMeta {
        fn parse (input: ParseStream) -> Result<Self> {
            Ok(Self(input.parse()?))
        }
    }
    impl Parse for CustomAttributeItem {
        fn parse (input: ParseStream) -> Result<Self> {
            let mut item: ItemEnum = input.parse()?;
            let mut branches: HashMap<Ident, (Fields, Expr)> = Default::default();
            let mut variants_filtered = item.variants.clone();
            variants_filtered.clear();
            for variant in item.variants.iter_mut() {
                let Variant { attrs, ident, fields, discriminant: _ } = variant;
                let mut attrs_filtered = attrs.clone();
                attrs_filtered.clear();
                for attr in attrs.iter() {
                    if let syn::Meta::List(MetaList { ref path, ref tokens, .. }) = attr.meta
                    && path == &Path::from(Ident::new("command", Span::call_site()))
                    && let Ok(handler) = syn::parse2::<Expr>(tokens.clone())
                    {
                        branches.insert(ident.clone(), (fields.clone(), handler));
                    } else {
                        attrs_filtered.push(attr.clone());
                    }
                }
                *attrs = attrs_filtered;
                variants_filtered.push(variant.clone())
            }
            item.variants = variants_filtered;
            Ok(Self(item, branches))
        }
    }
    impl ToTokens for CustomAttribute {
        fn to_tokens (&self, out: &mut TokenStream2) {
            let Self(
                CustomAttributeMeta(state),
                CustomAttributeItem(item, branches)
            ) = self;
            let ident = &item.ident;
            let mut body = quote! {};
            for (variant, (fields, handler)) in branches.iter() {
                let (handler, wrapper) = if let Expr::Path(ExprPath { path, .. }) = handler {
                    (path, None)
                } else if let Expr::Binary(ExprBinary {
                    op: BinOp::BitOr(_), left, right, ..
                }) = handler && let Expr::Path(ExprPath {
                    path: handler, ..
                }) = &**left && let Expr::Path(ExprPath {
                    path: wrapper, ..
                }) = &**right {
                    (handler, Some(wrapper))
                } else {
                    panic!()
                };
                match fields {
                    Fields::Named(_fields)   => todo!("named command fields"),
                    Fields::Unnamed(fields) => {
                        let mut params = quote! {};
                        let mut values = quote! {};
                        for (index, _field) in fields.unnamed.iter().enumerate() {
                            let name = Ident::new(&format!("arg{index}"), Span::call_site());
                            append(&mut params, quote! { #name, });
                            append(&mut values, quote! { #name, });
                        }
                        let invocation = if let Some(wrapper) = wrapper {
                            quote! { Ok(Some(#wrapper(#handler(state, #values)))) }
                        } else {
                            quote! { #handler(state, #values) }
                        };
                        append(&mut body, quote! {
                            #ident::#variant (#params) => #invocation,
                        });
                    },
                    _ => {
                        append(&mut body, quote! {
                            #ident::#variant => #handler(state),
                        });
                    }
                }
            }
            append(out, quote! {
                #item

                impl dizzle::Act<#state> for #ident {
                    fn act (&self, state: &mut #state) -> Perhaps<Self> {
                        match self {
                            #body
                            _ => unreachable!()
                        }
                    }
                }
            })
        }
    }
});

attribute!(keyword {
    #[derive(Debug, Clone)] pub struct CustomAttribute(
        pub CustomAttributeMeta, pub CustomAttributeItem
    );
    #[derive(Debug, Clone)] pub struct CustomAttributeMeta(
        pub Path
    );
    #[derive(Debug, Clone)] pub struct CustomAttributeItem(
        pub ItemEnum, pub BTreeMap<Ident, Vec<(Literal, BTreeMap<Ident, Type>)>>
    );
    impl Parse for CustomAttributeMeta {
        fn parse (input: ParseStream) -> Result<Self> {
            Ok(Self(input.parse()?))
        }
    }
    impl Parse for CustomAttributeItem {
        fn parse (input: ParseStream) -> Result<Self> {
            let mut item: ItemEnum = input.parse()?;
            for Variant {
                attrs, ident: _, fields: _, discriminant: _
            } in item.variants.iter_mut() {
                attrs.retain(|attr|if let syn::Meta::List(MetaList {
                    ref path, tokens: ref _tokens, ..
                }) = attr.meta && path == &Path::from(Ident::new(
                    "keyword", Span::call_site()
                )) {
                    // TODO
                    false
                } else {
                    true
                });
            }
            Ok(Self(item, Default::default()))
        }
    }
    impl ToTokens for CustomAttribute {
        fn to_tokens (&self, out: &mut TokenStream2) {
            let Self(
                CustomAttributeMeta(state),
                CustomAttributeItem(item, _variants)
            ) = self;
            let ident = &item.ident;
            let body = quote! {};
            append(out, quote! {
                #item

                impl Namespace<#ident> for #state {
                    def_namespace_symbols!(|_state: #state| -> #ident { #body, });
                }
            })
        }
    }
});

fn write <T: ToTokens> (t: T) -> TokenStream {
    let mut out = TokenStream2::new();
    t.to_tokens(&mut out);
    out.into()
}

fn append (out: &mut TokenStream2, quote: TokenStream2) {
    for token in quote {
        out.append(token);
    }
}