👀 hot shader reloading

crates/spirv-builder/Cargo.toml

@@ -10,6 +10,7 @@ license = "MIT OR Apache-2.0"
 default = ["use-compiled-tools"]
 use-installed-tools = ["rustc_codegen_spirv/use-installed-tools"]
 use-compiled-tools = ["rustc_codegen_spirv/use-compiled-tools"]
+watch = ["notify"]
 
 [dependencies]
 memchr = "2.3"
@@ -18,3 +19,5 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 # See comment in lib.rs invoke_rustc for why this is here
 rustc_codegen_spirv = { path = "../rustc_codegen_spirv", default-features = false }
+
+notify = { version = "5.0.0-pre.10", optional = true }

crates/spirv-builder/src/lib.rs

@@ -54,6 +54,8 @@
 #![allow()]
 
 mod depfile;
+#[cfg(feature = "watch")]
+mod watch;
 
 use raw_string::{RawStr, RawString};
 use serde::Deserialize;
@@ -71,10 +73,12 @@ pub use rustc_codegen_spirv::rspirv::spirv::Capability;
 pub use rustc_codegen_spirv::{CompileResult, ModuleResult};
 
 #[derive(Debug)]
+#[non_exhaustive]
 pub enum SpirvBuilderError {
     CratePathDoesntExist(PathBuf),
     BuildFailed,
     MultiModuleWithPrintMetadata,
+    WatchWithPrintMetadata,
     MetadataFileMissing(std::io::Error),
     MetadataFileMalformed(serde_json::Error),
 }
@@ -89,6 +93,9 @@ impl fmt::Display for SpirvBuilderError {
             SpirvBuilderError::MultiModuleWithPrintMetadata => f.write_str(
                 "Multi-module build cannot be used with print_metadata = MetadataPrintout::Full",
             ),
+            SpirvBuilderError::WatchWithPrintMetadata => {
+                f.write_str("Watching within build scripts will prevent build completion")
+            }
             SpirvBuilderError::MetadataFileMissing(_) => {
                 f.write_str("Multi-module metadata file missing")
             }
@@ -244,22 +251,47 @@ impl SpirvBuilder {
 
     /// Builds the module. If `print_metadata` is [`MetadataPrintout::Full`], you usually don't have to inspect the path
     /// in the result, as the environment variable for the path to the module will already be set.
-    pub fn build(self) -> Result<CompileResult, SpirvBuilderError> {
+    pub fn build(mut self) -> Result<CompileResult, SpirvBuilderError> {
+        self.validate_running_conditions()?;
+        let metadata_file = invoke_rustc(&self)?;
+        match self.print_metadata {
+            MetadataPrintout::Full | MetadataPrintout::DependencyOnly => {
+                leaf_deps(&metadata_file, |artifact| {
+                    println!("cargo:rerun-if-changed={}", artifact)
+                })
+                // Close enough
+                .map_err(SpirvBuilderError::MetadataFileMissing)?;
+            }
+            MetadataPrintout::None => (),
+        }
+        let metadata = self.parse_metadata_file(&metadata_file)?;
+
+        Ok(metadata)
+    }
+
+    pub(crate) fn validate_running_conditions(&mut self) -> Result<(), SpirvBuilderError> {
         if (self.print_metadata == MetadataPrintout::Full) && self.multimodule {
             return Err(SpirvBuilderError::MultiModuleWithPrintMetadata);
         }
         if !self.path_to_crate.is_dir() {
-            return Err(SpirvBuilderError::CratePathDoesntExist(self.path_to_crate));
+            return Err(SpirvBuilderError::CratePathDoesntExist(std::mem::take(
+                &mut self.path_to_crate,
+            )));
         }
-        let metadata_file = invoke_rustc(&self)?;
-        let metadata_contents =
-            File::open(&metadata_file).map_err(SpirvBuilderError::MetadataFileMissing)?;
+        Ok(())
+    }
+
+    pub(crate) fn parse_metadata_file(
+        &self,
+        at: &Path,
+    ) -> Result<CompileResult, SpirvBuilderError> {
+        let metadata_contents = File::open(&at).map_err(SpirvBuilderError::MetadataFileMissing)?;
         let metadata: CompileResult = serde_json::from_reader(BufReader::new(metadata_contents))
             .map_err(SpirvBuilderError::MetadataFileMalformed)?;
         match &metadata.module {
             ModuleResult::SingleModule(spirv_module) => {
                 assert!(!self.multimodule);
-                let env_var = metadata_file.file_name().unwrap().to_str().unwrap();
+                let env_var = at.file_name().unwrap().to_str().unwrap();
                 if self.print_metadata == MetadataPrintout::Full {
                     println!("cargo:rustc-env={}={}", env_var, spirv_module.display());
                 }
@@ -423,13 +455,8 @@ fn invoke_rustc(builder: &SpirvBuilder) -> Result<PathBuf, SpirvBuilderError> {
     // we do that even in case of an error, to let through any useful messages
     // that ended up on stdout instead of stderr.
     let stdout = String::from_utf8(build.stdout).unwrap();
-    let artifact = get_last_artifact(&stdout);
-
     if build.status.success() {
-        match builder.print_metadata {
-            MetadataPrintout::Full | MetadataPrintout::DependencyOnly => print_deps_of(&artifact),
-            MetadataPrintout::None => (),
-        }
+        let artifact = get_last_artifact(&stdout);
         Ok(artifact)
     } else {
         Err(SpirvBuilderError::BuildFailed)
@@ -467,23 +494,28 @@ fn get_last_artifact(out: &str) -> PathBuf {
     filename.into()
 }
 
-fn print_deps_of(artifact: &Path) {
+/// Internally iterate through the leaf dependencies of the artifact at `artifact`
+fn leaf_deps(artifact: &Path, mut handle: impl FnMut(&RawStr)) -> std::io::Result<()> {
     let deps_file = artifact.with_extension("d");
     let mut deps_map = HashMap::new();
     depfile::read_deps_file(&deps_file, |item, deps| {
         deps_map.insert(item, deps);
         Ok(())
-    })
-    .expect("Could not read dep file");
-    fn recurse(map: &HashMap<RawString, Vec<RawString>>, artifact: &RawStr) {
+    })?;
+    fn recurse(
+        map: &HashMap<RawString, Vec<RawString>>,
+        artifact: &RawStr,
+        handle: &mut impl FnMut(&RawStr),
+    ) {
         match map.get(artifact) {
             Some(entries) => {
                 for entry in entries {
-                    recurse(map, entry)
+                    recurse(map, entry, handle)
                 }
             }
-            None => println!("cargo:rerun-if-changed={}", artifact),
+            None => handle(artifact),
         }
     }
-    recurse(&deps_map, artifact.to_str().unwrap().into());
+    recurse(&deps_map, artifact.to_str().unwrap().into(), &mut handle);
+    Ok(())
 }

crates/spirv-builder/src/watch.rs

@@ -0,0 +1,104 @@
+use std::{collections::HashSet, sync::mpsc::sync_channel};
+
+use notify::{Event, RecursiveMode, Watcher};
+use rustc_codegen_spirv::CompileResult;
+
+use crate::{leaf_deps, SpirvBuilder, SpirvBuilderError};
+
+impl SpirvBuilder {
+    /// Watches the module for changes using [`notify`](https://crates.io/crates/notify).
+    ///
+    /// This is a blocking operation, wand should never return in the happy path
+    pub fn watch(
+        mut self,
+        on_compilation_finishes: impl Fn(CompileResult),
+    ) -> Result<(), SpirvBuilderError> {
+        self.validate_running_conditions()?;
+        if !matches!(self.print_metadata, crate::MetadataPrintout::None) {
+            return Err(SpirvBuilderError::WatchWithPrintMetadata);
+        }
+        let metadata_result = crate::invoke_rustc(&self);
+        // Load the dependencies of the thing
+        let metadata_file = match metadata_result {
+            Ok(path) => path,
+            Err(_) => {
+                let (tx, rx) = sync_channel(0);
+                // Fall back to watching from the crate root if the inital compilation fails
+                let mut watcher =
+                    notify::immediate_watcher(move |event: notify::Result<Event>| match event {
+                        Ok(e) => match e.kind {
+                            notify::EventKind::Access(_) => (),
+                            notify::EventKind::Any
+                            | notify::EventKind::Create(_)
+                            | notify::EventKind::Modify(_)
+                            | notify::EventKind::Remove(_)
+                            | notify::EventKind::Other => {
+                                let _ = tx.try_send(());
+                            }
+                        },
+                        Err(e) => println!("notify error: {:?}", e),
+                    })
+                    .expect("Could create watcher");
+                // This is likely to notice changes in the `target` dir, however, given that `cargo watch` doesn't seem to handle that,
+                watcher
+                    .watch(&self.path_to_crate, RecursiveMode::Recursive)
+                    .expect("Could watch crate root");
+                loop {
+                    rx.recv().expect("Watcher still alive");
+                    let metadata_file = crate::invoke_rustc(&self);
+                    if let Ok(f) = metadata_file {
+                        break f;
+                    }
+                }
+            }
+        };
+        let metadata = self.parse_metadata_file(&metadata_file)?;
+        on_compilation_finishes(metadata);
+        let mut watched_paths = HashSet::new();
+        let (tx, rx) = sync_channel(0);
+        let mut watcher =
+            notify::immediate_watcher(move |event: notify::Result<Event>| match event {
+                Ok(e) => match e.kind {
+                    notify::EventKind::Access(_) => (),
+                    notify::EventKind::Any
+                    | notify::EventKind::Create(_)
+                    | notify::EventKind::Modify(_)
+                    | notify::EventKind::Remove(_)
+                    | notify::EventKind::Other => {
+                        let _ = tx.try_send(());
+                    }
+                },
+                Err(e) => println!("notify error: {:?}", e),
+            })
+            .expect("Could create watcher");
+        leaf_deps(&metadata_file, |it| {
+            let path = it.to_path().unwrap();
+            if watched_paths.insert(path.to_owned()) {
+                watcher
+                    .watch(it.to_path().unwrap(), RecursiveMode::NonRecursive)
+                    .expect("Cargo dependencies are valid files");
+            }
+        })
+        .expect("Could read dependencies file");
+        loop {
+            rx.recv().expect("Watcher still alive");
+            let metadata_result = crate::invoke_rustc(&self);
+            if let Ok(file) = metadata_result {
+                // We can bubble this error up because it's an internal error  (e.g. rustc_codegen_spirv's version of CompileResult is somehow out of sync)
+                let metadata = self.parse_metadata_file(&file)?;
+
+                leaf_deps(&file, |it| {
+                    let path = it.to_path().unwrap();
+                    if watched_paths.insert(path.to_owned()) {
+                        watcher
+                            .watch(it.to_path().unwrap(), RecursiveMode::NonRecursive)
+                            .expect("Cargo dependencies are valid files");
+                    }
+                })
+                .expect("Could read dependencies file");
+
+                on_compilation_finishes(metadata);
+            }
+        }
+    }
+}

examples/runners/wgpu/Cargo.toml

@@ -25,7 +25,7 @@ clap = "3.0.0-beta.2"
 strum = { version = "0.20", default_features = false, features = ["derive"] }
 
 [target.'cfg(not(any(target_os = "android", target_arch = "wasm32")))'.dependencies]
-spirv-builder = { path = "../../../crates/spirv-builder", default-features = false }
+spirv-builder = { path = "../../../crates/spirv-builder", default-features = false, features = ["watch"] }
 
 [target.'cfg(target_os = "android")'.dependencies]
 ndk-glue = "0.2"