PicoBot/build.rs

use std::env;
use std::fs;
use std::io::Write;
use std::path::Path;

fn main() {
    let out_dir = env::var("OUT_DIR").unwrap();
    let skills_dir = Path::new("resources/skills");
    let skills_out_dir = Path::new(&out_dir).join("skills");
    fs::create_dir_all(&skills_out_dir).unwrap();

    let mut skills = Vec::new();

    if let Ok(entries) = fs::read_dir(skills_dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if !path.is_dir() {
                continue;
            }

            let skill_name = path.file_name().unwrap().to_str().unwrap().to_string();
            let compressed_path = skills_out_dir.join(format!("{}.tar.zst", skill_name));

            let compressed = compress_skill_dir(&path);
            fs::write(&compressed_path, &compressed).unwrap();

            skills.push(skill_name);
        }
    }

    let mut code = String::from(
        r#"pub struct EmbeddedSkill {
    pub name: &'static str,
    pub data: &'static [u8],
}

pub static EMBEDDED_SKILLS: &[EmbeddedSkill] = &[
"#,
    );

    for name in &skills {
        let file_path = skills_out_dir
            .join(format!("{}.tar.zst", name))
            .to_string_lossy()
            .to_string();
        code.push_str(&format!(
            "    EmbeddedSkill {{ name: \"{}\", data: include_bytes!(\"{}\") }},\n",
            name, file_path
        ));
    }

    code.push_str("];\n");

    let generated_path = Path::new(&out_dir).join("embedded_skills.rs");
    let mut f = fs::File::create(&generated_path).unwrap();
    f.write_all(code.as_bytes()).unwrap();
}

fn compress_skill_dir(dir: &Path) -> Vec<u8> {
    let mut buf = Vec::new();
    let mut builder = tar::Builder::new(&mut buf);
    builder.follow_symlinks(false);
    builder
        .append_dir_all(dir.file_name().unwrap().to_str().unwrap(), dir)
        .unwrap();
    drop(builder);

    zstd::encode_all(buf.as_slice(), 3).unwrap()
}