bumpikins/lib/src/lib.rs

#![feature(proc_macro_quote)]
extern crate proc_macro2;
use proc_macro2::*;

extern crate proc_quote;
use proc_quote::quote;

fn scrub_storage_macro(stream: TokenStream) -> TokenStream {
    let mut just_saw_hash = false;
    let mut without_storage = Vec::<TokenTree>::new();
    for item in stream.into_iter() {
        match item {
            TokenTree::Punct(ref p) if p.as_char() == '#' => just_saw_hash = true,
            TokenTree::Group(g) if just_saw_hash && g.delimiter() == Delimiter::Bracket => {
                just_saw_hash = false;
            }
            TokenTree::Group(g) => {
                let mut new_group = Group::new(g.delimiter(), scrub_storage_macro(g.stream()));
                new_group.set_span(g.span());
                without_storage.push(TokenTree::Group(new_group));
                just_saw_hash = false;
            }
            _ => {
                without_storage.push(item);
                just_saw_hash = false;
            }
        }
    }

    TokenStream::from_iter(without_storage)
}

#[proc_macro_attribute]
pub fn bump(
    attr: proc_macro::TokenStream,
    stream: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let attr: TokenStream = attr.into();
    let stream: TokenStream = stream.into();
    let refname = attr.into_iter().next().unwrap();

    let mut process = stream.clone().into_iter();
    if let Some(TokenTree::Ident(i)) = process.next() {
        if i.to_string() != "struct" {
            let stream = quote!(compile_error!(stringify!("The `bump` macro can only be used on `struct`s")););
            let stream = TokenStream::from_iter(stream.into_iter().map(|mut s| {
                s.set_span(i.span());
                s
            }));
            return stream.into();
        }
    }
    let struct_name = process.next().unwrap();

    let body = process.find(|s| match s {
        TokenTree::Group(g) if g.delimiter() == Delimiter::Brace => true,
        _ => false,
    });
    let body = if let Some(TokenTree::Group(b)) = body {
        b
    } else {
        return quote!(compile_error!("Unable to find the struct body");).into();
    };

    let storage_name;
    let storage_type;
    let mut body = body.stream().into_iter();
    loop {
        let item = if let Some(i) = body.next() {
            i
        } else {
            return quote!(compile_error!("Could not find `#[storage]`")).into();
        };

        if let TokenTree::Punct(p) = item {
            if p.as_char() != '#' {
                continue;
            }

            // storage macro
            // TODO: check
            let _ = body.next();

            let mut name = match body.next() {
                Some(TokenTree::Ident(i)) => i,
                _ => continue,
            };
            if name.to_string() == "pub" {
                name = match body.next() {
                    Some(TokenTree::Ident(i)) => i,
                    _ => continue,
                }
            }
            let _colon = body.next();
            let _vec = body.next();

            let _lbrac = body.next();
            let mut typ = Vec::new();
            while let Some(n) = body.next() {
                if let TokenTree::Punct(ref p) = n {
                    if p.as_char() == '>' {
                        break;
                    }
                }
                typ.push(n);
            }
            storage_name = name;
            storage_type = typ[0].clone();
            break;
        } else {
            continue;
        }
    }

    let idx_def = quote! {
        #[derive(Debug, PartialEq, Eq, Clone, Copy)]
        struct #refname {
            idx: usize,
        }
    };

    let impl_def = quote! {
        impl #struct_name {
            fn insert(&mut self, item: #storage_type) -> #refname {
                let idx = self.#storage_name.len();
                self.#storage_name.push(item);
                #refname { idx }
            }
        }

        impl ::std::ops::Index<#refname> for #struct_name {
            type Output = #storage_type;

            fn index(&self, rf: #refname) -> &Self::Output {
                &self.#storage_name[rf.idx]
            }
        }
    };

    let mut stream = scrub_storage_macro(stream);
    stream.extend(idx_def);
    stream.extend(impl_def);

    stream.into()
}