initial commit

src/lib.rs

@@ -0,0 +1,110 @@
+use proc_macro::TokenStream;
+use quote::quote;
+use syn::spanned::Spanned;
+use syn::{Attribute, DeriveInput, Expr, Ident, Lit, Meta, parse_macro_input};
+
+use std::fs;
+use std::path::{Path, PathBuf};
+
+#[proc_macro_derive(Embed, attributes(dir))]
+pub fn embed(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    let input = parse_macro_input!(input as DeriveInput);
+
+    let struct_name = &input.ident;
+
+    let attr = input
+        .attrs
+        .iter()
+        .find(|e| e.path().is_ident("dir"))
+        .expect("No #[dir = \"...\"] attribute found");
+
+    let base_path = PathBuf::from(extract_dir_path(attr));
+    let source_file = PathBuf::from(input.span().unwrap().file());
+
+    let source_dir = if let Some(parent) = source_file.parent() {
+        parent
+    } else {
+        // HACK: when running in rust-analyzer i can't seem to get the parent dir.
+        return TokenStream::from(generate_impl(struct_name, Vec::new(), Vec::new()));
+    };
+
+    let absolue_path = source_dir.join(&base_path);
+
+    let mut match_arms = Vec::new();
+    let mut entries = Vec::new();
+
+    for entry in collect_files(&absolue_path) {
+        let rel_path = entry
+            .strip_prefix(&absolue_path)
+            .unwrap()
+            .to_str()
+            .unwrap()
+            .replace("\\", "/");
+
+        let include_path = base_path.join(&rel_path);
+        let include_string = include_path.to_str();
+
+        match_arms.push(quote! {
+            #rel_path => Some(include_bytes!(#include_string) as &'static [u8]),
+        });
+
+        entries.push(quote! {
+            (#rel_path, include_bytes!(#include_string) as &'static [u8])
+        });
+    }
+
+    let expanded = generate_impl(struct_name, match_arms, entries);
+
+    proc_macro::TokenStream::from(expanded)
+}
+
+fn collect_files(dir: &Path) -> Vec<PathBuf> {
+    let mut files = Vec::new();
+    for entry in fs::read_dir(dir).unwrap() {
+        let path = entry.unwrap().path();
+        if path.is_file() {
+            files.push(path);
+        } else if path.is_dir() {
+            files.extend(collect_files(&path));
+        }
+    }
+    files
+}
+
+fn extract_dir_path(attr: &Attribute) -> String {
+    let meta = match &attr.meta {
+        Meta::NameValue(meta) => meta,
+        _ => panic!("Expected #[dir = \"...\"] as a name-value attribute."),
+    };
+
+    let expr_lit = match &meta.value {
+        Expr::Lit(expr_lit) => expr_lit,
+        _ => panic!("Expected #[dir = \"...\"] with a string literal."),
+    };
+
+    match &expr_lit.lit {
+        Lit::Str(str) => str.value(),
+        _ => panic!("Expected #[dir = \"...\"] to be a string."),
+    }
+}
+
+fn generate_impl(
+    struct_name: &Ident,
+    match_arms: Vec<proc_macro2::TokenStream>,
+    entries: Vec<proc_macro2::TokenStream>,
+) -> proc_macro2::TokenStream {
+    quote! {
+        impl #struct_name {
+            pub fn get(name: &str) -> Option<&'static [u8]> {
+                match name {
+                    #(#match_arms)*
+                    _ => None,
+                }
+            }
+
+            pub fn iter() -> impl Iterator<Item = (&'static str, &'static [u8])> {
+                [#(#entries),*].into_iter()
+            }
+        }
+    }
+}