strudel/packages/shader/shader.mjs

/*
shader.mjs - implements the `loadShader` function
Copyright (C) 2024 Strudel contributors
This program is free software: you can redistribute it and/or modify it under the terms of the GNU Affero General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Affero General Public License for more details. You should have received a copy of the GNU Affero General Public License along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

import { PicoGL } from 'picogl';
import { logger } from '@strudel/core';

// The standard fullscreen vertex shader.
const vertexShader = `#version 300 es
precision highp float;
layout(location=0) in vec2 position;
void main() {
  gl_Position = vec4(position, 1, 1);
}
`;

// Make the fragment source, similar to the one from shadertoy.
function mkFragmentShader(code) {
  return `#version 300 es
precision highp float;
out vec4 oColor;
uniform float iTime;
uniform vec2 iResolution;

#define STRUDEL 1

${code}

void main(void) {
  mainImage(oColor, gl_FragCoord.xy);
}
`;
}

// Modulation helpers to smooth the values.
class UniformValue {
  constructor() {
    this.value = 0;
    this.desired = 0;
    this.slow = 10;
  }

  get(elapsed) {
    // Adjust the value according to the rate of change
    const offset = (this.desired - this.value) / (this.slow * Math.min(1, elapsed * 60));
    // Ignore small changes
    if (Math.abs(offset) > 1e-3) this.value += offset;
    return this.value;
  }
}

// Set an uniform value (from a pattern).
export function setUniform(instanceName, name, value, incr, position, slow) {
  const instance = _instances[instanceName];
  if (!instance) {
    logger('[shader] not loaded yet', 'warning');
    return;
  }

  // console.log('setUniform: ', name, value, position, slow);
  const uniform = instance.uniforms[name];
  if (uniform) {
    let uniformValue;
    if (uniform.count == 0) {
      // This is a single value
      uniformValue = uniform.value;
    } else {
      // This is an array
      const idx = position % uniform.value.length;
      uniformValue = uniform.value[idx];
    }
    uniformValue.slow = slow;
    // TODO: handle direct assignment, this is incrementing by default
    if (incr) uniformValue.desired += value;
    else uniformValue.desired = value;
  } else {
    logger('[shader] unknown uniform: ' + name);
  }

  // Ensure the instance is drawn
  instance.age = 0;
  if (!instance.drawing) {
    instance.drawing = requestAnimationFrame(instance.update);
  }
}

// Update the uniforms for a given drawFrame call.
function updateUniforms(drawFrame, elapsed, uniforms) {
  Object.values(uniforms).forEach((uniform) => {
    const value = uniform.count == 0 ? uniform.value.get(elapsed) : uniform.value.map((v) => v.get(elapsed));

    // Send the value to the GPU
    // console.log('updateUniforms:', uniform.name, value);
    drawFrame.uniform(uniform.name, value);
  });
}

// Setup the instance's uniform after shader compilation.
function setupUniforms(uniforms, program) {
  Object.entries(program.uniforms).forEach(([name, uniform]) => {
    if (name != 'iTime' && name != 'iResolution') {
      // remove array suffix
      const uname = name.replace('[0]', '');
      const count = uniform.count | 0;
      if (!uniforms[uname] || uniforms[uname].count != count) {
        // TODO: keep the previous values when the count change:
        // if the count decreased, then drop the excess, else append new values
        uniforms[uname] = {
          name,
          count,
          value: count == 0 ? new UniformValue() : new Array(count).fill().map(() => new UniformValue()),
        };
      }
    }
  });
  // TODO: remove previous uniform that are no longer used...
  return uniforms;
}

// Setup the canvas and return the WebGL context.
function setupCanvas(name) {
  // TODO: support custom size
  const width = 400;
  const height = 300;
  const canvas = document.createElement('canvas');
  canvas.id = 'cnv-' + name;
  canvas.width = width;
  canvas.height = height;
  const top = 60 + Object.keys(_instances).length * height;
  canvas.style = `pointer-events:none;width:${width}px;height:${height}px;position:fixed;top:${top}px;right:23px`;
  document.body.append(canvas);
  return canvas.getContext('webgl2');
}

// Setup the shader instance
async function initializeShaderInstance(name, code) {
  // Setup PicoGL app
  const ctx = setupCanvas(name);
  const app = PicoGL.createApp(ctx);

  // Setup buffers
  const resolution = new Float32Array([ctx.canvas.width, ctx.canvas.height]);

  // Two triangle to cover the whole canvas
  const positionBuffer = app.createVertexBuffer(
    PicoGL.FLOAT,
    2,
    new Float32Array([-1, -1, -1, 1, 1, 1, 1, 1, 1, -1, -1, -1]),
  );

  // Setup the arrays
  const arrays = app.createVertexArray().vertexAttributeBuffer(0, positionBuffer);

  return app
    .createPrograms([vertexShader, code])
    .then(([program]) => {
      const drawFrame = app.createDrawCall(program, arrays);
      const instance = { app, code, program, arrays, drawFrame, uniforms: setupUniforms({}, program) };

      // Render frame logic
      let prev = performance.now() / 1000;
      instance.age = 0;
      instance.update = () => {
        const now = performance.now() / 1000;
        const elapsed = instance.age == 0 ? 1 / 60 : now - prev;
        prev = now;
        // console.log("drawing!")
        app.clear();
        instance.drawFrame.uniform('iResolution', resolution).uniform('iTime', now);

        updateUniforms(instance.drawFrame, elapsed, instance.uniforms);

        instance.drawFrame.draw();
        // After sometime, if no update happened, stop the animation loop
        if (instance.age++ < 100) requestAnimationFrame(instance.update);
        else instance.drawing = false;
      };
      return instance;
    })
    .catch((err) => {
      ctx.canvas.remove();
      throw err;
    });
}

// Update the instance program
async function reloadShaderInstanceCode(instance, code) {
  return instance.app.createPrograms([vertexShader, code]).then(([program]) => {
    instance.program.delete();
    instance.program = program;
    instance.uniforms = setupUniforms(instance.uniforms, program);
    instance.drawFrame = instance.app.createDrawCall(program, instance.arrays);
  });
}

// Keep track of the running shader instances
let _instances = {};
export async function loadShader(code = '', name = 'default') {
  if (code) {
    code = mkFragmentShader(code);
  }
  if (!_instances[name]) {
    _instances[name] = await initializeShaderInstance(name, code);
    logger('[shader] ready');
  } else if (_instances[name].code != code) {
    await reloadShaderInstanceCode(_instances[name], code);
    logger('[shader] reloaded');
  }
}