embrace the jank

.gitignore

@@ -0,0 +1,2 @@
+/target
+/Cargo.lock

Cargo.toml

@@ -0,0 +1,6 @@
+[workspace]
+
+members = [
+  "lib",
+  "test",
+]

lib/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "bumpikins"
+version = "0.1.0"
+edition = "2021"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro2 = "*"
+proc-quote = "*"

lib/src/lib.rs

@@ -0,0 +1,145 @@
+#![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()
+}

test/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "bumpikins-test"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+bumpikins = { path = "../lib" }

test/src/lib.rs

@@ -0,0 +1,27 @@
+#[cfg(test)]
+mod tests {
+    use bumpikins::bump;
+
+    #[bump(ExampleRef)]
+    struct Example {
+        #[storage]
+        thingy: Vec<u32>,
+    }
+
+    #[test]
+    fn it_works() {
+        let mut example = Example { thingy: Vec::new() };
+
+        let i1 = example.insert(5);
+        assert_eq!(i1, ExampleRef { idx: 0 });
+        assert_eq!(example[i1], 5);
+
+        let i2 = example.insert(7);
+        assert_eq!(i2, ExampleRef { idx: 1 });
+        assert_eq!(example[i2], 7);
+
+        let i3 = example.insert(9);
+        assert_eq!(i3, ExampleRef { idx: 2 });
+        assert_eq!(example[i3], 9);
+    }
+}