Rust 编译期嵌入指定目录下的所有文件

原理

通过宏实现，代码来自 macro-log：

use proc_macro::TokenStream;
use quote::quote;
use syn::LitStr;

pub fn read_dir(args: TokenStream) -> TokenStream {
    let path = syn::parse_macro_input!(args as LitStr).value() + "/";
    // plan 1
    // let files = get_files(&path).iter().map(|it| format!("r#\"{it}\"#")).collect::<Vec<String>>().join(",");
    // plan 2
    // let files = get_files(&path).iter().map(|it| quote!(#it,).to_string()).collect::<Vec<String>>().join("");
    // plan 3
    // let files = get_files(&path).iter().map(|it| quote!(#it).to_string()).collect::<Vec<String>>().join(",");
    let (workspace, files) = get_files(&path);
    #[cfg(windows)]
    let workspace = workspace.replacen("\\\\?\\", "", 1);
    println!("wrokspace: {workspace}");
    let len = files.len();
    let files = files.iter()
        .map(|it| {
            let file = workspace.clone() + "/" + it;
            quote!((#it, include_bytes!(#file))).to_string()
        }).collect::<Vec<String>>().join(",");
    let files = files.parse::<proc_macro2::TokenStream>().unwrap();
    quote! {
        [#files] as [(&str, &[u8]); #len]
    }.into()
}

use std::io;
use std::fs::{self, DirEntry};
use std::path::Path;

// one possible implementation of walking a directory only visiting files
fn visit_dirs(dir: &Path, cb: &dyn Fn(DirEntry, &mut Vec<DirEntry>), vec: &mut Vec<DirEntry>) -> io::Result<()> {
    if dir.is_dir() {
        for entry in fs::read_dir(dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.is_dir() {
                visit_dirs(&path, cb, vec)?;
            } else {
                cb(entry, vec);
            }
        }
    }
    Ok(())
}

fn get_files(_path: &str) -> (String, Vec<String>) {
    let mut files = vec![];
    let path = Path::new(&_path);
    println!("visit dir: {:?}", path);
    // 规范化工作区路径
    let Ok(workspace_path) = path.canonicalize() else {
        eprintln!("Failed to canonicalize the directory!");
        return (Default::default(), vec![]);
    };
    let Ok(_) = visit_dirs(path, &|entry, files| {
        // println!("file => {}", entry.path().to_string_lossy());
        println!("file => {}", entry.path().to_string_lossy().replacen(_path, "", 1));
        files.push(entry);
    }, &mut files) else {
        eprintln!("Failed to read the directory!");
        return (Default::default(), vec![]);
    };
    (
        workspace_path.to_string_lossy().to_string(),
        files.iter()
            .map(|it| it.path().to_string_lossy().to_string().replacen(_path, "", 1))
            .collect::<Vec<String>>()
    )
}

代码

use super::api;
use super::config::Config;

async fn exit() -> impl axum::response::IntoResponse {
    tokio::spawn(async {
        tokio::time::sleep(std::time::Duration::from_millis(1000)).await;
        std::process::exit(0);
    });
    "程序即将退出"
}

fn route_assets() -> axum::Router {
    let assets = macro_log::read_dir!("assets");
    let mut router = axum::Router::new();
    for (path, bin) in assets {
        #[cfg(windows)]
        let path = format!("/{}", path.replace("\\", "/"));
        macro_log::i!("serve: {path}");
        let mime = match () {
            _ if path.ends_with(".html") => axum::response::TypedHeader(axum::headers::ContentType::html()),
            _ if path.ends_with(".css") => axum::response::TypedHeader(axum::headers::ContentType::from(mime::TEXT_CSS)),
            _ if path.ends_with(".js") => axum::response::TypedHeader(axum::headers::ContentType::from(mime::TEXT_JAVASCRIPT)),
            _ if path.ends_with(".json") => axum::response::TypedHeader(axum::headers::ContentType::from(mime::APPLICATION_JSON)),
            _ => axum::response::TypedHeader(axum::headers::ContentType::octet_stream()),
        };
        router = router.route(&path, axum::routing::get(|| async {
            (
                mime,
                bin.as_ref()
            )
        }));
    }
    router
}

pub fn router(config: Config) -> axum::Router {
    axum::Router::new()
        // .nest_service("/", tower_http::services::ServeDir::new("assets"))
        .nest_service("/", route_assets())
        .route("/exit", axum::routing::get(exit))
        .nest("/api", api::router())
        // .with_state(config)
        .with_state(std::sync::Arc::new(tokio::sync::Mutex::new(config)))
        .layer(
            tower_http::cors::CorsLayer::new()
                .allow_origin(tower_http::cors::AllowOrigin::any())
                .allow_headers(tower_http::cors::Any)
        )
}