Fix p5.strands uniform calls, add instance mode construct

Author: davepagurek

src/webgl/ShaderGenerator.js

@@ -16,8 +16,10 @@ function shadergenerator(p5, fn) {
 
   const oldModify = p5.Shader.prototype.modify
 
-  p5.Shader.prototype.modify = function(shaderModifier, options = { parser: true, srcLocations: false }) {
+  p5.Shader.prototype.modify = function(shaderModifier, scope = {}) {
     if (shaderModifier instanceof Function) {
+      // TODO make this public. Currently for debugging only.
+      const options = { parser: true, srcLocations: false };
       let generatorFunction;
       if (options.parser) {
         const sourceString = shaderModifier.toString()
@@ -27,13 +29,17 @@ function shadergenerator(p5, fn) {
         });
         ancestor(ast, ASTCallbacks, undefined, { varyings: {} });
         const transpiledSource = escodegen.generate(ast);
-        generatorFunction = new Function(
+        const scopeKeys = Object.keys(scope);
+        const internalGeneratorFunction = new Function(
+          'p5',
+          ...scopeKeys,
           transpiledSource
           .slice(
             transpiledSource.indexOf('{') + 1,
             transpiledSource.lastIndexOf('}')
           ).replaceAll(';', '')
         );
+        generatorFunction = () => internalGeneratorFunction(p5, ...scopeKeys.map(key => scope[key]));
       } else {
         generatorFunction = shaderModifier;
       }
@@ -64,15 +70,24 @@ function shadergenerator(p5, fn) {
     }
   }
 
-  function ancestorIsUniform(ancestor) {
+  function nodeIsUniform(ancestor) {
     return ancestor.type === 'CallExpression'
-      && ancestor.callee?.type === 'Identifier'
-      && ancestor.callee?.name.startsWith('uniform');
+      && (
+        (
+          // Global mode
+          ancestor.callee?.type === 'Identifier' &&
+          ancestor.callee?.name.startsWith('uniform')
+        ) || (
+          // Instance mode
+          ancestor.callee?.type === 'MemberExpression' &&
+          ancestor.callee?.property.name.startsWith('uniform')
+        )
+      );
   }
 
   const ASTCallbacks = {
-    UnaryExpression(node, _state, _ancestors) {
-      if (_ancestors.some(ancestorIsUniform)) { return; }
+    UnaryExpression(node, _state, ancestors) {
+      if (ancestors.some(nodeIsUniform)) { return; }
 
       const signNode = {
         type: 'Literal',
@@ -83,7 +98,7 @@ function shadergenerator(p5, fn) {
           node.type = 'CallExpression'
           node.callee = {
             type: 'Identifier',
-            name: 'unaryNode',
+            name: 'p5.unaryNode',
           }
           node.arguments = [node.argument, signNode]
       }
@@ -106,7 +121,7 @@ function shadergenerator(p5, fn) {
             type: 'CallExpression',
             callee: {
               type: 'Identifier',
-              name: 'unaryNode'
+              name: 'p5.unaryNode'
             },
             arguments: [node.argument.object, signNode],
           };
@@ -123,8 +138,9 @@ function shadergenerator(p5, fn) {
       delete node.argument;
       delete node.operator;
     },
-    VariableDeclarator(node, _state, _ancestors) {
-      if (node.init.callee && node.init.callee.name?.startsWith('uniform')) {
+    VariableDeclarator(node, _state, ancestors) {
+      if (ancestors.some(nodeIsUniform)) { return; }
+      if (nodeIsUniform(node.init)) {
         const uniformNameLiteral = {
           type: 'Literal',
           value: node.id.name
@@ -140,7 +156,8 @@ function shadergenerator(p5, fn) {
         _state.varyings[node.id.name] = varyingNameLiteral;
       }
     },
-    Identifier(node, _state, _ancestors) {
+    Identifier(node, _state, ancestors) {
+      if (ancestors.some(nodeIsUniform)) { return; }
       if (_state.varyings[node.name]
           && !_ancestors.some(a => a.type === 'AssignmentExpression' && a.left === node)) {
         node.type = 'ExpressionStatement';
@@ -163,16 +180,18 @@ function shadergenerator(p5, fn) {
     },
     // The callbacks for AssignmentExpression and BinaryExpression handle
     // operator overloading including +=, *= assignment expressions
-    ArrayExpression(node, _state, _ancestors) {
+    ArrayExpression(node, _state, ancestors) {
+      if (ancestors.some(nodeIsUniform)) { return; }
       const original = JSON.parse(JSON.stringify(node));
       node.type = 'CallExpression';
       node.callee = {
         type: 'Identifier',
-        name: 'dynamicNode',
+        name: 'p5.dynamicNode',
       };
       node.arguments = [original];
     },
-    AssignmentExpression(node, _state, _ancestors) {
+    AssignmentExpression(node, _state, ancestors) {
+      if (ancestors.some(nodeIsUniform)) { return; }
       if (node.operator !== '=') {
         const methodName = replaceBinaryOperator(node.operator.replace('=',''));
         const rightReplacementNode = {
@@ -209,10 +228,10 @@ function shadergenerator(p5, fn) {
           }
         }
       },
-    BinaryExpression(node, _state, _ancestors) {
+    BinaryExpression(node, _state, ancestors) {
       // Don't convert uniform default values to node methods, as
       // they should be evaluated at runtime, not compiled.
-      if (_ancestors.some(ancestorIsUniform)) { return; }
+      if (ancestors.some(nodeIsUniform)) { return; }
       // If the left hand side of an expression is one of these types,
       // we should construct a node from it.
       const unsafeTypes = ['Literal', 'ArrayExpression', 'Identifier'];
@@ -221,7 +240,7 @@ function shadergenerator(p5, fn) {
           type: 'CallExpression',
           callee: {
             type: 'Identifier',
-            name: 'dynamicNode',
+            name: 'p5.dynamicNode',
           },
           arguments: [node.left]
         }
@@ -1010,7 +1029,7 @@ function shadergenerator(p5, fn) {
     return length
   }
 
-  fn.dynamicNode = function (input) {
+  p5.dynamicNode = function (input) {
     if (isShaderNode(input)) {
       return input;
     }
@@ -1023,8 +1042,8 @@ function shadergenerator(p5, fn) {
   }
 
   // For replacing unary expressions
-  fn.unaryNode = function(input, sign) {
-    input = dynamicNode(input);
+  p5.unaryNode = function(input, sign) {
+    input = p5.dynamicNode(input);
     return dynamicAddSwizzleTrap(new UnaryExpressionNode(input, sign));
   }
 
@@ -1131,6 +1150,7 @@ function shadergenerator(p5, fn) {
       }
 
       const windowOverrides = {};
+      const fnOverrides = {};
 
       Object.keys(availableHooks).forEach((hookName) => {
         const hookTypes = originalShader.hookTypes(hookName);
@@ -1166,7 +1186,7 @@ function shadergenerator(p5, fn) {
           // If the expected return type is a struct we need to evaluate each of its properties
           if (!isGLSLNativeType(expectedReturnType.typeName)) {
             Object.entries(returnedValue).forEach(([propertyName, propertyNode]) => {
-              propertyNode = dynamicNode(propertyNode);
+              propertyNode = p5.dynamicNode(propertyNode);
               toGLSLResults[propertyName] = propertyNode.toGLSLBase(this.context);
               this.context.updateComponents(propertyNode);
             });
@@ -1216,18 +1236,25 @@ function shadergenerator(p5, fn) {
           this.resetGLSLContext();
         }
         windowOverrides[hookTypes.name] = window[hookTypes.name];
+        fnOverrides[hookTypes.name] = fn[hookTypes.name];
 
         // Expose the Functions to global scope for users to use
         window[hookTypes.name] = function(userOverride) {
           GLOBAL_SHADER[hookTypes.name](userOverride);
         };
+        fn[hookTypes.name] = function(userOverride) {
+          GLOBAL_SHADER[hookTypes.name](userOverride);
+        };
       });
 
 
       this.cleanup = () => {
         for (const key in windowOverrides) {
           window[key] = windowOverrides[key];
         }
+        for (const key in fnOverrides) {
+          fn[key] = fnOverrides[key];
+        }
       };
     }
 
@@ -1636,14 +1663,14 @@ function shadergenerator(p5, fn) {
     if (!GLOBAL_SHADER?.isGenerating) {
       return originalNoise.apply(this, args); // fallback to regular p5.js noise
     }
-    
+
     GLOBAL_SHADER.output.vertexDeclarations.add(noiseGLSL);
     GLOBAL_SHADER.output.fragmentDeclarations.add(noiseGLSL);
     return fnNodeConstructor('noise', args, { args: ['vec2'], returnType: 'float' });
   };
 }
-  
-  
+
+
 export default shadergenerator;
 
 if (typeof p5 !== 'undefined') {

src/webgl/p5.Shader.js

@@ -279,13 +279,96 @@ class Shader {
    * of shader code replacing default behaviour.
    *
    * Each shader may let you override bits of its behavior. Each bit is called
-   * a *hook.* A hook is either for the *vertex* shader, if it affects the
-   * position of vertices, or in the *fragment* shader, if it affects the pixel
-   * color. You can inspect the different hooks available by calling
+   * a *hook.* For example, a hook can let you adjust positions of vertices, or
+   * the color of a pixel. You can inspect the different hooks available by calling
    * <a href="#/p5.Shader/inspectHooks">`yourShader.inspectHooks()`</a>. You can
    * also read the reference for the default material, normal material, color, line, and point shaders to
    * see what hooks they have available.
    *
+   * `modify()` can be passed a function as a parameter. Inside, you can override hooks
+   * by calling them as functions. Each hook will take in a callback that takes in inputs
+   * and is expected to return an output. For example, here is a function that changes the
+   * material color to red:
+   *
+   * ```js example
+   * let myShader;
+   *
+   * function setup() {
+   *   createCanvas(200, 200, WEBGL);
+   *   myShader = baseMaterialShader().modify(() => {
+   *     getPixelInputs((inputs) => {
+   *       inputs.color = [1, 0, 0, 1];
+   *       return inputs;
+   *     });
+   *   });
+   * }
+   *
+   * function draw() {
+   *   background(255);
+   *   shader(myShader); // Apply the custom shader
+   *   plane(width, height); // Draw a plane with the shader applied
+   * }
+   * ```
+   *
+   * In addition to calling hooks, you can create uniforms, which are special variables
+   * used to pass data from p5.js into the shader. They can be created by calling `uniform` + the
+   * type of the data, such as `uniformFloat` for a number of `uniformVector2` for a two-component vector.
+   * They take in a function that returns the data for the variable. You can then reference these
+   * variables in your hooks, and their values will update every time you apply
+   * the shader with the result of your function.
+   *
+   * ```js example
+   * let myShader;
+   *
+   * function setup() {
+   *   createCanvas(200, 200, WEBGL);
+   *   myShader = baseMaterialShader().modify(() => {
+   *     // Get the current time from p5.js
+   *     let t = uniformFloat(() => millis());
+   *
+   *     getPixelInputs((inputs) => {
+   *       inputs.color = [sin(t * 0.01) / 2 + 0.5, 0, 0, 1];
+   *       return inputs;
+   *     });
+   *   });
+   * }
+   *
+   * function draw() {
+   *   background(255);
+   *   shader(myShader); // Apply the custom shader
+   *   plane(width, height); // Draw a plane with the shader applied
+   * }
+   * ```
+   *
+   * p5.strands functions are special, since they get turned into a shader instead of being
+   * run like the rest of your code. They only have access to p5.js functions, and variables
+   * you declare inside the `modify` callback. If you need access to local variables, you
+   * can pass them into `modify` with an optional second parameter, `variables`. If you are
+   * using instance mode, you will need to pass your sketch object in this way.
+   *
+   * ```js example
+   * new p5((sketch) => {
+   *   let myShader;
+   *
+   *   sketch.setup = function() {
+   *     sketch.createCanvas(200, 200, sketch.WEBGL);
+   *     myShader = sketch.baseMaterialShader().modify(() => {
+   *       sketch.getPixelInputs((inputs) => {
+   *         inputs.color = [1, 0, 0, 1];
+   *         return inputs;
+   *       });
+   *     }, { sketch });
+   *   }
+   *
+   *   sketch.draw = function() {
+   *     sketch.background(255);
+   *     sketch.shader(myShader); // Apply the custom shader
+   *     sketch.plane(sketch.width, sketch.height); // Draw a plane with the shader applied
+   *   }
+   * });
+   * ```
+   *
+   * You can also write GLSL directly in `modify` if you need direct access. To do so,
    * `modify()` takes one parameter, `hooks`, an object with the hooks you want
    * to override. Each key of the `hooks` object is the name
    * of a hook, and the value is a string with the GLSL code for your hook.
@@ -298,18 +381,7 @@ class Shader {
    * a default value as its value. These will be automatically set when the shader is set
    * with `shader(yourShader)`.
    *
-   * You can also add a `declarations` key, where the value is a GLSL string declaring
-   * custom uniform variables, globals, and functions shared
-   * between hooks. To add declarations just in a vertex or fragment shader, add
-   * `vertexDeclarations` and `fragmentDeclarations` keys.
-   *
-   * @beta
-   * @param {Object} [hooks] The hooks in the shader to replace.
-   * @returns {p5.Shader}
-   *
-   * @example
-   * <div modernizr='webgl'>
-   * <code>
+   * ```js example
    * let myShader;
    *
    * function setup() {
@@ -334,12 +406,14 @@ class Shader {
    *   fill('red');      // Set fill color to red
    *   sphere(50);       // Draw a sphere with the shader applied
    * }
-   * </code>
-   * </div>
+   * ```
    *
-   * @example
-   * <div modernizr='webgl'>
-   * <code>
+   * You can also add a `declarations` key, where the value is a GLSL string declaring
+   * custom uniform variables, globals, and functions shared
+   * between hooks. To add declarations just in a vertex or fragment shader, add
+   * `vertexDeclarations` and `fragmentDeclarations` keys.
+   *
+   * ```js example
    * let myShader;
    *
    * function setup() {
@@ -364,8 +438,17 @@ class Shader {
    *   fill('red');
    *   sphere(50);
    * }
-   * </code>
-   * </div>
+   * ```
+   *
+   * @beta
+   * @param {Function} callback A function with p5.strands code to modify the shader.
+   * @param {Object} [variables] An optional object with local variables p5.strands
+   * should have access to.
+   * @returns {p5.Shader}
+   */
+  /**
+   * @param {Object} [hooks] The hooks in the shader to replace.
+   * @returns {p5.Shader}
    */
   modify(hooks) {
     // p5._validateParameters('p5.Shader.modify', arguments);