XilonenImpact/hk4e_proto/build.rs

use std::{
    fs::{self, File},
    io::{BufRead, BufReader},
    path::Path,
};

pub fn main() {
    let proto_files = [
        "packet_head",
        "server_only",
        "bin.server",
        "redis_data",
        "cmd",
    ];

    for proto_file in &proto_files {
        if Path::new(&format!("{proto_file}.proto")).exists() {
            println!("cargo:rerun-if-changed={proto_file}.proto");

            prost_build::Config::new()
                .out_dir("out/")
                .type_attribute(".", "#[derive(proto_gen::CmdID)]")
                .compile_protos(&[format!("{proto_file}.proto")], &[proto_file])
                .unwrap();
        }
    }

    implement_cmd_id(Path::new("out/cmd.rs")).unwrap();
    implement_cmd_id(Path::new("out/server_only.rs")).unwrap();
}

fn implement_cmd_id(path: &Path) -> std::io::Result<()> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);
    let mut output = Vec::new();

    let mut cmd_id_attr = None;
    for line in reader.lines() {
        let line = line?;
        if line.contains("CmdID: ") {
            cmd_id_attr = Some(make_cmd_id_attr(&line).unwrap());
        } else {
            output.push(line);
            if let Some(attr) = cmd_id_attr.take() {
                output.push(attr);
            }
        }
    }

    fs::write(path, output.join("\n").as_bytes())?;
    Ok(())
}

fn make_cmd_id_attr(line: &str) -> Option<String> {
    let cmd_id = line.split("CmdID: ").nth(1)?.parse::<u16>().ok()?;
    Some(format!("#[cmdid({cmd_id})]"))
}