feuerwerk/tests/core/shader-functions.test.ts

/**
 * Shader functions (#303): a graph of its own, used as a node in other shaders and set in by the
 * compiler where it is called.
 */
import { describe, expect, it } from 'vitest';
import { CommandError } from '../../src/core/commands/Command';
import { CreateAssetCommand } from '../../src/core/commands/assetCommands';
import { AddNodeCommand, ConnectCommand } from '../../src/core/commands/graphCommands';
import { AddPortCommand, RemovePortCommand, RenamePortCommand } from '../../src/core/commands/portCommands';
import { callProblem, FUNCTION_PARAM, resolveDef, SHADER_CALL } from '../../src/core/graph';
import { SHADER_FUNCTION_INPUT, SHADER_FUNCTION_OUTPUT } from '../../src/core/graph/nested';
import { graphLookup, shaderFunctionsOf } from '../../src/core/model/graphLookup';
import type { ShaderAsset, ShaderFunctionAsset } from '../../src/core/model/Asset';
import { findAssetOf, findGraphAsset } from '../../src/core/model/Project';
import { SHADER_PROPERTY } from '../../src/core/properties/shaders';
import { checkShaderFunction, compileShader } from '../../src/core/shader/compile';
import {
  SHADER_BLUR,
  SHADER_OUTPUT,
  SHADER_PASS,
  SHADER_TEXTURE_COLOR,
  SHADER_TIME,
  SHADER_VERTEX_OUTPUT,
} from '../../src/core/shader/nodes';
import { deserializeProject, serializeProject } from '../../src/core/store/serialize';
import { paletteEntries } from '../../src/editor/graph/paletteEntries';
import { shaderLibrary } from '../../src/viewport/ShaderLibrary';
import { makeSpriteScene, type SceneFixture } from '../helpers';

type Fixture = SceneFixture;

/** A graph asset with the helpers every test here builds with. */
function graphOf<T extends 'shader' | 'shaderfunction'>(f: Fixture, type: T, name: string) {
  const create = new CreateAssetCommand(type, name);
  f.stack.execute(create);
  const id = create.asset.id;
  const asset = () =>
    findAssetOf(f.project, id, type)! as T extends 'shader' ? ShaderAsset : ShaderFunctionAsset;
  const graph = () => findGraphAsset(f.project, id)!.graph;
  const add = (nodeType: string, params = {}) => {
    const command = new AddNodeCommand(id, nodeType, 0, 0, params);
    f.stack.execute(command);
    return command.node.id;
  };
  const link = (from: string, fromPort: string, to: string, toPort: string) =>
    f.stack.execute(new ConnectCommand(id, { node: from, port: fromPort }, { node: to, port: toPort }));
  /** An edge no command would lay, for what a hand-edited file may hold. */
  const wire = (from: string, fromPort: string, to: string, toPort: string) =>
    graph().edges.push({
      id: `w${graph().edges.length}`,
      from: { node: from, port: fromPort },
      to: { node: to, port: toPort },
    });
  const port = (node: string, key: string, kind: 'number' | 'color') =>
    f.stack.execute(new AddPortCommand(id, node, { key, kind }));
  return { id, asset, add, link, wire, port };
}

/** "Grayscale": a color in, its luminance out as a color. */
function grayscale(f: Fixture, name = 'Grayscale') {
  const fn = graphOf(f, 'shaderfunction', name);
  const input = fn.add(SHADER_FUNCTION_INPUT);
  const output = fn.add(SHADER_FUNCTION_OUTPUT);
  fn.port(input, 'color', 'color');
  fn.port(output, 'gray', 'color');
  const luminance = fn.add('ShaderLuminance');
  fn.link(input, 'color', luminance, 'color');
  fn.link(luminance, 'value', output, 'gray');
  return { ...fn, input, output, luminance };
}

/** A shader handing its texel through a call of `fnId` to the Output. */
function calling(f: Fixture, fnId: string, name = 'Uses') {
  const s = graphOf(f, 'shader', name);
  const out = s.add(SHADER_OUTPUT);
  const texel = s.add(SHADER_TEXTURE_COLOR);
  const call = s.add(SHADER_CALL, { [FUNCTION_PARAM]: fnId });
  const build = () => compileShader(s.asset().graph, shaderFunctionsOf(f.project));
  return { ...s, out, texel, call, build };
}

describe('shader functions (#303)', () => {
  it('draws through a call in two shaders exactly what the same nodes draw in place', () => {
    const f = makeSpriteScene();
    const fn = grayscale(f);
    const direct = graphOf(f, 'shader', 'Direct');
    const out = direct.add(SHADER_OUTPUT);
    const luminance = direct.add('ShaderLuminance');
    direct.link(direct.add(SHADER_TEXTURE_COLOR), 'color', luminance, 'color');
    direct.link(luminance, 'value', out, 'color');
    const expected = compileShader(direct.asset().graph).compiled!.fragment;

    for (const name of ['First', 'Second']) {
      const s = calling(f, fn.id, name);
      s.link(s.texel, 'color', s.call, 'color');
      s.link(s.call, 'gray', s.out, 'color');
      const build = s.build();
      expect(build.problems).toEqual({});
      expect(build.compiled!.fragment).toBe(expected);
    }
  });

  it('takes a change to the function into every shader calling it', () => {
    const f = makeSpriteScene();
    const fn = grayscale(f);
    const shaders = ['First', 'Second'].map((name) => {
      const s = calling(f, fn.id, name);
      s.link(s.texel, 'color', s.call, 'color');
      s.link(s.call, 'gray', s.out, 'color');
      return s;
    });
    const before = shaders.map((s) => shaderLibrary.buildFor(s.asset(), f.project).compiled!.fragment);
    const compiles = shaderLibrary.compiles;
    expect(shaders.map((s) => shaderLibrary.buildFor(s.asset(), f.project))).toBeTruthy();
    expect(shaderLibrary.compiles).toBe(compiles);

    const sin = fn.add('ShaderSin');
    fn.link(fn.luminance, 'value', sin, 'value');
    fn.link(sin, 'result', fn.output, 'gray');
    const after = shaders.map((s) => shaderLibrary.buildFor(s.asset(), f.project).compiled!.fragment);
    expect(shaderLibrary.compiles).toBe(compiles + 2);
    for (const [i, fragment] of after.entries()) {
      expect(fragment).not.toBe(before[i]);
      expect(fragment).toMatch(/= sin\(n\d+_value\);/);
    }
  });

  it('grows the call in another shader with a new input, and an undo takes it back in both', () => {
    const f = makeSpriteScene();
    const fn = grayscale(f);
    const s = calling(f, fn.id);
    const lookup = graphLookup(f.project);
    const inputsOfCall = () => {
      const node = s.asset().graph.nodes.find((n) => n.id === s.call)!;
      return resolveDef(node, s.asset().graph, lookup)!.inputs.map((p) => `${p.key}:${p.type}`);
    };
    expect(inputsOfCall()).toEqual(['color:vec3']);
    fn.port(fn.input, 'amount', 'number');
    expect(inputsOfCall()).toEqual(['color:vec3', 'amount:number']);
    f.stack.undo();
    expect(inputsOfCall()).toEqual(['color:vec3']);
    expect(
      fn
        .asset()
        .graph.nodes.find((n) => n.id === fn.input)!
        .ports!.map((p) => p.key),
    ).toEqual(['color']);
  });

  it('carries the edges at every call along when a port is renamed or removed', () => {
    const f = makeSpriteScene();
    const fn = grayscale(f);
    const s = calling(f, fn.id);
    s.link(s.texel, 'color', s.call, 'color');
    s.link(s.call, 'gray', s.out, 'color');
    const ports = () => s.asset().graph.edges.map((e) => (e.to.node === s.call ? e.to.port : e.from.port));

    f.stack.execute(new RenamePortCommand(fn.id, fn.input, 'color', 'tone'));
    expect(ports()).toEqual(['tone', 'gray']);
    expect(s.build().problems).toEqual({});

    f.stack.execute(new RemovePortCommand(fn.id, fn.output, 'gray'));
    expect(ports()).toEqual(['tone']);
    f.stack.undo();
    expect(ports()).toEqual(['tone', 'gray']);
    expect(s.build().compiled).not.toBeNull();
  });

  it('declares numbers and colors only', () => {
    const f = makeSpriteScene();
    const fn = grayscale(f);
    expect(() => f.stack.execute(new AddPortCommand(fn.id, fn.input, { key: 'name', kind: 'text' }))).toThrow(
      CommandError,
    );
    expect(() =>
      f.stack.execute(new AddPortCommand(fn.id, fn.output, { key: 'more', kind: 'flow' })),
    ).toThrow(CommandError);
  });

  it('reports a function calling itself, directly or through another, and compiles nothing endlessly', () => {
    const f = makeSpriteScene();
    const a = grayscale(f, 'A');
    const b = grayscale(f, 'B');
    const lookup = graphLookup(f.project);
    const self = a.add(SHADER_CALL, { [FUNCTION_PARAM]: a.id });
    expect(
      callProblem(
        a.asset().graph.nodes.find((n) => n.id === self)!,
        lookup,
        a.id,
      ),
    ).toBe('A function cannot call itself');
    a.asset().graph.nodes.splice(
      a.asset().graph.nodes.findIndex((n) => n.id === self),
      1,
    );
    const toB = a.add(SHADER_CALL, { [FUNCTION_PARAM]: b.id });
    const toA = b.add(SHADER_CALL, { [FUNCTION_PARAM]: a.id });
    a.port(a.output, 'twice', 'color');
    a.wire(a.input, 'color', toB, 'color');
    a.wire(toB, 'gray', a.output, 'twice');
    expect(
      callProblem(
        b.asset().graph.nodes.find((n) => n.id === toA)!,
        lookup,
        b.id,
      ),
    ).toBe('A function cannot call itself');
    expect(checkShaderFunction(a.asset().graph, shaderFunctionsOf(f.project), a.id).problems[toB]).toBe(
      'In B: A function cannot call itself',
    );

    const s = calling(f, a.id);
    s.link(s.texel, 'color', s.call, 'color');
    s.link(s.call, 'gray', s.out, 'color');
    const build = s.build();
    expect(build.compiled).toBeNull();
    expect(build.problems[s.call]).toBe('In A: In B: A function cannot call itself');
    expect(paletteEntries(f.project, b.asset()).some((e) => e.id === `${SHADER_CALL}:${a.id}`)).toBe(false);
  });

  it('feeds an unconnected input of the call, and an unconnected output, with the value typed in', () => {
    const f = makeSpriteScene();
    const fn = graphOf(f, 'shaderfunction', 'Scale');
    const input = fn.add(SHADER_FUNCTION_INPUT);
    const output = fn.add(SHADER_FUNCTION_OUTPUT, { spare: 0.25 });
    fn.port(input, 'amount', 'number');
    fn.port(output, 'result', 'number');
    fn.port(output, 'spare', 'number');
    fn.link(input, 'amount', output, 'result');
    const s = calling(f, fn.id);
    s.asset().graph.nodes.find((n) => n.id === s.call)!.params.amount = 0.75;
    const add = s.add('ShaderAdd');
    s.link(s.call, 'result', add, 'a');
    s.link(s.call, 'spare', add, 'b');
    s.link(add, 'result', s.out, 'alpha');
    const fragment = s.build().compiled!.fragment;
    expect(fragment).toMatch(/float n\d+_value = 0\.75;/);
    expect(fragment).toMatch(/float n\d+_value = 0\.25;/);
  });

  it('sets in two calls of one function apart, and a Time or a Pass inside counts for the caller', () => {
    const f = makeSpriteScene();
    const fn = graphOf(f, 'shaderfunction', 'Wide blur');
    const output = fn.add(SHADER_FUNCTION_OUTPUT);
    fn.port(output, 'color', 'color');
    const pass = fn.add(SHADER_PASS);
    const blur = fn.add(SHADER_BLUR);
    const time = fn.add(SHADER_TIME);
    fn.link(pass, 'picture', blur, 'picture');
    fn.link(time, 'value', blur, 'radius');
    fn.link(blur, 'color', output, 'color');
    const s = calling(f, fn.id);
    const second = s.add(SHADER_CALL, { [FUNCTION_PARAM]: fn.id });
    const mix = s.add('ShaderMix');
    s.link(s.call, 'color', mix, 'a');
    s.link(second, 'color', mix, 'b');
    s.link(mix, 'result', s.out, 'color');
    const compiled = s.build().compiled!;
    expect(compiled.passes).toHaveLength(2);
    expect(compiled.time).toBe(true);
    expect(compiled.fragment.match(/= blur\(/g)).toHaveLength(2);
  });

  it('reports what inside a function goes wrong at the call, and what has no place in it at the node', () => {
    const f = makeSpriteScene();
    const fn = grayscale(f);
    const property = fn.add(SHADER_PROPERTY);
    const own = checkShaderFunction(fn.asset().graph, shaderFunctionsOf(f.project), fn.id);
    expect(own.problems[property]).toBe('A function cannot hold Property');
    const s = calling(f, fn.id);
    s.link(s.call, 'gray', s.out, 'color');
    expect(s.build().problems[s.call]).toBe('In Grayscale: A function cannot hold Property');

    const lone = graphOf(f, 'shaderfunction', 'Lone');
    expect(checkShaderFunction(lone.asset().graph, shaderFunctionsOf(f.project), lone.id).reason).toBe(
      'No Function output node',
    );
    const typed = grayscale(f, 'Typed');
    typed.port(typed.output, 'level', 'number');
    typed.wire(typed.input, 'color', typed.output, 'level');
    expect(checkShaderFunction(typed.asset().graph, shaderFunctionsOf(f.project), typed.id).problems).toEqual(
      {
        [typed.output]: '"level" needs a number, not a color',
      },
    );
  });

  it('offers a function its own ends and no Output, Property or Vertex output, and a shader every function', () => {
    const f = makeSpriteScene();
    const fn = grayscale(f);
    const fresh = graphOf(f, 'shaderfunction', 'Fresh');
    const types = (asset: Parameters<typeof paletteEntries>[1]) =>
      paletteEntries(f.project, asset).map((e) => e.id);
    const inFunction = types(fresh.asset());
    expect(inFunction).toEqual(
      expect.arrayContaining([SHADER_FUNCTION_INPUT, SHADER_FUNCTION_OUTPUT, SHADER_TIME]),
    );
    for (const absent of [SHADER_OUTPUT, SHADER_VERTEX_OUTPUT, SHADER_PROPERTY, `${SHADER_CALL}:${fresh.id}`])
      expect(inFunction).not.toContain(absent);
    expect(inFunction).toContain(`${SHADER_CALL}:${fn.id}`);
    expect(types(fn.asset())).not.toContain(SHADER_FUNCTION_INPUT);

    const s = calling(f, fn.id);
    const inShader = types(s.asset());
    expect(inShader).not.toContain(SHADER_FUNCTION_INPUT);
    const entry = paletteEntries(f.project, s.asset()).find((e) => e.id === `${SHADER_CALL}:${fn.id}`)!;
    expect(entry.def.label).toBe('Grayscale');
  });

  it('keeps a function and its declared ports through saving and loading', () => {
    const f = makeSpriteScene();
    const fn = grayscale(f);
    const loaded = deserializeProject(serializeProject(f.project));
    const back = findAssetOf(loaded, fn.id, 'shaderfunction')!;
    expect(back.graph.nodes.map((n) => [n.type, n.ports?.map((p) => `${p.key}:${p.kind}`) ?? []])).toEqual(
      fn.asset().graph.nodes.map((n) => [n.type, n.ports?.map((p) => `${p.key}:${p.kind}`) ?? []]),
    );
    expect(back.graph.edges).toHaveLength(2);
  });
});