#![forbid(unsafe_code)]
#![deny(unused_imports)]
#![deny(missing_docs)]
use std::{
collections::{HashMap, HashSet},
path::PathBuf,
};
use proc_macro::TokenStream;
use pure_magic::{MagicDb, MagicSource};
use quote::quote;
use syn::{
Expr, ExprArray, ItemStruct, Meta, MetaNameValue, Token, parse::Parser, punctuated::Punctuated,
spanned::Spanned,
};
struct MetaParser {
attr: proc_macro2::TokenStream,
metas: HashMap<String, Meta>,
}
impl MetaParser {
fn parse_meta(attr: proc_macro2::TokenStream) -> Result<Self, syn::Error> {
let mut out = HashMap::new();
let parser = Punctuated::<Meta, Token![,]>::parse_terminated;
let metas = match parser.parse2(attr.clone()) {
Ok(m) => m,
Err(e) => return Err(syn::Error::new_spanned(attr, e.to_string())),
};
for meta in metas {
out.insert(
meta.path()
.get_ident()
.ok_or(syn::Error::new_spanned(
meta.clone(),
"failed to process meta",
))?
.to_string(),
meta,
);
}
Ok(Self {
attr: attr.clone(),
metas: out,
})
}
fn get_key_value(&self, key: &str) -> Result<Option<&MetaNameValue>, syn::Error> {
if let Some(meta) = self.metas.get(key) {
match meta {
Meta::NameValue(m) => return Ok(Some(m)),
_ => {
return Err(syn::Error::new_spanned(
&self.attr,
format!("expecting a key value attribute: {key}"),
));
}
}
}
Ok(None)
}
}
fn meta_name_value_to_string_vec(
nv: &MetaNameValue,
) -> Result<Vec<(proc_macro2::Span, String)>, syn::Error> {
if let Expr::Array(ExprArray { elems, .. }) = &nv.value {
Ok(elems
.into_iter()
.filter_map(|e| match e {
Expr::Lit(syn::ExprLit {
lit: syn::Lit::Str(lit_str),
..
}) => Some((lit_str.span(), lit_str.value())),
_ => None,
})
.collect::<Vec<_>>())
} else {
Err(syn::Error::new_spanned(
&nv.value,
"expected an array literal like [\"foo\", \"bar\"]",
))
}
}
fn impl_magic_embed(attr: TokenStream, item: TokenStream) -> Result<TokenStream, syn::Error> {
let input_struct: ItemStruct = syn::parse2(item.into())?;
let struct_name = &input_struct.ident;
let cs = proc_macro::Span::call_site();
let Some(source_file) = cs.local_file() else {
return Ok(quote! {}.into());
};
let source_dir = source_file.parent().unwrap();
let ts2: proc_macro2::TokenStream = attr.into();
let struct_vis = input_struct.vis;
let metas = MetaParser::parse_meta(ts2)?;
let exclude = if let Some(exclude) = metas.get_key_value("exclude")? {
meta_name_value_to_string_vec(exclude)?
.into_iter()
.map(|(s, p)| (s, source_dir.join(p)))
.collect()
} else {
vec![]
};
let include_nv = metas.get_key_value("include")?.ok_or(syn::Error::new(
struct_name.span(),
"expected a list of files or directory to include: \"include\" = [\"magdir\"]",
))?;
let include: Vec<(proc_macro2::Span, PathBuf)> = meta_name_value_to_string_vec(include_nv)?
.into_iter()
.map(|(s, p)| (s, source_dir.join(p)))
.collect();
let mut wo = fs_walk::WalkOptions::new();
wo.files().max_depth(0).sort(true);
let mut db = MagicDb::new();
let exclude_set: HashSet<PathBuf> = exclude.into_iter().map(|(_, p)| p).collect();
macro_rules! load_file {
($span: expr, $path: expr) => {
MagicSource::open($path).map_err(|e| {
syn::Error::new(
$span.clone(),
format!(
"failed to parse magic file={}: {e}",
$path.to_string_lossy()
),
)
})?
};
}
let mut rules = vec![];
for (s, p) in include.iter() {
if p.is_dir() {
for rule_file in wo.walk(p) {
let rule_file = rule_file
.map_err(|e| syn::Error::new(*s, format!("failed to list rule file: {e}")))?;
if exclude_set.contains(&rule_file) {
continue;
}
rules.push(load_file!(s, &rule_file));
}
} else if p.is_file() {
rules.push(load_file!(s, p));
}
}
db.load_bulk(rules.into_iter());
db.verify()
.map_err(|e| syn::Error::new(include_nv.span(), format!("inconsistent database: {e}")))?;
let mut ser = vec![];
db.serialize(&mut ser).map_err(|e| {
syn::Error::new(
struct_name.span(),
format!("failed to serialize database: {e}"),
)
})?;
let output = quote! {
#struct_vis struct #struct_name;
impl #struct_name {
const DB: &[u8] = &[ #( #ser ),* ];
#struct_vis fn open() -> Result<pure_magic::MagicDb, pure_magic::Error> {
pure_magic::MagicDb::deserialize(&mut Self::DB.as_ref())
}
}
};
Ok(output.into())
}
#[proc_macro_attribute]
pub fn magic_embed(attr: TokenStream, item: TokenStream) -> TokenStream {
match impl_magic_embed(attr, item) {
Ok(ts) => ts,
Err(e) => e.to_compile_error().into(),
}
}