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#[derive(Debug, Clone, Copy)]
10enum EmbedMode {
11    Bytes,
12    Str,
13}
14
15#[proc_macro_derive(Embed, attributes(dir, mode))]
16pub fn embed(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
17    let input = parse_macro_input!(input as DeriveInput);
18
19    let struct_name = &input.ident;
20
21    let dir_attr = input
22        .attrs
23        .iter()
24        .find(|e| e.path().is_ident("dir"))
25        .expect("No #[dir = \"...\"] attribute found");
26
27    let mode_attr = input.attrs.iter().find(|e| e.path().is_ident("mode"));
28
29    let mode = mode_attr.map(extract_mode).unwrap_or(EmbedMode::Bytes);
30
31    let base_path = PathBuf::from(extract_dir_path(dir_attr));
32    let source_file = PathBuf::from(input.span().unwrap().file());
33
34    let source_dir = if let Some(parent) = source_file.parent() {
35        parent
36    } else {
37        // HACK: when running in rust-analyzer i can't seem to get the parent dir.
38        return TokenStream::from(generate_byte_impl(struct_name, Vec::new()));
39    };
40
41    let absolue_path = source_dir.join(&base_path);
42
43    let mut match_arms = Vec::new();
44
45    for entry in collect_files(&absolue_path) {
46        let rel_path = entry
47            .strip_prefix(&absolue_path)
48            .unwrap()
49            .to_str()
50            .unwrap()
51            .replace("\\", "/");
52
53        let include_path = base_path.join(&rel_path);
54        let include_string = include_path.to_str().unwrap();
55
56        let arm = match mode {
57            EmbedMode::Bytes => generate_byte_arm(&rel_path, include_string),
58            EmbedMode::Str => generate_str_arm(&rel_path, include_string),
59        };
60
61        match_arms.push(arm);
62    }
63
64    let expanded = match mode {
65        EmbedMode::Bytes => generate_byte_impl(struct_name, match_arms),
66        EmbedMode::Str => generate_str_impl(struct_name, match_arms),
67    };
68
69    proc_macro::TokenStream::from(expanded)
70}
71
72fn collect_files(dir: &Path) -> Vec<PathBuf> {
73    let mut files = Vec::new();
74    for entry in fs::read_dir(dir).unwrap() {
75        let path = entry.unwrap().path();
76        if path.is_file() {
77            files.push(path);
78        } else if path.is_dir() {
79            files.extend(collect_files(&path));
80        }
81    }
82    files
83}
84
85fn extract_dir_path(attr: &Attribute) -> String {
86    let meta = match &attr.meta {
87        Meta::NameValue(meta) => meta,
88        _ => panic!("Expected #[dir = \"...\"] as a name-value attribute."),
89    };
90
91    let expr_lit = match &meta.value {
92        Expr::Lit(expr_lit) => expr_lit,
93        _ => panic!("Expected #[dir = \"...\"] with a string literal."),
94    };
95
96    match &expr_lit.lit {
97        Lit::Str(str) => str.value(),
98        _ => panic!("Expected #[dir = \"...\"] to be a string."),
99    }
100}
101
102fn extract_mode(attr: &Attribute) -> EmbedMode {
103    let meta = match &attr.meta {
104        Meta::NameValue(meta) => meta,
105        _ => panic!("Expected #[mode = \"bytes\"|\"str\"] as a name-value attribute."),
106    };
107
108    let expr_lit = match &meta.value {
109        Expr::Lit(expr_lit) => expr_lit,
110        _ => panic!("Expected #[mode = \"bytes\"|\"str\"] with a string literal."),
111    };
112
113    match &expr_lit.lit {
114        Lit::Str(str) => match str.value().as_str() {
115            "bytes" => EmbedMode::Bytes,
116            "str" => EmbedMode::Str,
117            other => panic!("Unknown mode: {other}. Use `bytes` or `str`."),
118        },
119        _ => panic!("Expected #[mode = \"bytes\"|\"str\"] to be a string."),
120    }
121}
122
123fn generate_byte_arm(rel: &str, include: &str) -> proc_macro2::TokenStream {
124    quote! {
125        #rel => Some(include_bytes!(#include)),
126    }
127}
128
129fn generate_byte_impl(
130    struct_name: &Ident,
131    match_arms: Vec<proc_macro2::TokenStream>,
132) -> proc_macro2::TokenStream {
133    quote! {
134        impl #struct_name {
135            pub fn get(name: &str) -> Option<&'static [u8]> {
136                match name {
137                    #(#match_arms)*
138                    _ => None,
139                }
140            }
141        }
142    }
143}
144
145fn generate_str_arm(rel: &str, include: &str) -> proc_macro2::TokenStream {
146    quote! {
147        #rel => Some(include_str!(#include)),
148    }
149}
150
151fn generate_str_impl(
152    struct_name: &Ident,
153    match_arms: Vec<proc_macro2::TokenStream>,
154) -> proc_macro2::TokenStream {
155    quote! {
156        impl #struct_name {
157            pub fn get(name: &str) -> Option<&'static str> {
158                match name {
159                    #(#match_arms)*
160                    _ => None,
161                }
162            }
163        }
164    }
165}
dir_embed/lib.rs

dir_embed/
lib.rs
