slave_codegen.rs

#![crate_type="dylib"]
#![feature(quote, plugin_registrar, rustc_private)]

extern crate slave;
extern crate proc_macro;
extern crate rustc_plugin;
extern crate syntax;

use rustc_plugin::Registry;

use syntax::ptr::P;
use syntax::ext::base::{SyntaxExtension, ExtCtxt, Annotatable};
use syntax::ext::quote::rt::Span;
use syntax::codemap::{self, Spanned};
use syntax::tokenstream::TokenTree;
use syntax::parse::token;
use syntax::ast;
use syntax::ast::*;

// TODO: Write proper required Worker struct fields message
//const REQUIRED_FIELDS_MSG: &str = "To be able to derive from Worker, struct needs to have ...";

#[plugin_registrar]
pub fn plugin_registrar(reg: &mut Registry) {
    reg.register_syntax_extension(Name::intern("worker_interface"),
                                  SyntaxExtension::MultiModifier(Box::new(worker_interface)));
}

#[allow(unused_variables)]
fn worker_interface(mut ecx: &mut ExtCtxt, sp: Span, meta_item: &MetaItem, annotated: Annotatable) -> Vec<Annotatable> {
    let struct_name = match meta_item.meta_item_list() {
        Some(ref items) => match &items[0] {
            &Spanned { ref node, .. } => match node {
                &NestedMetaItemKind::MetaItem( MetaItem { name, .. } ) => name.clone(),
                _ => panic!("Valid usage is #[worker_interface(StructName)]")
            },
            _ => panic!("Valid usage is #[worker_interface(StructName)]")
        },
        _ => panic!("Valid usage is #[worker_interface(StructName)]")
    };

    let struct_name = Ident::with_empty_ctxt(struct_name);

    let mut item = annotated.expect_item().unwrap();
    let mut commands: Vec<(MethodSig, Ident)> = vec![];
    let mut messages: Vec<(MethodSig, Ident)> = vec![];
    let impl_ty;
    
    item.node = if let ItemKind::Impl(unsafety, impl_polarity, defaultness, generics, trait_ref, ty, impl_items) = item.node {
        impl_ty = ty.clone();

        let new_impl_items = impl_items.clone().into_iter().map(|impl_item| {
            match impl_item.node {
                ImplItemKind::Method(ref sig, ref block) if is_command(&impl_item) || is_message(&impl_item) => {
                    let mut new_impl_item = impl_item.clone();
                    remove_attrs(&mut new_impl_item);

                    if is_command(&impl_item) {
                        commands.push((sig.clone(), impl_item.ident.clone()));
                    } else if is_message(&impl_item) {
                        messages.push((sig.clone(), impl_item.ident.clone()));
                    }

                    new_impl_item
                },

                _ => impl_item
            }
        }).collect::<Vec<_>>();

        ItemKind::Impl(unsafety, impl_polarity, defaultness, generics, trait_ref, ty, new_impl_items)
    } else {
        panic!();
    };

    let command_variants = commands.iter().chain(messages.iter()).flat_map(|cmd| create_command_variant(&mut ecx, &cmd.0, &cmd.1)).collect::<Vec<TokenTree>>();

    let message_variants = messages.iter().flat_map(|cmd| create_message_variant(&mut ecx, &cmd.0, &cmd.1)).collect::<Vec<TokenTree>>();

    let worker_handle = create_worker_handle(&mut ecx, &commands, &messages);

    let message_dispatcher = create_message_dispatcher(&mut ecx, &commands, &messages);

    let name = format!("{:?}", impl_ty);
    let item_ = quote_item!(
            ecx,
            pub mod handle {
                use slave::Worker;
                use std::sync::mpsc::{Sender, SyncSender};
                use std::thread::{self, JoinHandle};
                use super::$struct_name;

                #[derive(Debug)]
                pub enum Command {
                    $command_variants
                }
                    
                #[derive(Debug)]
                pub enum Message {
                    $message_variants
                }

                $worker_handle

                impl Worker for $struct_name {
                    type Handle = Handle;

                    fn name(&self) -> &str {
                        $name
                    }

                    fn run(mut self: Self) -> JoinHandle<()> {
                        thread::spawn(move || {
                            loop {
                                match self.receiver.recv() {
                                    $message_dispatcher
                                    _ => ()
                                }
                            }
                        })
                    }

                    fn get_handle(&self) -> Self::Handle {
                        Handle { sender: self.sender.clone() }
                    }
                }
            }
        );


    vec![Annotatable::Item(P(item)), Annotatable::Item(item_.unwrap())]
}

fn create_worker_handle(mut ecx: &mut ExtCtxt, commands: &Vec<(MethodSig, Ident)>, messages: &Vec<(MethodSig, Ident)>) -> Vec<TokenTree> {
    let handle_methods = commands.iter().chain(messages.iter()).flat_map(|cmd| {
        create_handle_method(&mut ecx, &cmd.0, &cmd.1)
    }).collect::<Vec<_>>();

    quote_tokens!(
        ecx,
        #[derive(Clone)]
        pub struct Handle {
            sender: SyncSender<Command>
        }

        impl Handle {
            $handle_methods
        }
    )
}

fn create_message_dispatcher(ecx: &mut ExtCtxt, commands: &[(MethodSig, Ident)], messages: &[(MethodSig, Ident)]) -> Vec<TokenTree> {
    let mut command_dispatchers: Vec<TokenTree> = commands.iter().flat_map(|&(ref sig, ref ident)| {
        let fn_name = Ident::from_str(&ident.name.as_str());
        let cmd_name = Ident::from_str(format!("Cmd_{}", fn_name).as_str());
        let args = args_to_tokens(sig.decl.inputs.iter().skip(1));
        quote_tokens!(ecx, Ok(Command::$cmd_name($args)) => self.$fn_name($args),)
    }).collect();

    let message_dispatcher: Vec<TokenTree> = messages.iter().flat_map(|&(ref sig, ref ident)| {
        let fn_name = Ident::from_str(&ident.name.as_str());
        let cmd_name = Ident::from_str(format!("Cmd_{}", fn_name).as_str());
        let args = args_to_tokens(sig.decl.inputs.iter().skip(1));

        quote_tokens!(ecx, 
                      Ok(Command::$cmd_name($args sender)) => {
                          let res = self.$fn_name($args);
                          sender.send(res);
                      })
    }).collect();

    command_dispatchers.extend(message_dispatcher.into_iter());
    command_dispatchers
}

fn create_handle_method(mut ecx: &mut ExtCtxt, sig: &MethodSig, ident: &Ident) -> Vec<TokenTree> {
    let arguments = &sig.decl.inputs.iter().skip(1).flat_map(|arg| {
        quote_tokens!(ecx, $arg,)
    }).collect::<Vec<_>>();

    let ret_ty = sig.decl.output.clone();

    let block = match ret_ty {
        FunctionRetTy::Ty(_) => create_handle_method_block_message(&mut ecx, &sig, &ident),
        _                    => create_handle_method_block_command(&mut ecx, &sig, &ident)
    };

    
    let ret = match ret_ty {
        FunctionRetTy::Ty(ref ty) => quote_tokens!(ecx, -> Result<$ty, ()>),
        _                         => vec![]
    };

    quote_tokens!(
        ecx,
        pub fn $ident(&self, $arguments) $ret {
            $block
        }
    )
}

fn arg_to_ident(arg: &Arg) -> Ident {
    match arg.pat.node {
        PatKind::Ident(_, ident, _) => ident.node,
        _                           => panic!()
    }
}

fn args_to_tokens<'a, T>(args: T) -> Vec<TokenTree> 
    where T : Iterator<Item = &'a Arg>
{
    let args = args.into_iter()
        .map(arg_to_ident)
        .map(|ident| vec![token::Ident(ident.clone())])
        .collect::<Vec<_>>();

    args.as_slice().join(&token::Comma)
        .into_iter()
        .map(|t| TokenTree::Token(codemap::DUMMY_SP, t))
        .collect::<Vec<_>>()
}

fn create_handle_method_block_command(ecx: &mut ExtCtxt, sig: &MethodSig, ident: &Ident) -> Vec<TokenTree> {
    let cmd_name = Ident::from_str(format!("Cmd_{}", ident.name).as_str());
    let args = args_to_tokens(sig.decl.inputs.iter().skip(1));

    quote_tokens!(ecx, self.sender.send(Command::$cmd_name($args));)
}

fn create_handle_method_block_message(ecx: &mut ExtCtxt, sig: &MethodSig, ident: &Ident) -> Vec<TokenTree> {
    let cmd_name = Ident::from_str(format!("Cmd_{}", ident.name).as_str());
    let msg_name = Ident::from_str(format!("Msg_{}", ident.name).as_str());

    let args = args_to_tokens(sig.decl.inputs.iter().skip(1));
    let ret_ty = match sig.decl.output {
        FunctionRetTy::Ty(ref ty) => ty.clone().unwrap(),
        _ => panic!("No return type for message")
    };

    quote_tokens!(
        ecx,
        use std::sync::mpsc;
        let (sender, receiver) = mpsc::sync_channel::<$ret_ty>(1000);
        self.sender.send(Command::$cmd_name($args sender.clone()));

        match receiver.recv() {
            Ok(res) => Ok(res),
            _       => Err(())
        }
    )
}

fn create_command_variant(ecx: &mut ExtCtxt, sig: &MethodSig, ident: &Ident) -> Vec<TokenTree> {
    let variant_name = Ident::from_str(format!("Cmd_{}", ident.name).as_str());
    let mut arguments = sig.decl.inputs.clone().into_iter().map(|arg| {
        arg.ty.unwrap()
    }).skip(1).collect::<Vec<_>>();

    if let FunctionRetTy::Ty(ref ty) = sig.decl.output {
        let ty = ty.clone().unwrap();
        let ty = quote_ty!(ecx, SyncSender<$ty>).unwrap();

        arguments.push(ty);
    }

    quote_tokens!(
        ecx,
        $variant_name($arguments),
    )
}

fn create_message_variant(ecx: &mut ExtCtxt, sig: &MethodSig, ident: &Ident) -> Vec<TokenTree> {
    let variant_name = Ident::from_str(format!("Msg_{}", ident.name).as_str());
    let argument = match sig.decl.output {
        FunctionRetTy::Ty(ref ty) => ty.clone().unwrap(),
        _ => panic!("Messages should return a value")
    };

    quote_tokens!(
        ecx,
        $variant_name($argument)
    )
}

fn remove_attrs(item: &mut ImplItem) {
    item.attrs = item.attrs.clone().into_iter().filter(|attr| attr.path != "command" && attr.path != "message").collect();
}

fn is_command(item: &ImplItem) -> bool {
    item.attrs.iter().find(|attr| attr.path == "command").is_some()
}

fn is_message(item: &ImplItem) -> bool {
    item.attrs.iter().find(|attr| attr.path == "message").is_some()
}