mwbot_derive/

lib.rs

// SPDX-License-Identifier: GPL-3.0-or-later
// Copyright (C) 2023 Kunal Mehta <legoktm@debian.org>
//! This crate provides the proc_macro for [`mwbot`](https://docs.rs/mwbot/),
//! please refer to its documentation for usage.
#![deny(clippy::all)]
#![deny(rustdoc::all)]

use proc_macro::TokenStream;
use proc_macro2::{Ident, TokenStream as TokenStream2};
use quote::{format_ident, quote};
use std::collections::HashMap;
use syn::meta::ParseNestedMeta;
use syn::spanned::Spanned;
use syn::{
    parse_macro_input, Data, DeriveInput, Error, Expr, ExprLit, Field, Fields,
    GenericArgument, Lit, LitBool, LitStr, Path, PathArguments, Result, Type,
};

// TODO: figure out how to link to the trait in the docs. See
// <https://doc.rust-lang.org/rustdoc/write-documentation/linking-to-items-by-name.html#namespaces-and-disambiguators>
/// macro to implement the `Generator` trait
///
/// See the `Generator` trait documentation for details on usage.
#[proc_macro_derive(
    Generator,
    attributes(generator, params, param, templated_param)
)]
pub fn generator(input: TokenStream) -> TokenStream {
    // Parse the input tokens into a syntax tree
    let input = parse_macro_input!(input as DeriveInput);
    let tokens = build(input).unwrap_or_else(|err| err.into_compile_error());
    tokens.into()
}

fn build(input: DeriveInput) -> Result<TokenStream2> {
    let name = &input.ident;
    let params = parse(&input)?;
    let fixed_params = parse_fixed_params(&input)?;
    let generator_option = parse_generator_option(&input)?;
    //dbg!(&params);
    let plain_impl = build_impl(name, &params, &generator_option.exports);
    //println!("{}", &plain_impl);
    let trait_impl =
        build_trait_impl(name, &params, fixed_params, &generator_option);

    // Build the output, possibly using quasi-quotation
    let expanded = quote! {
        #plain_impl
        #trait_impl
    };
    Ok(expanded)
}

struct Parameter {
    rust_name: Ident,
    rust_type: Type,
    doc: Option<LitStr>,
    mw_name: LitStr,
    required: bool,
}

impl Parameter {
    /// Whether rust_type is "bool"
    fn is_boolean(&self) -> bool {
        if let Type::Path(type_path) = &self.rust_type {
            if let Some(first) = type_path.path.segments.first() {
                return first.ident == "bool";
            }
        }

        false
    }
}

/// Peeks through Option<T> to determine the underlying type and if its required
/// Given `Option<foo>`, return `(foo, false)`
/// Given `foo`, return `(foo, true)`
fn parse_through_option(ty: &Type) -> (Type, bool) {
    if let Type::Path(type_path) = ty {
        if let Some(first) = type_path.path.segments.first() {
            if first.ident == "Option" {
                if let PathArguments::AngleBracketed(inside) = &first.arguments
                {
                    if let Some(GenericArgument::Type(ty)) = inside.args.first()
                    {
                        return (ty.clone(), false);
                    }
                }
            }
        }
    }
    (ty.clone(), true)
}

/// Build the `Generator` trait implementation
fn build_trait_impl(
    name: &Ident,
    params: &[Parameter],
    fixed_params: HashMap<String, LitStr>,
    generator_option: &GeneratorOption,
) -> TokenStream2 {
    let GeneratorOption {
        exports,
        return_type,
        response_type,
        transform_fn,
        wrap_in_vec,
    } = generator_option;

    let fixed_params: Vec<_> = fixed_params
        .into_iter()
        .map(|(key, value)| {
            quote! {
                map.insert(#key, #value.to_string());
            }
        })
        .collect();
    let params: Vec<_> = params
        .iter()
        .map(|param| {
            let mw_name = &param.mw_name;
            let rust_name = &param.rust_name;
            if param.required {
                quote! {
                    map.insert(#mw_name, #exports::ParamValue::stringify(&self.#rust_name));
                }
            } else if param.is_boolean() {
                // For boolean parameters, include it if it's true, otherwise
                // omit it entirely.
                quote! {
                    if self.#rust_name.unwrap_or(false) {
                        map.insert(#mw_name, "1".to_string());
                    }
                }
            } else {
                quote! {
                    if let Some(value) = &self.#rust_name {
                        map.insert(#mw_name, #exports::ParamValue::stringify(value));
                    }
                }
            }
        })
        .collect();

    let wrap_vec = if wrap_in_vec.value {
        quote! {
            values
        }
    } else {
        quote! {
            {
                let mut vec = Vec::new();
                vec.push(values);
                vec
            };
        }
    };

    quote! {
        impl #exports::Generator for #name {
            type Output = #exports::Result<#return_type>;

            fn params(&self) -> #exports::HashMap<&'static str, String> {
                let mut map = #exports::HashMap::new();
                #(#fixed_params)*
                #(#params)*
                map
            }

            fn generate(self, bot: &#exports::Bot) -> #exports::tokio::Receiver<Self::Output> {
                let (tx, rx) = #exports::tokio::channel(50);
                let bot: #exports::Bot = #exports::Clone::clone(&bot);
                #exports::tokio::spawn(async move {
                    let params = #exports::IntoIterator::into_iter(#exports::Generator::params(&self));
                    let params = #exports::Iterator::map(params, |(k, v)| (k.to_string(), v));
                    let params = #exports::Iterator::collect::<#exports::HashMap<String, String>>(params);

                    let mut params = #exports::Params {
                        main: params,
                        ..#exports::Default::default()
                    };
                    loop {
                        let resp: #response_type =
                            match #exports::mwapi_responses::query_api(&bot.api(), params.merged())
                                .await
                            {
                                Ok(resp) => resp,
                                Err(err) => {
                                    match tx.send(Err(<#exports::Error as #exports::From<_>>::from(err))).await {
                                        Ok(_) => break,
                                        // Receiver hung up, just abort
                                        Err(_) => return,
                                    }
                                }
                            };
                        params.continue_ = #exports::Clone::clone(&resp.continue_);
                        for item in #exports::mwapi_responses::ApiResponse::<_>::into_items(resp) {
                            let values = match #transform_fn(&bot, item)  {
                                Ok(values) => values,
                                Err(err) => {
                                    if tx.send(Err(err)).await.is_err() {
                                        return;
                                    };
                                    continue;
                                },
                            };
                            let values = #wrap_vec;

                            for value in values {
                                if tx.send(Ok(value)).await.is_err() {
                                    // Receiver hung up, just abort
                                    return;
                                }
                            }
                        }
                        if params.continue_.is_empty() {
                            // All done
                            break;
                        }
                    }
                });
                rx
            }
        }
    }
}

/// Build the type's implementation with `new()` and parameter builders
fn build_impl(
    name: &Ident,
    params: &[Parameter],
    exports: &Path,
) -> TokenStream2 {
    let type_def: Vec<_> = params
        .iter()
        .filter(|param| param.required)
        .map(|param| {
            let name = &param.rust_name;
            let ty = &param.rust_type;
            quote! { #name: impl #exports::Into<#ty> }
        })
        .collect();
    let fields: Vec<_> = params
        .iter()
        .map(|param| {
            let name = &param.rust_name;
            if param.required {
                quote! { #name: #exports::Into::into(#name) }
            } else {
                quote! { #name: #exports::Option::None }
            }
        })
        .collect();

    let setters: Vec<_> = params
        .iter()
        .filter(|param| !param.required)
        .map(|param| setter(param, exports))
        .collect();

    quote! {
        impl #name {
            pub fn new( #(#type_def),* ) -> Self {
                Self {
                    #(#fields),*
                }
            }

            #(#setters)*
        }
    }
}

/// Build a parameter setter function
fn setter(param: &Parameter, exports: &Path) -> TokenStream2 {
    let name = &param.rust_name;
    let ty = &param.rust_type;
    let doc = match &param.doc {
        Some(doc) => quote! {
            #[doc = #doc]
        },
        None => quote! {},
    };
    assert!(!param.required);
    let with_name = format_ident!("with_{}", name);
    let set_name = format_ident!("set_{}", name);
    quote! {
        #doc
        pub fn #name(mut self, value: impl #exports::Into<#ty>) -> Self {
            self.#name = Some(#exports::Into::into(value));
            self
        }

        pub fn #with_name(mut self, value: impl #exports::Into<Option<#ty>>) -> Self {
            self.#name = #exports::Into::into(value);
            self
        }

        pub fn #set_name(&mut self, value: impl #exports::Into<Option<#ty>>) {
            self.#name = #exports::Into::into(value);
        }
    }
}

/// Parse Rust struct types into our `Parameter` type
fn parse(input: &DeriveInput) -> Result<Vec<Parameter>> {
    let mut params = vec![];
    let data = if let Data::Struct(data) = &input.data {
        data
    } else {
        return Err(Error::new(input.ident.span(), "expected a struct"));
    };
    let fields = if let Fields::Named(fields) = &data.fields {
        fields
    } else {
        return Err(Error::new(
            input.ident.span(),
            "struct fields must have names",
        ));
    };
    for field in &fields.named {
        let (rust_type, required) = parse_through_option(&field.ty);

        let (doc, mw_name) = parse_mw_name(field)?;
        let param = Parameter {
            rust_name: field.ident.clone().ok_or_else(|| {
                Error::new(field.span(), "struct fields must have names")
            })?,
            doc,
            mw_name,
            rust_type,
            required,
        };
        params.push(param);
    }
    Ok(params)
}

/// Parse the generator attribute
fn parse_fixed_params(input: &DeriveInput) -> Result<HashMap<String, LitStr>> {
    let mut map = HashMap::new();
    for attr in &input.attrs {
        if attr.path().is_ident("params") {
            attr.parse_nested_meta(|meta| {
                let key = meta
                    .path
                    .get_ident()
                    .ok_or_else(|| meta.error("invalid parameter name"))?;
                let value: LitStr = meta.value()?.parse()?;
                map.insert(key.to_string(), value);
                Ok(())
            })?;
        }
    }
    if !map
        .keys()
        .any(|key| key == "generator" || key == "list" || key == "prop")
    {
        Err(Error::new(
            input.ident.span(),
            "Missing #[params(generator = \"...\")], #[params(list = \" ... \")] or #[params(prop = \" ... \")] attribute",
        ))
    } else {
        Ok(map)
    }
}

fn get_lit_str(ident: &str, meta: &ParseNestedMeta) -> Result<Option<LitStr>> {
    if !meta.path.is_ident(ident) {
        return Ok(None);
    }

    let expr: Expr = meta.value()?.parse()?;
    let mut value = &expr;
    while let Expr::Group(e) = value {
        value = &e.expr;
    }
    if let Expr::Lit(ExprLit {
        lit: Lit::Str(lit), ..
    }) = value
    {
        let suffix = lit.suffix();
        if !suffix.is_empty() {
            return Err(meta.error(format!(
                "unexpected suffix `{}` on string literal",
                suffix
            )));
        }
        Ok(Some(lit.clone()))
    } else {
        Err(meta.error(format!(
            "expected attribute to be a string: `{}` = \"...\"",
            ident
        )))
    }
}

fn get_path(meta: &ParseNestedMeta, ident: &str) -> Result<Option<Path>> {
    if !meta.path.is_ident(ident) {
        return Ok(None);
    }
    let string = match get_lit_str(ident, meta)? {
        Some(string) => string,
        None => return Ok(None),
    };

    string.parse().map(Some)
}

fn get_type(meta: &ParseNestedMeta, ident: &str) -> Result<Option<Type>> {
    if !meta.path.is_ident(ident) {
        return Ok(None);
    }
    let string = match get_lit_str(ident, meta)? {
        Some(string) => string,
        None => return Ok(None),
    };

    string.parse().map(Some)
}

struct GeneratorOption {
    exports: Path,
    return_type: Type,
    response_type: Type,
    transform_fn: Path,
    wrap_in_vec: LitBool,
}

fn parse_generator_option(input: &DeriveInput) -> Result<GeneratorOption> {
    let mut crate_: Option<Path> = None;
    let mut return_type: Option<Type> = None;
    let mut response_type: Option<Type> = None;
    let mut transform_fn: Option<Path> = None;
    let mut wrap_in_vec: Option<LitBool> = None;
    for attr in &input.attrs {
        if !attr.path().is_ident("generator") {
            continue;
        }
        attr.parse_nested_meta(|meta| {
            if let Some(path) = get_path(&meta, "crate")? {
                crate_ = Some(path);
                return Ok(());
            }
            if let Some(path) = get_type(&meta, "return_type")? {
                return_type = Some(path);
                return Ok(());
            }
            if let Some(path) = get_type(&meta, "response_type")? {
                response_type = Some(path);
                return Ok(());
            }
            if let Some(path) = get_path(&meta, "transform_fn")? {
                transform_fn = Some(path);
                return Ok(());
            }
            if meta.path.is_ident("wrap_in_vec") {
                wrap_in_vec = Some(meta.value()?.parse()?);
                return Ok(());
            }
            Ok(())
        })?;
    }
    let crate_ = crate_.unwrap_or_else(|| syn::parse_quote!(crate));
    let exports: Path = syn::parse_quote!(#crate_::generators::__exports);
    Ok(GeneratorOption {
        exports: exports.clone(),
        return_type: return_type
            .unwrap_or_else(|| syn::parse_quote!(#exports::Page)),
        response_type: response_type
            .unwrap_or_else(|| syn::parse_quote!(#exports::InfoResponse)),
        transform_fn: transform_fn
            .unwrap_or_else(|| syn::parse_quote!(#exports::transform_to_page)),
        wrap_in_vec: wrap_in_vec.unwrap_or_else(|| syn::parse_quote!(false)),
    })
}

/// Parse the parameter attribute
fn parse_mw_name(field: &Field) -> Result<(Option<LitStr>, LitStr)> {
    let mut doc = None;
    let mut param = None;
    for attr in &field.attrs {
        if attr.path().is_ident("doc") {
            if let Expr::Lit(expr) = &attr.meta.require_name_value()?.value {
                if let Lit::Str(lit) = &expr.lit {
                    doc = Some(lit.clone());
                }
            }
            // else: error?
        } else if attr.path().is_ident("param") {
            let parsed: LitStr = attr.parse_args()?;
            param = Some(parsed);
        }
    }
    match param {
        Some(param) => Ok((doc, param)),
        None => Err(Error::new(
            field.ident.span(),
            "Missing #[param(\"...\")] attribute",
        )),
    }
}