diff --git a/src/compose/mod.rs b/src/compose/mod.rs
index af6672f..0d2c7a4 100644
--- a/src/compose/mod.rs
+++ b/src/compose/mod.rs
@@ -362,7 +362,7 @@ impl Default for Composer {
             )
             .unwrap(),
             virtual_fn_regex: Regex::new(
-                r"(?P<lead>[\s]*virtual\s+fn\s+)(?P<function>[^\s]+)(?P<trail>\s*)\(",
+                r"(?P<lead>[\s]*)(?P<virtual_keyword>virtual\s+)(?<fn_keyword>fn\s+)(?P<function_name>[^\s]+)(?P<trail>\s*)\(",
             )
             .unwrap(),
             override_fn_regex: Regex::new(
@@ -858,17 +858,13 @@ impl Composer {
         let source = self
             .virtual_fn_regex
             .replace_all(source, |cap: &regex::Captures| {
-                let target_function = cap.get(2).unwrap().as_str().to_owned();
-
-                let replacement_str = format!(
-                    "{}fn {}{}(",
-                    " ".repeat(cap.get(1).unwrap().range().len() - 3),
-                    target_function,
-                    " ".repeat(cap.get(3).unwrap().range().len()),
-                );
-
-                virtual_functions.insert(target_function);
+                let function_name = cap.name("function_name").unwrap().as_str();
+                let lead = cap.name("lead").unwrap().as_str();
+                let trail = cap.name("trail").unwrap().as_str();
+                let fn_keyword = cap.name("fn_keyword").unwrap().as_str();
 
+                let replacement_str = format!("{lead}{fn_keyword}{function_name}{trail}(");
+                virtual_functions.insert(function_name.to_string());
                 replacement_str
             });
 
diff --git a/src/compose/tests/overrides/mod.wgsl b/src/compose/tests/overrides/mod.wgsl
index a64f3ab..d279fe1 100644
--- a/src/compose/tests/overrides/mod.wgsl
+++ b/src/compose/tests/overrides/mod.wgsl
@@ -1,9 +1,11 @@
 #define_import_path mod
 
+// This is a comment
+
 virtual fn inner(arg: f32) -> f32 {
     return arg * 2.0;
 }
 
 fn outer() -> f32 {
     return inner(1.0);
-}
\ No newline at end of file
+}
diff --git a/src/compose/tests/overrides/top.wgsl b/src/compose/tests/overrides/top.wgsl
index bdabfe5..27afeee 100644
--- a/src/compose/tests/overrides/top.wgsl
+++ b/src/compose/tests/overrides/top.wgsl
@@ -1,9 +1,10 @@
 #import mod
 
+// This is a comment
 override fn mod::inner(arg: f32) -> f32 {
     return arg * 3.0;
 }
 
 fn top() -> f32 {
     return mod::outer();
-}
\ No newline at end of file
+}