rknpu-sys 0.1.2

use std::env;
use std::path::PathBuf;
use std::fs::read_dir;

#[cfg(feature = "bindgen")]
use bindgen::callbacks::{ItemInfo, ParseCallbacks};

#[derive(Debug)]
pub struct StripPrefixCallback {
    remove_prefix: Option<String>,
}

impl StripPrefixCallback {
    pub fn new(prefix: &str) -> StripPrefixCallback {
        StripPrefixCallback {
            remove_prefix: Some(prefix.to_string()),
        }
    }
}

#[cfg(feature = "bindgen")]
impl ParseCallbacks for StripPrefixCallback {
    fn generated_name_override(&self, item_info: ItemInfo<'_>) -> Option<String> {
        self.remove_prefix
            .as_ref()
            .and_then(|s| item_info.name.strip_prefix(s.as_str()).map(String::from))
    }
}

fn main() {
    println!("cargo:rerun-if-env-changed={}", "RKNPU_INCLUDE_DIR");
    println!("cargo:rerun-if-env-changed={}", "RKNPU_LINKING");
    println!("cargo:rerun-if-env-changed={}", "RKNPU_LIB_DIR");
    println!("cargo:rerun-if-env-changed={}", "RKNPU_BINDING_OUT");

    if env::var("DOCS_RS").ok().is_some() {
        return;
    }

    let out_dir = PathBuf::from(env::var_os("OUT_DIR").unwrap());

    let include_dir = env::var_os("RKNPU_INCLUDE_DIR")
        .expect("Env `RKNPU_INCLUDE_DIR` should be set to locate RKNPU headers.");

    if read_dir(&include_dir).is_err() {
        panic!("The directory specified by env `RKNPU_INCLUDE_DIR` does not exist.");
    }

    println!("cargo:include={}", include_dir.to_string_lossy());

    let target = env::var("TARGET").map(|t| t.to_lowercase());
    let linking_opt = env::var("RKNPU_LINKING").ok().map(|v| v.to_lowercase());
    let linking_disabled = linking_opt.as_ref().map(|l| l.eq_ignore_ascii_case("off")).unwrap_or(false);

    if linking_disabled {
        println!("cargo:warning=Library linking is disabled for debug/testing purpose.")
    } else if !target.map_or(false, |t| t.starts_with("arm") || t.starts_with("aarch64")) {
        println!("cargo:warning=Library linking is disabled on non-native architectures.")
    } else {
        let lib_dir = env::var_os("RKNPU_LIB_DIR")
            .expect("Env `RKNPU_LIB_DIR` should be set to locate RKNPU libraries.");

        let lib_dir_entries: Vec<_> = read_dir(&lib_dir).unwrap()
            .filter_map(|entry| entry.ok())
            .collect();

        let libs: Vec<_> = lib_dir_entries.into_iter()
            .filter_map(|entry| entry.file_name().into_string().ok())
            .collect();

        let (lib_name, lib_dylib, lib_static) = if cfg!(feature = "rv1106") {
            ("rknnmrt", "librknnmrt.so", "librknnmrt.a")
        } else {
            ("rknnrt", "librknnrt.so", "librknnrt.a")
        };

        let has_dylib = libs.iter().any(|lib| lib == lib_dylib);
        let has_static = libs.iter().any(|lib| lib == lib_static);

        let linking_kind: &'static str = match linking_opt.as_ref().map(|s| s.as_ref()) {
            Some("dylib" | "dynamic" | "d") => {
                if has_dylib {
                    "dylib"
                } else {
                    panic!("Cannot find dynamic RKNPU library ({}).", lib_dylib)
                }
            }
            Some("static" | "s") => {
                if has_static {
                    "static"
                } else {
                    panic!("Cannot find static RKNPU library ({}).", lib_static)
                }
            }
            Some(_) | None => {
                if has_dylib {
                    "dylib"
                } else if has_static {
                    "static"
                } else {
                    panic!(
                        "Cannot find any RKNPU library ({}, {}), ensure `RKNPU_LIB_DIR` is correctly set.", lib_dylib, lib_static
                    )
                }
            }
        };

        println!("cargo:rustc-link-search=native={}", lib_dir.to_string_lossy());

        println!("cargo:rustc-link-lib={}={}", linking_kind, lib_name);
    }

    #[cfg(feature = "bindgen")]
    {
        #[allow(unused_mut)]
        let mut includes = Vec::new();

        #[cfg(feature = "rknn")]
        includes.push("rknn_api.h");

        #[cfg(feature = "custom-op")]
        includes.push("rknn_custom_op.h");

        #[cfg(feature = "matmul")]
        includes.push("rknn_matmul_api.h");

        let enums = [
            "_rknn_query_cmd",
            "_rknn_tensor_type",
            "_rknn_tensor_qnt_type",
            "_rknn_tensor_format",
            "_rknn_core_mask",
            "_rknn_tensor_mem_flags",
            "_rknn_mem_alloc_flags",
            "_rknn_mem_sync_mode",
        ];

        let mut builder = bindgen::builder()
            .clang_arg("-I")
            .clang_arg(include_dir.to_string_lossy())
            .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
            .header_contents(
                "bindings.h",
                includes.iter()
                    .map(|header| format!("#include <{}>", header))
                    .collect::<Vec<_>>()
                    .join("\n")
                    .as_ref()
            );

        for name in enums {
            builder = builder.newtype_enum(name);
        }

        builder
            .generate()
            .unwrap()
            .write_to_file(out_dir.join("bindings.rs"))
            .unwrap();

        if let Some(path) = env::var_os("RKNPU_BINDING_OUT") {
            std::fs::copy(out_dir.join("bindings.rs"), path).unwrap();
        }
    }
}