From a2f823f2837a46154bdf5aa8cff148ba3e774800 Mon Sep 17 00:00:00 2001
From: Ralf Biedert <rb@xr.io>
Date: Tue, 25 Oct 2022 10:16:50 +0200
Subject: [PATCH] On NixOS find cudatoolkit via `which nvcc`, fixes #92.

---
 crates/find_cuda_helper/src/lib.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/crates/find_cuda_helper/src/lib.rs b/crates/find_cuda_helper/src/lib.rs
index b3943be..a93c9da 100644
--- a/crates/find_cuda_helper/src/lib.rs
+++ b/crates/find_cuda_helper/src/lib.rs
@@ -1,5 +1,6 @@
 //! Tiny crate for common logic for finding and including CUDA.
 
+use std::process::Command;
 use std::{
     env,
     path::{Path, PathBuf},
@@ -57,6 +58,7 @@ pub fn find_cuda_root() -> Option<PathBuf> {
 
 #[cfg(target_os = "windows")]
 pub fn find_cuda_lib_dirs() -> Vec<PathBuf> {
+    let x = detect_cuda_root_via_which_nvcc();
     if let Some(root_path) = find_cuda_root() {
         // To do this the right way, we check to see which target we're building for.
         let target = env::var("TARGET")
@@ -131,6 +133,7 @@ pub fn find_cuda_lib_dirs() -> Vec<PathBuf> {
         candidates.push(e)
     }
     candidates.push(PathBuf::from("/usr/lib/cuda"));
+    candidates.push(detect_cuda_root_via_which_nvcc());
 
     let mut valid_paths = vec![];
     for base in &candidates {
@@ -150,6 +153,24 @@ pub fn find_cuda_lib_dirs() -> Vec<PathBuf> {
     valid_paths
 }
 
+#[cfg(not(target_os = "windows"))]
+fn detect_cuda_root_via_which_nvcc() -> PathBuf {
+    let output = Command::new("which")
+        .arg("nvcc")
+        .output()
+        .expect("Command `which` must be available on *nix like systems.")
+        .stdout;
+
+    let path: PathBuf = String::from_utf8(output)
+        .expect("Result must be valid UTF-8")
+        .trim()
+        .to_string()
+        .into();
+
+    // The above finds `CUDASDK/bin/nvcc`, so we have to go 2 up for the SDK root.
+    path.parent().unwrap().parent().unwrap().to_path_buf()
+}
+
 #[cfg(target_os = "windows")]
 pub fn find_optix_root() -> Option<PathBuf> {
     // the optix SDK installer sets OPTIX_ROOT_DIR whenever it installs.