Skip to main content

rust_silos_macros/
lib.rs

1//! Proc-macro for rust-silos: generates a sorted array of entries for fast binary search lookup.
2
3extern crate proc_macro;
4use proc_macro::TokenStream;
5use quote::{quote, quote_spanned};
6use std::fs;
7use std::path::Path;
8use syn::{
9    parse::{Parse, ParseStream},
10    parse_macro_input, LitStr, Token,
11};
12use walkdir::WalkDir;
13
14type EmbedMeta = (String, String, usize, u64);
15type CollectResult = (Vec<EmbedMeta>, Vec<proc_macro2::TokenStream>);
16
17/// Internal: Macro input parser for `silo!` macro. Accepts a path and optional force argument.
18/// Path must be a string literal. Force is a bool literal.
19struct SiloMacroInput {
20    path: LitStr,
21    force: Option<(syn::Ident, syn::LitBool)>,
22    crate_path: Option<syn::Path>,
23}
24
25/// Parse implementation for macro input. Handles path and optional force argument.
26impl Parse for SiloMacroInput {
27    fn parse(input: ParseStream) -> syn::Result<Self> {
28        let path: LitStr = input.parse()?;
29        let mut force = None;
30        let mut crate_path = None;
31        while input.peek(Token![,]) {
32            input.parse::<Token![,]>()?;
33            let ident: syn::Ident = input.parse()?;
34            input.parse::<Token![=]>()?;
35            if ident == "force" {
36                let value: syn::LitBool = input.parse()?;
37                force = Some((ident, value));
38            } else if ident == "crate" {
39                let path: syn::Path = input.parse()?;
40                crate_path = Some(path);
41            } else {
42                return Err(syn::Error::new(ident.span(), "Unknown argument to embed_silo!"));
43            }
44        }
45        Ok(SiloMacroInput { path, force, crate_path })
46    }
47}
48
49/// Macro to embed all files in a directory as a sorted array for fast, allocation-free binary search access.
50///
51/// Usage: `let silo = embed_silo!("assets");` or `let silo = embed_silo!("assets", force = true);`
52/// In debug mode, uses dynamic loading unless `force = true`.
53/// Directory path must exist at build time for embedding.
54#[proc_macro]
55pub fn embed_silo(input: TokenStream) -> TokenStream {
56    let SiloMacroInput { path, force, crate_path } = parse_macro_input!(input as SiloMacroInput);
57    let dir_path = path.value();
58    let call_span = path.span();
59    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap_or_else(|_| String::new());
60    if manifest_dir.is_empty() {
61        return compile_error("embed_silo!: CARGO_MANIFEST_DIR not set", call_span);
62    }
63    let manifest_dir_canon = match Path::new(&manifest_dir).canonicalize() {
64        Ok(p) => p,
65        Err(_) => return compile_error("embed_silo!: failed to resolve CARGO_MANIFEST_DIR", call_span),
66    };
67
68    let abs_path = manifest_dir_canon.join(&dir_path);
69    let abs_path = match abs_path.canonicalize() {
70        Ok(p) => p,
71        Err(_) => {
72            return compile_error(
73                format!("embed_silo!: failed to resolve path: {}", dir_path),
74                call_span,
75            )
76        }
77    };
78    let abs_path_str = match abs_path.to_str() {
79        Some(p) => p,
80        None => return compile_error("embed_silo!: path must be valid UTF-8", call_span),
81    };
82
83    // Path-safe containment check (avoid prefix-string bugs like /foo/bar matching /foo/bar2).
84    if !abs_path.starts_with(&manifest_dir_canon) {
85        let msg = format!(
86            "embed_silo!: directory not found:\n  {}\n  expected to be inside crate root:\n  {}\n  relative path: {}",
87            abs_path_str,
88            manifest_dir_canon.display(),
89            dir_path
90        );
91        return compile_error(&msg, call_span);
92    }
93
94    let force_embed = force.as_ref().is_some_and(|(_, v)| v.value());
95    let crate_root = crate_path
96        .map(|p| quote! { #p })
97        .unwrap_or_else(|| quote! { ::rust_silos });
98
99    // Keep a stable absolute root for dynamic fallback and for `into_dynamic()` conversions.
100    let abs_root_lit = syn::LitStr::new(abs_path_str, call_span);
101
102    // Always collect entries (needed for both embedded and cfg-gated code generation)
103    let (entries, errors) = collect_embed_entries(abs_path_str, call_span);
104    if !errors.is_empty() {
105        return quote! { #(#errors)* }.into();
106    }
107
108    // Generate unique identifier using hash + length for better collision resistance
109    let mut hasher = std::collections::hash_map::DefaultHasher::new();
110    use std::hash::{Hash, Hasher};
111    abs_path_str.hash(&mut hasher);
112    let hash = hasher.finish();
113    let array_ident = quote::format_ident!("__EMBED_ARRAY_{:x}_{}", hash, abs_path_str.len());
114
115    let array_entries = generate_sorted_array(&entries, &crate_root);
116    let entry_count = entries.len();
117
118    if force_embed {
119        // Force embed - always use embedded data regardless of debug/release
120        let expanded = quote! {
121            {
122                static #array_ident: [(&str, #crate_root::EmbedEntry); #entry_count] = [
123                    #array_entries
124                ];
125                #crate_root::Silo::from_embedded(&#array_ident, #abs_root_lit)
126            }
127        };
128        expanded.into()
129    } else {
130        // Let consumer's cfg!(debug_assertions) decide at their compile time
131        let expanded = quote! {
132            {
133                #[cfg(debug_assertions)]
134                let __silo = #crate_root::Silo::from_static(#abs_root_lit);
135
136                #[cfg(not(debug_assertions))]
137                let __silo = {
138                    static #array_ident: [(&str, #crate_root::EmbedEntry); #entry_count] = [
139                        #array_entries
140                    ];
141                    #crate_root::Silo::from_embedded(&#array_ident, #abs_root_lit)
142                };
143                __silo
144            }
145        };
146        expanded.into()
147    }
148}
149
150/// Recursively collects all files in the given directory for embedding.
151/// Returns (entries, errors):
152///   - entries: Vec<(relative_path, abs_path, size, modified)>
153///   - errors: Vec<TokenStream> for compile_error!s
154fn collect_embed_entries(dir: &str, span: proc_macro2::Span) -> CollectResult {
155    let mut entries = Vec::new();
156    let mut errors = Vec::new();
157    let root = Path::new(dir);
158    for entry in WalkDir::new(root).into_iter() {
159        let entry = match entry {
160            Ok(e) => e,
161            Err(e) => {
162                let msg = format!("embed_silo!: failed to read entry: {}", e);
163                errors.push(quote_spanned! {span=> compile_error!(#msg); });
164                continue;
165            }
166        };
167        if entry.file_type().is_file() {
168            let path = entry.path();
169            let rel_path = match path.strip_prefix(root) {
170                Ok(r) => r.to_string_lossy().replace('\\', "/"),
171                Err(_) => {
172                    let msg = "embed_silo!: failed to get relative path";
173                    errors.push(quote_spanned! {span=> compile_error!(#msg); });
174                    continue;
175                }
176            };
177            let abs_path = match path.canonicalize() {
178                Ok(p) => p.to_string_lossy().to_string(),
179                Err(_) => {
180                    let msg = format!("embed_silo!: failed to canonicalize file: {}", path.display());
181                    errors.push(quote_spanned! {span=> compile_error!(#msg); });
182                    continue;
183                }
184            };
185            // Single metadata call for both size and modified time
186            let meta = fs::metadata(path).ok();
187            let size = meta.as_ref().map(|m| m.len() as usize).unwrap_or(0);
188            let modified = meta
189                .and_then(|m| m.modified().ok())
190                .and_then(|mtime| mtime.duration_since(std::time::UNIX_EPOCH).ok())
191                .map(|d| d.as_secs())
192                .unwrap_or(0);
193            entries.push((rel_path, abs_path, size, modified));
194        }
195    }
196
197    // Make builds more reproducible across platforms/filesystems.
198    entries.sort_by(|(a, _, _, _), (b, _, _, _)| a.cmp(b));
199    (entries, errors)
200}
201
202// emit_compile_error removed; use quote_spanned! inline instead
203
204/// Emit compile_error! and return from macro expansion.
205fn compile_error<S: AsRef<str>>(msg: S, span: proc_macro2::Span) -> proc_macro::TokenStream {
206    let lit = syn::LitStr::new(msg.as_ref(), span);
207    let tokens = quote!(compile_error!(#lit));
208    tokens.into()
209}
210
211/// Generates a sorted array of entries for binary search lookup.
212/// Used internally by the macro. Expects (rel_path, abs_path, size, modified) tuples.
213/// Entries must already be sorted by rel_path.
214fn generate_sorted_array(entries: &[EmbedMeta], crate_root: &proc_macro2::TokenStream) -> proc_macro2::TokenStream {
215    let pairs = entries.iter().map(|(rel_path, abs_path, size, modified)| {
216        let rel_path_lit = syn::LitStr::new(rel_path, proc_macro2::Span::call_site());
217        let abs_path_lit = syn::LitStr::new(abs_path, proc_macro2::Span::call_site());
218        let size_lit = syn::LitInt::new(&size.to_string(), proc_macro2::Span::call_site());
219        let mod_lit = syn::LitInt::new(&modified.to_string(), proc_macro2::Span::call_site());
220        quote! {
221            (#rel_path_lit, #crate_root::EmbedEntry {
222                path: #rel_path_lit,
223                contents: include_bytes!(#abs_path_lit),
224                size: #size_lit,
225                modified: #mod_lit,
226            }),
227        }
228    });
229    quote! {
230        #(#pairs)*
231    }
232}