molt_argparse_procmacro/

lib.rs

//! This proc-macro crate is intended to be re-exported by `molt_argparse`. Please do NOT use it alone.

use proc_macro::TokenStream;
use quote::{ quote, ToTokens };
use indexmap::IndexMap;

struct ArgInfo {
    funcarg_ident: syn::Ident,
    funcarg_def: syn::PatType,
    /// if keyed, this is specified.
    keyarg_name: Option<String>,
    is_optional: bool,
    is_flag: bool,
    raw_type: proc_macro2::TokenStream,
    value_converter: proc_macro2::TokenStream
}

fn tokens_equal(
    a: &(impl Clone + IntoIterator<Item = proc_macro2::TokenTree>),
    b: impl IntoIterator<Item = proc_macro2::TokenTree>
) -> bool
{
    use proc_macro2::*;
    let a = a.clone().into_iter().collect::<Vec<_>>();
    let b = b.into_iter().collect::<Vec<_>>();

    if a.len() != b.len() { return false; }
    for (ai, bi) in a.iter().zip(b.iter()) {
        match (ai, bi) {
            (TokenTree::Ident(ia), TokenTree::Ident(ib)) if ia == ib => { },
            (TokenTree::Punct(pa), TokenTree::Punct(pb))
                if pa.as_char() == pb.as_char() => { },
            _ => return false
        }
    }
    true
}

fn tokens_match(
    a: &(impl Clone + IntoIterator<Item = proc_macro2::TokenTree>),
    l: impl IntoIterator<Item = proc_macro2::TokenTree>,
    r: impl IntoIterator<Item = proc_macro2::TokenTree>
) -> Option<proc_macro2::TokenStream>
{
    use proc_macro2::*;
    let mut a = a.clone().into_iter().collect::<Vec<_>>();
    let l = l.into_iter().collect::<Vec<_>>();
    let r = r.into_iter().collect::<Vec<_>>();
    
    if a.len() < l.len() + r.len() { return None; }
    for i in 0..l.len() {
        match (&a[i], &l[i]) {
            (TokenTree::Ident(ia), TokenTree::Ident(ib)) if ia == ib => { },
            (TokenTree::Punct(pa), TokenTree::Punct(pb))
                if pa.as_char() == pb.as_char() => { },
            _ => return None,
        }
    }
    for i in 0..r.len() {
        match (&a[a.len() - r.len() + i], &r[i]) {
            (TokenTree::Ident(ia), TokenTree::Ident(ib)) if ia == ib => { },
            (TokenTree::Punct(pa), TokenTree::Punct(pb))
                if pa.as_char() == pb.as_char() => { },
            _ => return None,
        }
    }

    let ret = a.splice(l.len()..(a.len() - r.len()), []);
    Some(ret.into_iter().collect())
}

#[proc_macro_attribute]
pub fn molt_argparse(args: TokenStream, decl: TokenStream) -> TokenStream {
    // println!("orig decl: {:#?}", decl);
    let mut args = syn::parse_macro_input!(args as syn::LitStr).value()
        .split_whitespace()
        .map(|s| s.to_string())
        .collect::<Vec<String>>();
    let decl = syn::parse_macro_input!(decl as syn::ItemFn);
    // println!("{:?}", args);
    // println!("{:#?}", decl);

    let mut funcargs: IndexMap<String, (syn::PatType, syn::PatIdent)> =
        decl.sig.inputs.iter().map(|fnarg| match fnarg {
            syn::FnArg::Typed(typ) => match &*typ.pat {
                syn::Pat::Ident(patident) => {
                    (patident.ident.to_string(),
                     (typ.clone(), patident.clone()))
                }
                _ => panic!("unexpected arg def")
            },
            syn::FnArg::Receiver(_) =>
                panic!("self argument is unsupported by molt_argparse")
        })
        .collect();

    let contextdecl = {
        if args[0].starts_with("{") {
            let s = args.remove(0);
            let p = &s[1..s.len() - 1];
            let (pattype, _) = funcargs.remove(p).expect("context arg not found");
            quote! {
                let #pattype = _interp.context(_ctxid);
            }
        }
        else { quote! { } }
    };

    let args = args.into_iter().map(|arg| {
        let (funcarg_name, keyarg_name) = match arg.starts_with("-") {
            true => match arg.find(':') {
                Some(i) => (arg[i + 1..].to_string(),
                            Some(arg[0..i].to_string())),
                None => (arg[1..].to_string(), Some(arg))
            }
            false => (arg, None)
        };
        let (pt, pi) = funcargs.remove(&funcarg_name)
            .expect("arg not found");

        let type_tokens = pt.ty.to_token_stream();
        let (raw_type, is_optional) = match tokens_match(
            &type_tokens,
            quote!{ Option< }, quote!{ > }
        ) {
            Some(ts) => (ts, true),
            None => (type_tokens, false)
        };
        
        let raw_type_vec = raw_type.clone().into_iter().collect::<Vec<_>>();

        let is_flag = keyarg_name.is_some() &&
            !is_optional &&
            tokens_equal(&raw_type, quote!{ bool });

        let value_converter = {
            if tokens_equal(&raw_type_vec, quote!{ bool }) {
                quote!{ _molt_v.as_bool().ok() }
            }
            else if tokens_equal(&raw_type_vec, quote!{ &str }) {
                quote!{ Some(_molt_v.as_str()) }
            }
            else if let Some(_) = tokens_match(
                &raw_type_vec,
                quote!{ &mut }, quote!{ }
            ) {
                panic!("mutable reference is not allowed in molt_argparse");
            }
            else if let Some(t) = tokens_match(
                &raw_type_vec,
                quote!{ &Vec< }, quote!{ > }
            ) {
                quote!{ _molt_v.as_other::<molt_argparse::WrapVec<#t>>()
                         .map(|v| &v.0) }
            }
            else if let Some(t) = tokens_match(
                &raw_type_vec,
                quote!{ Vec< }, quote!{ > }
            ) {
                quote!{ _molt_v.as_other::<molt_argparse::WrapVec<#t>>()
                         .map(|v| v.0.clone()) }
            }
            else if let Some(v) = tokens_match(
                &raw_type_vec,
                quote!{ & }, quote!{ }
            ) {
                quote!{ _molt_v.as_other::<#v>().map(|r| r.as_ref()) }
            }
            else if let Some(v) = tokens_match(
                &raw_type_vec,
                quote!{ Rc< }, quote!{ > }
            ) {
                quote!{ _molt_v.as_other::<#v>() }
            }
            else {
                quote!{ _molt_v.as_other::<#(#raw_type_vec)*>()
                         .map(|r| r.as_ref().clone()) }
            }
        };

        ArgInfo {
            funcarg_ident: pi.ident,
            funcarg_def: pt,
            keyarg_name,
            is_optional,
            is_flag,
            raw_type,
            value_converter
        }
    }).collect::<Vec<_>>();

    assert_eq!(funcargs.len(), 0, "excess funcargs");

    let argdecls = args.iter().map(|ai| {
        let ident = &ai.funcarg_ident;
        let typ = &ai.raw_type;
        match ai.is_flag {
            true => quote! { let mut #ident: bool = false; },
            false => quote! { let mut #ident: Option<#typ> = None; }
        }
    });

    let argpostprocess = args.iter().filter(|ai| !ai.is_optional && !ai.is_flag).map(|ai| {
        let ident = &ai.funcarg_ident;
        let error_info = format!("argparse: {} is required.", ident);
        quote! {
            let #ident = match #ident {
                Some(#ident) => #ident,
                None => { return molt_ng::molt_err!(#error_info); }
            };
        }
    });

    let argfixmatch = args.iter().filter(|ai| ai.keyarg_name.is_none()).enumerate().map(|(i, ai)| {
        let ident = &ai.funcarg_ident;
        let value_converter = &ai.value_converter;
        quote! {
            #i => {
                if let Some(_v) = #value_converter {
                    #ident = Some(_v);
                    _molt_argparse_i += 1;
                    continue;
                }
            }
        }
    });

    let argkeymatch = args.iter().filter(|ai| ai.keyarg_name.is_some()).map(|ai| {
        let keyarg_name = ai.keyarg_name.as_ref().unwrap().as_str();
        let ident = &ai.funcarg_ident;
        let value_converter = &ai.value_converter;
        match ai.is_flag {
            true => quote! {
                #keyarg_name => {
                    #ident = true;
                    Ok(true)
                }
            },
            false => quote! {
                #keyarg_name => {
                    if let Some(_molt_v) = _molt_argparse_vals.next() {
                        if let Some(_v) = #value_converter {
                            #ident = Some(_v);
                            Ok(true)
                        }
                        else {
                            // although not MoltResult, we can still
                            // use this macro as it expands to Err(..).
                            return molt_ng::molt_err!(
                                "argparse: {} does not accept {}",
                                #keyarg_name, _molt_v.as_str());
                        }
                    }
                    else {
                        return molt_ng::molt_err!(
                            "argparse: expected value after {}",
                            #keyarg_name);
                    }
                }
            }
        }
    });

    let orig_args = args.iter().map(|ai| ai.funcarg_ident.clone());
    let orig_arg_defs = args.iter().map(|ai| ai.funcarg_def.clone());
    let orig_stmts = &decl.block.stmts;
    let (orig_ret_type, is_ret_empty) = match &decl.sig.output {
        syn::ReturnType::Default => (quote!{ () }, true),
        syn::ReturnType::Type(_, typ) => (typ.to_token_stream().into(), false)
    };
    let (orig_raw_rettype, orig_ret_is_option) = {
        if let Some(v) = tokens_match(
            &orig_ret_type, quote!{ Option< }, quote!{ > }
        ) {
            (v, true)
        }
        else { (orig_ret_type.clone(), false) }
    };

    let retstmt = {
        if is_ret_empty {
            quote!{ molt_ng::Value::empty() }
        }
        else if tokens_equal(&orig_raw_rettype, quote!{ &str }) {
            quote!{ molt_ng::Value::from(_ret) }
        }
        else if tokens_equal(&orig_raw_rettype, quote!{ String }) {
            quote!{ molt_ng::Value::from(_ret) }
        }
        else if let Some(t) = tokens_match(
            &orig_raw_rettype, quote!{ Vec< }, quote!{ > }
        ) {
            quote!{ molt_ng::Value::from_other(
                molt_argparse::WrapVec<#t>(_ret)) }
        }
        else {
            quote!{ molt_ng::Value::from_other(_ret) }
        }
    };
    let retstmt = match orig_ret_is_option {
        true => quote!{
            match _ret {
                Some(_ret) => Ok(#retstmt),
                None => Err(molt_ng::Exception::molt_err(molt_ng::Value::empty()))
            }
        },
        false => quote!{ Ok(#retstmt) }
    };

    let parser = quote!{
        fn parser_function(_interp: &mut molt_ng::Interp, _ctxid: molt_ng::ContextID, _argv: &[molt_ng::Value]) -> molt_ng::MoltResult {
            #contextdecl
            #(#argdecls)*
            
            let mut _molt_argparse_i: usize = 0;
            let mut _molt_argparse_vals = _argv.iter();
            let _ = _molt_argparse_vals.next().unwrap();  // skip command name
            let mut _molt_f_parse_key =
                |_molt_key: &str, _molt_argparse_vals: &mut std::slice::Iter<molt_ng::Value>| ->
                std::result::Result<bool, molt_ng::Exception>
            {
                match _molt_key {
                    #(#argkeymatch,)*
                    _ => Ok(false)
                }
            };
            while let Some(_molt_v) = _molt_argparse_vals.next() {
                if let Some(_molt_v_str) = _molt_v.try_as_str() {
                    if _molt_f_parse_key(_molt_v_str, &mut _molt_argparse_vals)? {
                        continue;
                    }
                }
                match _molt_argparse_i {
                    #(#argfixmatch,)*
                    _ => { }
                }
                let s = _molt_v.as_str();
                if !_molt_f_parse_key(s, &mut _molt_argparse_vals)? {
                    return molt_ng::molt_err!("argparse: stray {:?}", s);
                }
            }
            #(#argpostprocess)*

            let _ret = (|#(#orig_arg_defs),*| -> #orig_ret_type {
                #(#orig_stmts)*
            })(#(#orig_args),*);
            #retstmt
        }
    };
    // println!("orig parser: {}", parser);

    let mut decl = decl;
    let mut parser = syn::parse2::<syn::ItemFn>(parser).unwrap();

    std::mem::swap(&mut decl.sig.inputs, &mut parser.sig.inputs);
    std::mem::swap(&mut decl.sig.output, &mut parser.sig.output);
    std::mem::swap(&mut decl.block, &mut parser.block);

    let decl = decl.to_token_stream();
    // println!("rebuilt code: {}", decl);
    decl.into()
}