dir-embed 0.1.0

Like include_bytes! for directories
Documentation
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()
            }
        }
    }
}