Retain original shader source

hakolao · May 12, 2024 · 79b3a98 · 79b3a98
1 parent 73582a9
commit 79b3a98
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/src/utils.rs b/src/utils.rs
@@ -46,6 +46,7 @@ impl std::fmt::Display for ShaderError {
 
 pub struct WatchedShaderModule {
     source: ShaderSource,
+    original_source: Option<ShaderSource>,
     _watchers: HashMap<String, Option<RecommendedWatcher>>,
     _receivers: HashMap<String, Receiver<notify::Result<Event>>>,
 }
@@ -80,6 +81,7 @@ impl WatchedShaderModule {
         };
         Ok(WatchedShaderModule {
             source,
+            original_source: None,
             _watchers: watchers,
             _receivers: receivers,
         })
@@ -133,9 +135,13 @@ impl WatchedShaderModule {
 
     pub fn reload_with_modified_source(
         &mut self,
-        mut modify_fn: impl FnMut(&mut ShaderSource) -> Result<(), ShaderError>,
+        mut modify_fn: impl FnMut(&ShaderSource) -> Result<ShaderSource, ShaderError>,
     ) -> Result<(), ShaderError> {
-        modify_fn(&mut self.source)?;
+        if self.original_source.is_none() {
+            self.original_source = Some(self.source.clone());
+        } else {
+            self.source = modify_fn(self.original_source.as_ref().unwrap())?;
+        }
         self.reload()
     }
 
@@ -590,7 +596,7 @@ const TEST: u32 = u32(1);
 
 const TEST2: u32 = u32(2);
 "#
-        .to_string();
+            .to_string();
         assert_eq!(result.source, should_be);
     }