Allow configuring more make parameters from outside the build

boncheolgu · Mar 10, 2021 · d2f567c · d2f567c
1 parent 5346165
commit d2f567c
Showing 1 changed file with 31 additions and 3 deletions.
diff --git a/build.rs b/build.rs
@@ -6,6 +6,7 @@ use std::env;
 use std::path::{Path, PathBuf};
 #[cfg(feature = "build")]
 use std::time::Instant;
+use std::env::VarError;
 
 fn manifest_dir() -> PathBuf {
     PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap())
@@ -118,9 +119,7 @@ fn prepare_tensorflow_library() {
 
             let make_dir = tflite.parent().unwrap();
 
-            make.arg(format!("TARGET={}", target))
-                .arg(format!("TARGET_ARCH={}", arch))
-                .arg("-j")
+            make.arg("-j")
                 // allow parallelism to be overridden
                 .arg(
                     env::var("TFLITE_RS_MAKE_PARALLELISM").unwrap_or_else(|_| {
@@ -131,6 +130,35 @@ fn prepare_tensorflow_library() {
                 .arg("-f")
                 .arg("tensorflow/lite/tools/make/Makefile");
 
+            for &make_var in &[
+                "TARGET_TOOLCHAIN_PREFIX",
+                "EXTRA_CFLAGS"
+            ] {
+                let env_var = format!("TFLITE_RS_MAKE_{}", make_var);
+                println!("cargo:rerun-if-env-changed={}", env_var);
+
+                match env::var(&env_var) {
+                    Ok(result) => {
+                        make.arg(format!("{}={}", make_var, result));
+                    }
+                    Err(VarError::NotPresent) => {
+                        // Try and set some reasonable default values
+                        match make_var {
+                            "TARGET" => {
+                                make.arg(format!("TARGET={}", target));
+                            }
+                            "TARGET_ARCH" => {
+                                make.arg(format!("TARGET_ARCH={}", arch));
+                            }
+                            _ => {}
+                        }
+                    }
+                    Err(VarError::NotUnicode(_)) => {
+                        panic!("Provided a non-unicode value for {}", env_var)
+                    }
+                }
+            }
+
             if cfg!(feature = "no_micro") {
                 println!("Building lib but no micro");
                 make.arg("lib");