dir_embed/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::spanned::Spanned;
4use syn::{Attribute, DeriveInput, Expr, Ident, Lit, Meta, parse_macro_input};
5
6use std::fs;
7use std::path::{Path, PathBuf};
8
9#[proc_macro_derive(Embed, attributes(dir))]
10pub fn embed(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
11    let input = parse_macro_input!(input as DeriveInput);
12
13    let struct_name = &input.ident;
14
15    let attr = input
16        .attrs
17        .iter()
18        .find(|e| e.path().is_ident("dir"))
19        .expect("No #[dir = \"...\"] attribute found");
20
21    let base_path = PathBuf::from(extract_dir_path(attr));
22    let source_file = PathBuf::from(input.span().unwrap().file());
23
24    let source_dir = if let Some(parent) = source_file.parent() {
25        parent
26    } else {
27        // HACK: when running in rust-analyzer i can't seem to get the parent dir.
28        return TokenStream::from(generate_impl(struct_name, Vec::new(), Vec::new()));
29    };
30
31    let absolue_path = source_dir.join(&base_path);
32
33    let mut match_arms = Vec::new();
34    let mut entries = Vec::new();
35
36    for entry in collect_files(&absolue_path) {
37        let rel_path = entry
38            .strip_prefix(&absolue_path)
39            .unwrap()
40            .to_str()
41            .unwrap()
42            .replace("\\", "/");
43
44        let include_path = base_path.join(&rel_path);
45        let include_string = include_path.to_str();
46
47        match_arms.push(quote! {
48            #rel_path => Some(include_bytes!(#include_string) as &'static [u8]),
49        });
50
51        entries.push(quote! {
52            (#rel_path, include_bytes!(#include_string) as &'static [u8])
53        });
54    }
55
56    let expanded = generate_impl(struct_name, match_arms, entries);
57
58    proc_macro::TokenStream::from(expanded)
59}
60
61fn collect_files(dir: &Path) -> Vec<PathBuf> {
62    let mut files = Vec::new();
63    for entry in fs::read_dir(dir).unwrap() {
64        let path = entry.unwrap().path();
65        if path.is_file() {
66            files.push(path);
67        } else if path.is_dir() {
68            files.extend(collect_files(&path));
69        }
70    }
71    files
72}
73
74fn extract_dir_path(attr: &Attribute) -> String {
75    let meta = match &attr.meta {
76        Meta::NameValue(meta) => meta,
77        _ => panic!("Expected #[dir = \"...\"] as a name-value attribute."),
78    };
79
80    let expr_lit = match &meta.value {
81        Expr::Lit(expr_lit) => expr_lit,
82        _ => panic!("Expected #[dir = \"...\"] with a string literal."),
83    };
84
85    match &expr_lit.lit {
86        Lit::Str(str) => str.value(),
87        _ => panic!("Expected #[dir = \"...\"] to be a string."),
88    }
89}
90
91fn generate_impl(
92    struct_name: &Ident,
93    match_arms: Vec<proc_macro2::TokenStream>,
94    entries: Vec<proc_macro2::TokenStream>,
95) -> proc_macro2::TokenStream {
96    quote! {
97        impl #struct_name {
98            pub fn get(name: &str) -> Option<&'static [u8]> {
99                match name {
100                    #(#match_arms)*
101                    _ => None,
102                }
103            }
104
105            pub fn iter() -> impl Iterator<Item = (&'static str, &'static [u8])> {
106                [#(#entries),*].into_iter()
107            }
108        }
109    }
110}
dir_embed/lib.rs

dir_embed/
lib.rs
