include_zstd_derive/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro_crate::{FoundCrate, crate_name};
3use quote::quote;
4use std::fs;
5use std::path::{Path, PathBuf};
6use syn::parse::{Parse, ParseStream};
7use syn::{LitByteStr, LitStr, Token, parse_macro_input};
8
9struct FileMacroInput {
10    source_file: Option<LitStr>,
11    target_path: LitStr,
12}
13
14impl Parse for FileMacroInput {
15    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
16        let first: LitStr = input.parse()?;
17        if input.is_empty() {
18            return Ok(Self {
19                source_file: None,
20                target_path: first,
21            });
22        }
23
24        let _comma: Token![,] = input.parse()?;
25        let second: LitStr = input.parse()?;
26        if !input.is_empty() {
27            return Err(input.error("expected one string literal path or 'source_file, path'"));
28        }
29
30        Ok(Self {
31            source_file: Some(first),
32            target_path: second,
33        })
34    }
35}
36
37#[proc_macro]
38pub fn r#str(input: TokenStream) -> TokenStream {
39    let value = parse_macro_input!(input as LitStr);
40    let data = value.value().into_bytes();
41    expand_from_data(data, true)
42}
43
44#[proc_macro]
45pub fn bytes(input: TokenStream) -> TokenStream {
46    let value = parse_macro_input!(input as LitByteStr);
47    let data = value.value();
48    expand_from_data(data, false)
49}
50
51#[proc_macro]
52pub fn file_str(input: TokenStream) -> TokenStream {
53    let input = parse_macro_input!(input as FileMacroInput);
54    let source_file = input.source_file.as_ref().map(LitStr::value);
55    let source_path = input.target_path.value();
56
57    let absolute_path = match resolve_path(source_file.as_deref(), &source_path) {
58        Ok(path) => path,
59        Err(err) => {
60            return syn::Error::new(input.target_path.span(), err)
61                .to_compile_error()
62                .into();
63        }
64    };
65
66    let data = match fs::read(&absolute_path) {
67        Ok(data) => data,
68        Err(err) => {
69            return syn::Error::new(
70                input.target_path.span(),
71                format!("failed to read '{}': {err}", absolute_path.display()),
72            )
73            .to_compile_error()
74            .into();
75        }
76    };
77
78    expand_from_data(data, true)
79}
80
81#[proc_macro]
82pub fn file_bytes(input: TokenStream) -> TokenStream {
83    let input = parse_macro_input!(input as FileMacroInput);
84    let source_file = input.source_file.as_ref().map(LitStr::value);
85    let source_path = input.target_path.value();
86
87    let absolute_path = match resolve_path(source_file.as_deref(), &source_path) {
88        Ok(path) => path,
89        Err(err) => {
90            return syn::Error::new(input.target_path.span(), err)
91                .to_compile_error()
92                .into();
93        }
94    };
95
96    let data = match fs::read(&absolute_path) {
97        Ok(data) => data,
98        Err(err) => {
99            return syn::Error::new(
100                input.target_path.span(),
101                format!("failed to read '{}': {err}", absolute_path.display()),
102            )
103            .to_compile_error()
104            .into();
105        }
106    };
107
108    expand_from_data(data, false)
109}
110
111fn expand_from_data(data: Vec<u8>, decode_utf8: bool) -> TokenStream {
112    let compressed = match zstd::stream::encode_all(data.as_slice(), 0) {
113        Ok(compressed) => compressed,
114        Err(err) => {
115            return syn::Error::new(
116                proc_macro2::Span::call_site(),
117                format!("failed to compress data: {err}"),
118            )
119            .to_compile_error()
120            .into();
121        }
122    };
123
124    let include_zstd_crate = match crate_name("include-zstd") {
125        // In a package with both lib+bin, proc-macros expanded inside the bin
126        // should still target the library crate namespace.
127        Ok(FoundCrate::Itself) => quote!(::include_zstd),
128        Ok(FoundCrate::Name(name)) => {
129            let ident = syn::Ident::new(&name, proc_macro2::Span::call_site());
130            quote!(::#ident)
131        }
132        Err(_) => quote!(::include_zstd),
133    };
134
135    let expanded = if decode_utf8 {
136        quote! {
137            {
138                static __INCLUDE_ZSTD_COMPRESSED: &[u8] = &[#(#compressed),*];
139                static __INCLUDE_ZSTD_CACHE: ::std::sync::OnceLock<::std::boxed::Box<[u8]>> = ::std::sync::OnceLock::new();
140
141                #include_zstd_crate::__private::decode_utf8(
142                    __INCLUDE_ZSTD_CACHE
143                        .get_or_init(|| #include_zstd_crate::__private::decompress_bytes(__INCLUDE_ZSTD_COMPRESSED))
144                        .as_ref(),
145                )
146            }
147        }
148    } else {
149        quote! {
150            {
151                static __INCLUDE_ZSTD_COMPRESSED: &[u8] = &[#(#compressed),*];
152                static __INCLUDE_ZSTD_CACHE: ::std::sync::OnceLock<::std::boxed::Box<[u8]>> = ::std::sync::OnceLock::new();
153
154                __INCLUDE_ZSTD_CACHE
155                    .get_or_init(|| #include_zstd_crate::__private::decompress_bytes(__INCLUDE_ZSTD_COMPRESSED))
156                    .as_ref()
157            }
158        }
159    };
160
161    expanded.into()
162}
163
164fn resolve_path(source_file: Option<&str>, source_path: &str) -> Result<PathBuf, String> {
165    let target_path = Path::new(source_path);
166    if target_path.is_absolute() {
167        return Ok(target_path.to_path_buf());
168    }
169
170    // Match `include_str!` semantics: always resolve relative paths against the
171    // parent directory of the invocation's source file, using an absolute path
172    // so the result is independent of the compiler's current working directory.
173    let source_file_abs = if let Some(source_file) = source_file {
174        absolutize_source_file(Path::new(source_file))
175    } else {
176        invocation_source_file_abs()
177    };
178
179    let source_dir = source_file_abs.parent().ok_or_else(|| {
180        format!(
181            "failed to resolve include path '{}': invocation source file '{}' has no parent directory",
182            source_path,
183            source_file_abs.display()
184        )
185    })?;
186
187    Ok(source_dir.join(target_path))
188}
189
190/// Return the absolute path of the source file that contains the macro
191/// invocation, mirroring how `include_str!` locates its base directory.
192fn invocation_source_file_abs() -> PathBuf {
193    let call_site = proc_macro::Span::call_site();
194
195    // `local_file()` returns the canonical absolute on-disk path when the span
196    // originates from a real source file; this is the same information rustc
197    // uses internally to resolve `include_str!`.
198    if let Some(path) = call_site.local_file() {
199        return path;
200    }
201
202    // Fallback: `Span::file()` typically yields a path relative to the crate
203    // root (e.g. "src/lib.rs"). Anchor it to `CARGO_MANIFEST_DIR` so the final
204    // path is absolute, matching `include_str!`'s file-relative resolution.
205    absolutize_source_file(Path::new(&call_site.file()))
206}
207
208fn absolutize_source_file(source_file: &Path) -> PathBuf {
209    if source_file.is_absolute() {
210        return source_file.to_path_buf();
211    }
212
213    if let Ok(manifest_dir) = std::env::var("CARGO_MANIFEST_DIR") {
214        return PathBuf::from(manifest_dir).join(source_file);
215    }
216
217    source_file.to_path_buf()
218}
include_zstd_derive/lib.rs

include_zstd_derive/
lib.rs
