expressions in overrides (#311)

* expressions in overrides

* fix workgroup parse

* fix lint

* fix (1.0) completely

* checks before bigint converts

Author: munrocket

lib/engine/index.ts

@@ -251,12 +251,12 @@ fn passSampleLevelBilinearRepeat(pass_index: int, uv: float2, lod: float) -> flo
      * Preprocess shader source code
      */
     async preprocess(shader: string): Promise<SourceMap> {
-        const defines = new Map<string, string>([
+        const overrides = new Map<string, string>([
             ['SCREEN_WIDTH', this.screenWidth.toString()],
             ['SCREEN_HEIGHT', this.screenHeight.toString()]
         ]);
 
-        return new Preprocessor(defines).preprocess(shader);
+        return new Preprocessor(overrides).preprocess(shader);
     }
 
     /**

lib/engine/preprocessor.ts

@@ -1,11 +1,11 @@
 /**
  * WGSL shader preprocessor implementation
  */
-import { fetchInclude, parseUint32, WGSLError } from './utils';
+import { evalMathExpression, fetchInclude, parseUint32, WGSLError } from './utils';
 
 // Regular expressions for preprocessing
 const RE_COMMENT = /(\/\/.*|\/\*[\s\S]*?\*\/)/g;
-const RE_WORD = /[a-zA-Z_][a-zA-Z0-9_]*/g;
+const RE_WORD = /\b\w+\b/g;
 
 const STRING_MAX_LEN = 20;
 
@@ -35,15 +35,17 @@ export class SourceMap {
  * Handles WGSL preprocessing including includes, defines, etc.
  */
 export class Preprocessor {
+    private overrides: Map<string, string>;
     private defines: Map<string, string>;
     private source: SourceMap;
     private storageCount: number;
     // private assertCount: number;
     private specialStrings: boolean;
 
-    constructor(defines: Map<string, string>) {
-        this.defines = new Map(defines);
-        this.defines.set('STRING_MAX_LEN', STRING_MAX_LEN.toString());
+    constructor(overrides: Map<string, string>) {
+        this.overrides = new Map(overrides);
+        this.overrides.set('STRING_MAX_LEN', STRING_MAX_LEN.toString());
+        this.defines = new Map();
         this.source = new SourceMap();
         this.storageCount = 0;
         // this.assertCount = 0;
@@ -57,20 +59,14 @@ export class Preprocessor {
         return source.replace(RE_COMMENT, '');
     }
 
-    /**
-     * Substitute defines in source text
-     */
-    private substDefines(source: string): string {
-        return source.replace(RE_WORD, match => {
-            return this.defines.get(match) ?? match;
-        });
-    }
-
     /**
      * Process a single line of shader source
      */
     private async processLine(lineOrig: string, lineNum: number): Promise<void> {
-        let line = this.substDefines(lineOrig);
+        // Substitute overrides and defines
+        let line = lineOrig
+            .replace(RE_WORD, match => this.overrides.get(match) ?? match)
+            .replace(RE_WORD, match => this.defines.get(match) ?? match);
 
         // Handle enable directives
         if (line.trimStart().startsWith('enable')) {
@@ -79,6 +75,14 @@ export class Preprocessor {
             return;
         }
 
+        // Handle override in one line and evaluate to number
+        if (line.trimStart().startsWith('override')) {
+            line = line.replace(RE_COMMENT, '');
+            const tokens = line.trim().replace('=', ' = ').replace(/\s+/g, ' ').split(' ');
+            this.handleOverride(tokens, lineNum);
+            return;
+        }
+
         // Handle preprocessor directives
         if (line.trimStart().startsWith('#')) {
             line = line.replace(RE_COMMENT, '');
@@ -103,7 +107,7 @@ export class Preprocessor {
                     break;
 
                 case '#define':
-                    this.handleDefine(lineOrig, tokens, lineNum);
+                    this.handleDefine(tokens, lineNum);
                     break;
 
                 case '#storage':
@@ -232,20 +236,41 @@ export class Preprocessor {
     /**
      * Handle #define directive
      */
-    private handleDefine(lineOrig: string, tokens: string[], lineNum: number): void {
-        const name = lineOrig.trim().split(' ')[1];
+    private handleDefine(tokens: string[], lineNum: number): void {
+        const name = tokens[1];
         if (!name) {
             throw new WGSLError('Invalid #define syntax', lineNum);
         }
 
         const value = tokens.slice(2).join(' ');
-        if (this.defines.has(name)) {
+        if (value.includes(name)) {
             throw new WGSLError(`Cannot redefine ${name}`, lineNum);
         }
 
         this.defines.set(name, value);
     }
 
+    // /**
+    //  * Handle override in one line and evaluate to number
+    //  */
+    private handleOverride(tokens: string[], lineNum: number): void {
+        const name = tokens[1];
+        let expression = tokens.slice(3).join('');
+        if (!name || !expression.endsWith(';')) {
+            throw new WGSLError(
+                `Please write override in single line ${lineNum} and terminated!`,
+                lineNum
+            );
+        }
+        if (expression.includes(name)) {
+            throw new WGSLError(`Cannot redefine ${name}`, lineNum);
+        }
+
+        expression = expression.slice(0, -1);
+        expression = evalMathExpression(expression, lineNum);
+        this.defines.set(name, expression);
+    }
+
     /**
      * Handle #storage directive
      */

lib/engine/utils.ts

@@ -82,3 +82,83 @@ export async function fetchInclude(name: string): Promise<string | null> {
 export function countNewlines(text: string): number {
     return (text.match(/\n/g) || []).length;
 }
+
+/**
+ * Evaluate a mathematical expression with bigint/float support
+ */
+export function evalMathExpression(expression: string, lineNumber: number): string {
+    const VALID_EXPRESSIONS = /^[0-9+\-*/%^.,()]+$/;
+    const INTEGER_FUNCTIONS = /(select|abs|min|max|sign)/g;
+    const FLOAT_FUNCTIONS =
+        /(asin|acos|atan|atan2|asinh|acosh|atanh|sin|cos|tan|sinh|cosh|tanh|round|floor|ceil|pow|sqrt|log|log2|exp|exp2|fract|clamp|mix|smoothstep|radians|degrees)/g;
+    const ONE_FLOAT = /[+-]?(?=\d*[.eE])(?=\.?\d)\d*\.?\d*(?:[eE][+-]?\d+)?/;
+    const ONE_INTEGER = /^[-+]?\d+$/;
+    const ALL_INTEGERS = /[-+]?\d+/g;
+
+    let result = expression.replace(/\s+/g, '');
+    if (result === '') {
+        return '';
+    }
+
+    // safety check for eval vulnerabilities
+    let stripped = result.replace(INTEGER_FUNCTIONS, '').replace(FLOAT_FUNCTIONS, '');
+    if (!VALID_EXPRESSIONS.test(stripped)) {
+        stripped = stripped.replace(/[0-9+\-*/%^.,()]+/g, '');
+        throw new WGSLError(`Unsafe symbols '${stripped}' in expression ${expression}`, lineNumber);
+    }
+
+    // convert to bigint if there are no any floating point numbers or float functions
+    let isAbstractInteger = false;
+    if (!ONE_FLOAT.test(expression) && !FLOAT_FUNCTIONS.test(expression)) {
+        result = result.replace(ALL_INTEGERS, match => `${match}n`);
+        isAbstractInteger = true;
+    }
+
+    const mathReplacements = {
+        'select(': '((f,t,cond) => cond ? t : f)(',
+        'abs(': '((x) => Math.abs(Number(x)))(',
+        'min(': '((x,y) => Math.min(Number(x),Number(y)))(',
+        'max(': '((x,y) => Math.max(Number(x),Number(y)))(',
+        'sign(': '((x) => Math.sign(Number(x)))(',
+        'sin(': 'Math.sin(',
+        'cos(': 'Math.cos(',
+        'tan(': 'Math.tan(',
+        'sinh(': 'Math.sinh(',
+        'cosh(': 'Math.cosh(',
+        'tanh(': 'Math.tanh(',
+        'asin(': 'Math.asin(',
+        'acos(': 'Math.acos(',
+        'atan(': 'Math.atan(',
+        'atan2(': 'Math.atan2(',
+        'asinh(': 'Math.asinh(',
+        'acosh(': 'Math.acosh(',
+        'atanh(': 'Math.atanh(',
+        'round(': 'Math.round(',
+        'floor(': 'Math.floor(',
+        'ceil(': 'Math.ceil(',
+        'pow(': 'Math.pow(',
+        'sqrt(': 'Math.sqrt(',
+        'log(': 'Math.log(',
+        'log2(': 'Math.log2(',
+        'exp(': 'Math.exp(',
+        'exp2(': 'Math.pow(2,',
+        'fract(': '((x) => x - Math.floor(x))(',
+        'clamp(': '((x,min,max) => Math.min(Math.max(x,min),max))(',
+        'mix(': '((x,y,a) => x*(1-a) + y*a)(',
+        'smoothstep(':
+            '((e0,e1,x) => { let t = Math.min(Math.max((x-e0)/(e1-e0),0),1); return t*t*(3-2*t); })(',
+        'step(': '((edge,x) => x < edge ? 0 : 1)(',
+        'radians(': '((x) => x * Math.PI / 180)(',
+        'degrees(': '((x) => x * 180 / Math.PI)('
+    };
+
+    result = result.replaceAll('^', '**');
+    result = result.replace(/\b\w+\(/g, match => mathReplacements[match] || '');
+
+    try {
+        result = Function(`'use strict'; return (${result}).toString();`)();
+        return !isAbstractInteger && ONE_INTEGER.test(result) ? `${result}.0` : result;
+    } catch (error) {
+        throw new WGSLError(`Invalid eval '${expression}' -> '${result}' (${error})`, lineNumber);
+    }
+}