compile_codegen: allow SpirvBuilder watch to return something on first completion

Author: Firestar99

crates/spirv-builder/src/watch.rs

@@ -8,11 +8,15 @@ use std::{collections::HashSet, sync::mpsc::sync_channel};
 impl SpirvBuilder {
     /// Watches the module for changes using [`notify`](https://crates.io/crates/notify),
     /// and rebuild it upon changes. Calls `on_compilation_finishes` after each
-    /// successful compilation.
-    pub fn watch(
+    /// successful compilation. The second `Option<AcceptFirstCompile<T>>`
+    /// param allows you to return some `T` on the first compile, which is
+    /// then returned by this function (wrapped in Option).
+    pub fn watch<T>(
         &self,
-        mut on_compilation_finishes: impl FnMut(CompileResult) + Send + 'static,
-    ) -> Result<(), SpirvBuilderError> {
+        mut on_compilation_finishes: impl FnMut(CompileResult, Option<AcceptFirstCompile<'_, T>>)
+        + Send
+        + 'static,
+    ) -> Result<Option<T>, SpirvBuilderError> {
         let path_to_crate = self
             .path_to_crate
             .as_ref()
@@ -42,7 +46,8 @@ impl SpirvBuilder {
             }
         };
         let metadata = self.parse_metadata_file(&metadata_file)?;
-        on_compilation_finishes(metadata);
+        let mut out = None;
+        on_compilation_finishes(metadata, Some(AcceptFirstCompile(&mut out)));
 
         let builder = self.clone();
         let thread = std::thread::spawn(move || {
@@ -57,12 +62,24 @@ impl SpirvBuilder {
                         .parse_metadata_file(&file)
                         .expect("Metadata file is correct");
                     watcher.watch_leaf_deps(&metadata_file);
-                    on_compilation_finishes(metadata);
+                    on_compilation_finishes(metadata, None);
                 }
             }
         });
         drop(thread);
-        Ok(())
+        Ok(out)
+    }
+}
+
+pub struct AcceptFirstCompile<'a, T>(&'a mut Option<T>);
+
+impl<'a, T> AcceptFirstCompile<'a, T> {
+    pub fn new(write: &'a mut Option<T>) -> Self {
+        Self(write)
+    }
+
+    pub fn submit(self, t: T) {
+        *self.0 = Some(t);
     }
 }
 

examples/runners/wgpu/src/lib.rs

@@ -164,13 +164,7 @@ fn maybe_watch(
             // HACK(eddyb) needed because of `debugPrintf` instrumentation limitations
             // (see https://github.com/KhronosGroup/SPIRV-Tools/issues/4892).
             .multimodule(has_debug_printf);
-        let initial_result = if let Some(mut f) = on_watch {
-            builder
-                .watch(move |compile_result| f(handle_compile_result(compile_result)))
-                .expect("Configuration is correct for watching")
-        } else {
-            builder.build().unwrap()
-        };
+
         fn handle_compile_result(compile_result: CompileResult) -> CompiledShaderModules {
             let load_spv_module = |path| {
                 let data = std::fs::read(path).unwrap();
@@ -194,7 +188,22 @@ fn maybe_watch(
                 },
             }
         }
-        handle_compile_result(initial_result)
+
+        if let Some(mut f) = on_watch {
+            builder
+                .watch(move |compile_result, accept| {
+                    let modules = handle_compile_result(compile_result);
+                    if let Some(accept) = accept {
+                        accept.submit(modules);
+                    } else {
+                        f(modules);
+                    }
+                })
+                .expect("Configuration is correct for watching")
+                .unwrap()
+        } else {
+            handle_compile_result(builder.build().unwrap())
+        }
     }
     #[cfg(any(target_os = "android", target_arch = "wasm32"))]
     {