fn_mut 0.1.0

//! # Usage
//!
//! ```rust,ignore
//!
//! #[fn_mut(disable_self, disable_output, enable_attrs = "attr1,attr2,...")]
//! fn(attr1: T1, attr2: T2, ...) -> OutT { ... }
//!
//! ```
//!
//! By default `fn_mut` macro generates function which takes mutable reference to
//! self and returns mutable reference, if in original function there is any
//! reference.
//!
//! Possible options are:
//!
//! - `disable_self`: do not change `&self` to `&mut self`
//! - `disable_output`: do not change  output value
//! - `enable_attrs = "attr1,attr2,..."`: change mutability in attr
//!
//! Requires nightly compiler and `#![feature(proc_macro)]`.
//!
//! # Examples
//!
//! ```rust,ignore
//!
//! #![feature(proc_macro)]
//!
//! extern crate fn_mut;
//! use fn_mut::fn_mut;
//!
//! struct Test(u64);
//!
//! impl Test {
//!     #[fn_mut(enable_attrs = "text")]
//!     fn test(&self, text: &str) -> Option<&u64> {
//!         if_mut! {
//!             println!("This is mut fn: {}", text);
//!         }
//!         if_const! {
//!             println!("This is const fn: {}", text);
//!         }
//!         Some(ptr!(self.0))
//!     }
//! }
//!
//! ```
//!
//! This example expands to:
//!
//! ```rust
//!
//! struct Test(u64);
//!
//! impl Test {
//!     fn test(&self, text: &str) -> Option<&u64> {
//!         println!("This is const fn: {}", text);
//!         Some(&self.0)
//!     }
//!     fn test_mut(&mut self, text: &mut str) -> Option<&mut u64> {
//!         println!("This is mut fn: {}", text);
//!         Some(&mut self.0)
//!     }
//! }
//!
//! ```

#![feature(proc_macro)]

extern crate proc_macro;
use proc_macro::TokenStream;

extern crate syn;

#[macro_use]
extern crate quote;

extern crate regex;

#[macro_use]
extern crate lazy_static;

mod config;
use config::*;

#[proc_macro_attribute]
pub fn fn_mut(attrs: TokenStream, input: TokenStream) -> TokenStream {
    let attrs = attrs.to_string();
    let config = get_config(&attrs);

    let s = input.to_string();
    let ast = syn::parse_item(&s).unwrap();
    let gen = gen_mut_fn(ast, config);
    gen.parse().unwrap()
}

fn into_mut_ty(ty: &mut syn::Ty) {
    match ty {
        &mut syn::Ty::Ptr(ref mut mut_ty) |
         &mut syn::Ty::Rptr(_, ref mut mut_ty) => {
            mut_ty.mutability = syn::Mutability::Mutable;
        }
        &mut syn::Ty::Array(ref mut inner, ..) |
         &mut syn::Ty::Slice(ref mut inner) |
         &mut syn::Ty::Paren(ref mut inner) => return into_mut_ty(inner),
         &mut syn::Ty::Tup(ref mut vec) => {
            for ty in vec.iter_mut() {
                into_mut_ty(ty);
            }
        }
        &mut syn::Ty::Path(ref mut qualified, ref mut path) => {
            if let &mut Some(ref mut qualified) = qualified {
                // TODO: should this be implemented?
                unimplemented!()
            }
            for seg in path.segments.iter_mut() {
                match seg.parameters {
                    syn::PathParameters::AngleBracketed(ref mut param) => {
                        for ty in param.types.iter_mut() {
                            into_mut_ty(ty);
                        }
                    },
                    _ => ()
                }
            }
        }
        _ => ()
    };
}

fn add_code(ast: &mut syn::Item, is_mut: bool) {
    let code = if is_mut {
        quote! {
            fn code() {
                macro_rules! if_mut {
                    { $( $code:tt )* } => (
                        $( $code )*
                    )
                }
                macro_rules! if_const {
                    { $( $code:tt )* } => ()
                }
                macro_rules! ptr {
                    ($v:expr) => (
                        &mut $v
                    )
                }
                macro_rules! raw_ptr {
                    ($v:expr) => (
                        *mut $v
                    )
                }
            }
        }
    } else {
        quote! {
            fn code() {
                macro_rules! if_const {
                    { $( $code:tt )* } => (
                        $( $code )*
                    )
                }
                macro_rules! if_mut {
                    { $( $code:tt )* } => ()
                }
                macro_rules! ptr {
                    ($v:expr) => (
                        &$v
                    )
                }
                macro_rules! raw_ptr {
                    ($v:expr) => (
                        *const $v
                    )
                }
            }
        }
    };
    let code = code.into_string();
    let syn_code = syn::parse_item(&code).unwrap();
    let mut stmts = match syn_code.node {
        syn::ItemKind::Fn(_, _, _, _, _, block) => {
            block.stmts
        }
        _ => unreachable!()
    };
    match ast.node {
        syn::ItemKind::Fn(_, _, _, _, _, ref mut block) => {
            stmts.extend_from_slice(&block.stmts);
            block.stmts = stmts;
        },
        _ => ()
    }
}

fn gen_mut_fn(ast: syn::Item, config: Config) -> quote::Tokens {
    let mut const_fn = ast.clone();
    add_code(&mut const_fn, false);
    let mut mut_fn = ast;
    add_code(&mut mut_fn, true);
    mut_fn.ident = syn::Ident::from(format!("{}_mut", &mut_fn.ident).as_ref());
    match mut_fn.node {
        syn::ItemKind::Fn(ref mut header, ..) => {
            if header.variadic {
                panic!("#[mut_fn] does not support variadic functions");
            }
            if config.enable_output {
                match header.output {
                    syn::FunctionRetTy::Default => (),
                    syn::FunctionRetTy::Ty(ref mut t) => {
                        into_mut_ty(t);
                    }
                }
            }
            for input in header.inputs.iter_mut() {
                match input {
                    &mut syn::FnArg::SelfRef(ref _lifetime, ref mut mutability) => if config.enable_self {
                        *mutability = syn::Mutability::Mutable
                    },
                    &mut syn::FnArg::SelfValue(ref mut mutability) => if config.enable_self {
                        *mutability = syn::Mutability::Mutable
                    },
                    &mut syn::FnArg::Captured(ref pat, ref mut ty) => {
                        match pat {
                            &syn::Pat::Ident(_, ref ident, _) => {
                                if config.enable_attrs.contains(ident.as_ref()) {
                                    into_mut_ty(ty)
                                }
                            },
                            _ => ()
                        }
                    },
                    _ => ()
                }
            }
        },
        _ => panic!("#[mut_fn] used on non-fn type")
    }
    quote! {
        #const_fn
        #mut_fn
    }
}