embed-dir/src/lib.rs

use mime::Mime;
use proc_macro::TokenStream;
use quote::quote;
use syn::spanned::Spanned;
use syn::{Attribute, DeriveInput, Expr, Ident, Lit, Meta, parse_macro_input};

use std::fs;
use std::path::{Path, PathBuf};

#[derive(Debug, Clone, Copy)]
enum EmbedMode {
    Bytes,
    Str,
    BytesMime,
}

#[proc_macro_derive(Embed, attributes(dir, mode))]
pub fn embed(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = parse_macro_input!(input as DeriveInput);

    let struct_name = &input.ident;

    let dir_attr = input
        .attrs
        .iter()
        .find(|e| e.path().is_ident("dir"))
        .expect("No #[dir = \"...\"] attribute found");

    let mode_attr = input.attrs.iter().find(|e| e.path().is_ident("mode"));

    let mode = mode_attr.map(extract_mode).unwrap_or(EmbedMode::Bytes);

    let base_path = PathBuf::from(extract_dir_path(dir_attr));
    let source_file = PathBuf::from(input.span().unwrap().file());

    let source_dir = if let Some(parent) = source_file.parent() {
        parent
    } else {
        // HACK: when running in rust-analyzer i can't seem to get the parent dir.
        return TokenStream::from(generate_byte_impl(struct_name, Vec::new()));
    };

    let absolue_path = source_dir.join(&base_path);

    let mut match_arms = Vec::new();

    for entry in collect_files(&absolue_path) {
        let rel_path = entry
            .0
            .strip_prefix(&absolue_path)
            .unwrap()
            .to_str()
            .unwrap()
            .replace("\\", "/");

        let include_path = base_path.join(&rel_path);
        let include_string = include_path.to_str().unwrap();

        let arm = match mode {
            EmbedMode::Bytes => generate_byte_arm(&rel_path, include_string),
            EmbedMode::Str => generate_str_arm(&rel_path, include_string),
            EmbedMode::BytesMime => generate_mime_arm(
                &rel_path,
                include_string,
                entry.1.unwrap_or(mime::APPLICATION_OCTET_STREAM),
            ),
        };

        match_arms.push(arm);
    }

    #[allow(unused_mut)]
    let mut expanded = match mode {
        EmbedMode::Bytes => generate_byte_impl(struct_name, match_arms),
        EmbedMode::Str => generate_str_impl(struct_name, match_arms),
        EmbedMode::BytesMime => generate_mime_impl(struct_name, match_arms),
    };

    proc_macro::TokenStream::from(expanded)
}

fn collect_files(dir: &Path) -> Vec<(PathBuf, Option<Mime>)> {
    let mut files = Vec::new();

    for entry in fs::read_dir(dir).unwrap() {
        let path = entry.unwrap().path();

        if path.is_file() {
            let mime_type = mime_guess::from_path(&path);
            files.push((path, mime_type.first()));
        } else if path.is_dir() {
            files.extend(collect_files(&path));
        }
    }

    files
}

fn extract_dir_path(attr: &Attribute) -> String {
    let meta = match &attr.meta {
        Meta::NameValue(meta) => meta,
        _ => panic!("Expected #[dir = \"...\"] as a name-value attribute."),
    };

    let expr_lit = match &meta.value {
        Expr::Lit(expr_lit) => expr_lit,
        _ => panic!("Expected #[dir = \"...\"] with a string literal."),
    };

    match &expr_lit.lit {
        Lit::Str(str) => str.value(),
        _ => panic!("Expected #[dir = \"...\"] to be a string."),
    }
}

fn extract_mode(attr: &Attribute) -> EmbedMode {
    let meta = match &attr.meta {
        Meta::NameValue(meta) => meta,
        _ => panic!("Expected #[mode = \"bytes\"|\"str\"|\"mime\"] as a name-value attribute."),
    };

    let expr_lit = match &meta.value {
        Expr::Lit(expr_lit) => expr_lit,
        _ => panic!("Expected #[mode = \"bytes\"|\"str\"|\"mime\"] with a string literal."),
    };

    match &expr_lit.lit {
        Lit::Str(str) => match str.value().as_str() {
            "bytes" => EmbedMode::Bytes,
            "str" => EmbedMode::Str,
            "mime" => EmbedMode::BytesMime,
            other => panic!("Unknown mode: {other}. Use `bytes`,`str` or `mime`."),
        },
        _ => panic!("Expected #[mode = \"bytes\"|\"str\"|\"mime\"] to be a string."),
    }
}

fn generate_byte_arm(rel: &str, include: &str) -> proc_macro2::TokenStream {
    quote! {
        #rel => Some(include_bytes!(#include)),
    }
}

fn generate_byte_impl(
    struct_name: &Ident,
    match_arms: Vec<proc_macro2::TokenStream>,
) -> proc_macro2::TokenStream {
    quote! {
        impl #struct_name {
            pub fn get(name: &str) -> Option<&'static [u8]> {
                match name {
                    #(#match_arms)*
                    _ => None,
                }
            }
        }
    }
}

fn generate_str_arm(rel: &str, include: &str) -> proc_macro2::TokenStream {
    quote! {
        #rel => Some(include_str!(#include)),
    }
}

fn generate_str_impl(
    struct_name: &Ident,
    match_arms: Vec<proc_macro2::TokenStream>,
) -> proc_macro2::TokenStream {
    quote! {
        impl #struct_name {
            pub fn get(name: &str) -> Option<&'static str> {
                match name {
                    #(#match_arms)*
                    _ => None,
                }
            }
        }
    }
}

fn generate_mime_arm(rel: &str, include: &str, mime_type: Mime) -> proc_macro2::TokenStream {
    let mime_str = mime_type.essence_str();
    quote! {
        #rel => Some((include_bytes!(#include),#mime_str)),
    }
}

fn generate_mime_impl(
    struct_name: &Ident,
    match_arms: Vec<proc_macro2::TokenStream>,
) -> proc_macro2::TokenStream {
    quote! {
        impl #struct_name {
            pub fn get(name: &str) -> Option<(&'static [u8],&'static str)> {
                match name {
                    #(#match_arms)*
                    _ => None,
                }
            }
        }
    }
}