@lightningjs/renderer/src/core/renderers/webgl/WebGlShaderProgram.ts

Lightning 3 Renderer

/*
 * If not stated otherwise in this file or this component's LICENSE file the
 * following copyright and licenses apply:
 *
 * Copyright 2023 Comcast Cable Communications Management, LLC.
 *
 * Licensed under the Apache License, Version 2.0 (the License);
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
import type { GlContextWrapper } from '../../platforms/GlContextWrapper.js';
import { Default } from '../../shaders/webgl/Default.js';
import type { CoreShaderProgram } from '../CoreShaderProgram.js';
import type { WebGlCtxTexture } from './WebGlCtxTexture.js';
import type { WebGlRenderer, WebGlNodeRenderOp } from './WebGlRenderer.js';
import type { WebGlShaderType } from './WebGlShaderNode.js';
import { WebGlShaderNode } from './WebGlShaderNode.js';
import type { BufferCollection } from './internal/BufferCollection.js';
import {
  createProgram,
  createShader,
  type UniformSet1Param,
  type UniformSet2Params,
  type UniformSet3Params,
  type UniformSet4Params,
} from './internal/ShaderUtils.js';
import { CoreNode } from '../../CoreNode.js';
import type { Dimensions } from '../../../common/CommonTypes.js';
import type { Stage } from '../../Stage.js';

/**
 * Structural surface shared by node render ops (CoreNode / CoreTextNode) and
 * batched SdfRenderOps that {@link WebGlShaderProgram.bindRenderOp} consumes.
 *
 * Node ops expose `parentFramebufferDimensions`, SdfRenderOp exposes
 * `framebufferDimensions`; `isCoreNode` picks the active one. Keeping this as
 * an interface (rather than widening to the full `WebGlRenderOp` union)
 * avoids demanding properties from StencilClipRenderOp, which is never bound
 * through this path.
 */
export interface WebGlBoundRenderOp {
  renderOpTextures: WebGlCtxTexture[];
  quadBufferCollection: BufferCollection;
  isCoreNode: boolean;
  rtt: boolean;
  parentHasRenderTexture: boolean;
  parentFramebufferDimensions?: Dimensions | null;
  framebufferDimensions?: Dimensions | null;
  stage: Stage;
  time: number;
  worldAlpha: number;
  w: number;
  h: number;
  shader: WebGlShaderNode | null;
}

export class WebGlShaderProgram implements CoreShaderProgram {
  protected program: WebGLProgram | null;
  /**
   * Vertex Array Object
   *
   * @remarks
   * Used by WebGL2 Only
   */
  protected vao: WebGLVertexArrayObject | undefined;
  protected renderer: WebGlRenderer;
  glw: GlContextWrapper;
  protected attributeLocations: string[];
  protected uniformLocations: Record<string, WebGLUniformLocation> | null;
  protected lifecycle: Pick<WebGlShaderType, 'update' | 'canBatch'>;
  protected useSystemAlpha = false;
  protected useSystemDimensions = false;
  protected useTimeValue = false;
  public isDestroyed = false;
  supportsIndexedTextures = false;

  /**
   * Shadow copies of system uniform values. Used by bindRenderOp to skip
   * redundant gl.uniform* calls when the value hasn't changed.
   *
   * Each gl.uniform* call crosses into the GPU-process command buffer, so
   * skipping value-identical re-uploads is a real per-op CPU saving on
   * embedded targets.
   *
   * Sentinel -1 never collides with real values (all >= 0).
   */
  private lastPixelRatio = -1;
  private lastResolutionW = -1;
  private lastResolutionH = -1;
  private lastAlpha = -1;
  private lastDimensionsW = -1;
  private lastDimensionsH = -1;
  private lastTime = -1;

  /**
   * Last-bound shader uniform collection by identity reference.
   *
   * Uniform collections are created once, filled once, and shared by
   * reference across shader nodes with the same value key — so reference
   * equality implies value equality. Allows skipping the entire for-in
   * loop over uniform buckets when the same collection is already bound.
   */
  private lastBoundUniforms: unknown = null;

  constructor(
    renderer: WebGlRenderer,
    config: WebGlShaderType,
    resolvedProps: Record<string, any>,
  ) {
    this.renderer = renderer;
    const glw = (this.glw = renderer.glw);

    // Check that extensions are supported
    const webGl2 = glw.isWebGl2();
    let requiredExtensions: string[] = [];

    this.supportsIndexedTextures =
      config.supportsIndexedTextures || this.supportsIndexedTextures;
    requiredExtensions =
      (webGl2 && config.webgl2Extensions) ||
      (!webGl2 && config.webgl1Extensions) ||
      [];

    const glVersion = webGl2 ? '2.0' : '1.0';
    requiredExtensions.forEach((extensionName) => {
      if (!glw.getExtension(extensionName)) {
        throw new Error(
          `Shader "${this.constructor.name}" requires extension "${extensionName}" for WebGL ${glVersion} but wasn't found`,
        );
      }
    });

    let vertexSource =
      config.vertex instanceof Function
        ? config.vertex(renderer, resolvedProps)
        : config.vertex;

    if (vertexSource === undefined) {
      vertexSource = Default.vertex as string;
    }

    const fragmentSource =
      config.fragment instanceof Function
        ? config.fragment(renderer, resolvedProps)
        : config.fragment;

    const vertexShader = createShader(glw, glw.VERTEX_SHADER, vertexSource);
    if (!vertexShader) {
      throw new Error('Vertex shader creation failed');
    }

    const fragmentShader = createShader(
      glw,
      glw.FRAGMENT_SHADER,
      fragmentSource,
    );

    if (!fragmentShader) {
      throw new Error('fragment shader creation failed');
    }

    const program = createProgram(glw, vertexShader, fragmentShader);
    this.program = program;
    this.attributeLocations = glw.getAttributeLocations(program);

    const uniLocs = (this.uniformLocations = glw.getUniformLocations(program));

    this.useSystemAlpha = uniLocs['u_alpha'] !== undefined;
    this.useSystemDimensions = uniLocs['u_dimensions'] !== undefined;

    this.useTimeValue =
      this.glw.getUniformLocation(program, 'u_dimensions') !== null &&
      config.time !== undefined;

    this.lifecycle = {
      update: config.update,
      canBatch: config.canBatch,
    };
  }

  disableAttribute(location: number) {
    this.glw.disableVertexAttribArray(location);
  }

  disableAttributes() {
    const glw = this.glw;
    const attribLen = this.attributeLocations.length;
    for (let i = 0; i < attribLen; i++) {
      glw.disableVertexAttribArray(i);
    }
  }

  reuseRenderOp(node: CoreNode, currentRenderOp: WebGlNodeRenderOp): boolean {
    if (this.lifecycle.canBatch !== undefined) {
      return this.lifecycle.canBatch(node, currentRenderOp);
    }

    if (this.useTimeValue === true) {
      if (node.time !== currentRenderOp.time) {
        return false;
      }
    }

    if (this.useSystemAlpha === true) {
      if (node.worldAlpha !== currentRenderOp.worldAlpha) {
        return false;
      }
    }

    if (this.useSystemDimensions === true) {
      if (node.w !== currentRenderOp.w || node.h !== currentRenderOp.h) {
        return false;
      }
    }

    const shader = node.props.shader as WebGlShaderNode | null;
    const opShader = currentRenderOp.shader as WebGlShaderNode | null;

    // Same shader node — same resolved props by definition.
    if (shader === opShader) {
      return true;
    }

    if (shader === null || opShader === null) {
      return false;
    }

    // Uniform collections are shared by reference across shader nodes with
    // equal value keys — reference equality implies the resolved prop values
    // match without a key-by-key compare.
    if (shader.uniforms === opShader.uniforms) {
      return true;
    }

    const shaderPropsA = shader.resolvedProps as
      | Record<string, unknown>
      | undefined;
    const shaderPropsB = opShader.resolvedProps as
      | Record<string, unknown>
      | undefined;

    if (
      (shaderPropsA === undefined && shaderPropsB !== undefined) ||
      (shaderPropsA !== undefined && shaderPropsB === undefined)
    ) {
      return false;
    }

    if (shaderPropsA !== undefined && shaderPropsB !== undefined) {
      for (const key in shaderPropsA) {
        if (shaderPropsA[key] !== shaderPropsB[key]) {
          return false;
        }
      }
    }

    return true;
  }

  bindRenderOp(renderOp: WebGlBoundRenderOp) {
    this.bindTextures(renderOp.renderOpTextures);
    this.bindBufferCollection(renderOp.quadBufferCollection);

    const parentHasRenderTexture = renderOp.parentHasRenderTexture;
    const framebufferDimensions = renderOp.isCoreNode
      ? renderOp.parentFramebufferDimensions
      : renderOp.framebufferDimensions;

    // Skip if the parent and current operation both have render textures
    if (renderOp.rtt === true && parentHasRenderTexture === true) {
      return;
    }

    // Resolve target pixel ratio / resolution, then compare-and-set against
    // the program's shadow state to skip value-identical re-uploads.
    let pixelRatio: number;
    let resolutionW: number;
    let resolutionH: number;
    if (parentHasRenderTexture === true && framebufferDimensions) {
      pixelRatio = 1.0;
      resolutionW = framebufferDimensions.w;
      resolutionH = framebufferDimensions.h;
    } else {
      pixelRatio = renderOp.stage.pixelRatio;
      resolutionW = this.glw.canvas.width;
      resolutionH = this.glw.canvas.height;
    }

    if (pixelRatio !== this.lastPixelRatio) {
      this.glw.uniform1f('u_pixelRatio', pixelRatio);
      this.lastPixelRatio = pixelRatio;
    }

    if (
      resolutionW !== this.lastResolutionW ||
      resolutionH !== this.lastResolutionH
    ) {
      this.glw.uniform2f('u_resolution', resolutionW, resolutionH);
      this.lastResolutionW = resolutionW;
      this.lastResolutionH = resolutionH;
    }

    if (this.useTimeValue === true && renderOp.time !== this.lastTime) {
      this.glw.uniform1f('u_time', renderOp.time);
      this.lastTime = renderOp.time;
    }

    if (
      this.useSystemAlpha === true &&
      renderOp.worldAlpha !== this.lastAlpha
    ) {
      this.glw.uniform1f('u_alpha', renderOp.worldAlpha);
      this.lastAlpha = renderOp.worldAlpha;
    }

    if (this.useSystemDimensions === true) {
      if (
        renderOp.w !== this.lastDimensionsW ||
        renderOp.h !== this.lastDimensionsH
      ) {
        this.glw.uniform2f('u_dimensions', renderOp.w, renderOp.h);
        this.lastDimensionsW = renderOp.w;
        this.lastDimensionsH = renderOp.h;
      }
    }

    const shader = renderOp.shader as WebGlShaderNode;
    const uniforms = shader.uniforms;

    if (shader.beforeDraw !== undefined) {
      shader.beforeDraw();
    }

    if (uniforms.hasStoredUniforms === true) {
      // Uniform collections are immutable after creation and shared by
      // reference across shader nodes with equal value keys — when the same
      // object is already bound, the GL program still holds these values.
      if ((uniforms as unknown) === this.lastBoundUniforms) {
        return;
      }
      this.lastBoundUniforms = uniforms;

      for (const key in uniforms.single) {
        const { method, value } = uniforms.single[key]!;
        this.glw[method as keyof UniformSet1Param](key, value as never);
      }

      for (const key in uniforms.vec2) {
        const { method, value } = uniforms.vec2[key]!;
        this.glw[method as keyof UniformSet2Params](key, value[0], value[1]);
      }

      for (const key in uniforms.vec3) {
        const { method, value } = uniforms.vec3[key]!;
        this.glw[method as keyof UniformSet3Params](
          key,
          value[0],
          value[1],
          value[2],
        );
      }

      for (const key in uniforms.vec4) {
        const { method, value } = uniforms.vec4[key]!;
        this.glw[method as keyof UniformSet4Params](
          key,
          value[0],
          value[1],
          value[2],
          value[3],
        );
      }
    }
  }

  bindBufferCollection(buffer: BufferCollection) {
    const { glw } = this;
    const attribs = this.attributeLocations;
    const attribLen = attribs.length;

    for (let i = 0; i < attribLen; i++) {
      const name = attribs[i]!;
      const resolvedBuffer = buffer.getBuffer(name);
      const resolvedInfo = buffer.getAttributeInfo(name);
      if (resolvedBuffer === undefined || resolvedInfo === undefined) {
        continue;
      }
      glw.enableVertexAttribArray(i);
      glw.vertexAttribPointer(
        resolvedBuffer,
        i,
        resolvedInfo.size,
        resolvedInfo.type,
        resolvedInfo.normalized,
        resolvedInfo.stride,
        resolvedInfo.offset,
      );
    }
  }

  bindTextures(textures: WebGlCtxTexture[]) {
    const t = textures[0];
    if (t === undefined) return;
    this.glw.activeTexture(0);
    this.glw.bindTexture(t.ctxTexture);
  }

  attach(): void {
    if (this.isDestroyed === true) {
      return;
    }
    this.glw.useProgram(this.program, this.uniformLocations!);
    if (this.glw.isWebGl2() && this.vao) {
      this.glw.bindVertexArray(this.vao);
    }
  }

  /**
   * Activate this program and bind a buffer collection without going through
   * the shader manager's detach/attach cycle. Used exclusively by the stencil
   * write pass so it does not disrupt the currently attached scene shader.
   *
   * The caller is responsible for restoring `shManager.attachedShader` to null
   * after the pass so the next real draw triggers a full re-attach.
   */
  bindForStencil(bufferCollection: BufferCollection): void {
    if (this.isDestroyed === true) {
      return;
    }
    this.glw.useProgram(this.program, this.uniformLocations!);
    this.bindBufferCollection(bufferCollection);
  }

  detach(): void {
    this.disableAttributes();
  }

  destroy() {
    if (this.isDestroyed === true) {
      return;
    }
    const glw = this.glw;

    this.detach();

    glw.deleteProgram(this.program!);
    this.program = null;
    this.uniformLocations = null;

    const attribs = this.attributeLocations;
    const attribLen = this.attributeLocations.length;
    for (let i = 0; i < attribLen; i++) {
      this.glw.deleteBuffer(attribs[i]!);
    }
  }
}