Find first function for preamble

visgl · Feb 19, 2020 · 0714612 · 0714612
1 parent ac7da55
commit 0714612
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/modules/shadertools/src/lib/assemble-shaders.js b/modules/shadertools/src/lib/assemble-shaders.js
@@ -5,7 +5,7 @@ import injectShader, {DECLARATION_INJECT_MARKER} from './inject-shader';
 import transpileShader from './transpile-shader';
 import {assert} from '../utils';
 
-const REGEX_START_OF_MAIN = /void main\s*\([^)]*\)\s*\{\n?/;
+const REGEX_FIRST_FUNCTION = /[ \t]*\w+[ \t]+\w+\s*\([^)]*\)\s*\{\n?/;
 const INJECT_SHADER_DECLARATIONS = `\n\n${DECLARATION_INJECT_MARKER}\n\n`;
 
 const SHADER_TYPE = {
@@ -67,9 +67,14 @@ function assembleShader(
     versionLine = `#version ${glslVersion}`;
   }
 
-  const mainIndex = sourceLines.findIndex(l => l.match(REGEX_START_OF_MAIN));
-  const corePreamble = sourceLines.slice(0, mainIndex).join('\n');
-  const coreMain = sourceLines.slice(mainIndex).join('\n');
+  let functionIndex = sourceLines.findIndex(l => l.match(REGEX_FIRST_FUNCTION));
+
+  if (functionIndex === -1) {
+    functionIndex = 0;
+  }
+
+  const corePreamble = sourceLines.slice(0, functionIndex).join('\n');
+  const coreMain = sourceLines.slice(functionIndex).join('\n');
 
   // Combine Module and Application Defines
   const allDefines = {};

diff --git a/modules/shadertools/src/lib/inject-shader.js b/modules/shadertools/src/lib/inject-shader.js
@@ -10,7 +10,7 @@ const MODULE_INJECTORS = {
 
 export const DECLARATION_INJECT_MARKER = '__LUMA_INJECT_DECLARATIONS__'; // Uniform/attribute declarations
 
-const REGEX_START_OF_MAIN = /void main\s*\([^)]*\)\s*\{\n?/; // Beginning of main
+const REGEX_START_OF_MAIN = /void\s+main\s*\([^)]*\)\s*\{\n?/; // Beginning of main
 const REGEX_END_OF_MAIN = /}\n?[^{}]*$/; // End of main, assumes main is last function
 const fragments = [];