feuerwerk/tests/core/shaders.test.ts

/**
 * A shader is an asset whose graph is GLSL in nodes: `Property` nodes declare the rows a sprite
 * wearing it shows, the compiler turns everything reachable from `Output` into a fragment. The
 * one graph every test builds is the "Adjust colors" recipe (brightness, contrast, saturation).
 */
import { describe, expect, it } from 'vitest';
import { CreateAssetCommand, SetAssetPropertyCommand } from '../../src/core/commands/assetCommands';
import { AddNodeCommand, ConnectCommand } from '../../src/core/commands/graphCommands';
import { SetOverrideCommand } from '../../src/core/commands/sceneObjectCommands';
import type { Graph } from '../../src/core/graph';
import { BUILT_IN_NODES } from '../../src/core/graph/nodes';
import type { ShaderAsset } from '../../src/core/model/Asset';
import { isPlaceableType } from '../../src/core/model/AssetTypes';
import { findAssetOf } from '../../src/core/model/Project';
import { assetDeviates, objectDeviates } from '../../src/core/properties/deviation';
import {
  allShaderParamRows,
  schemaFor,
  schemaForAsset,
  schemaForObject,
} from '../../src/core/properties/PropertySchema';
import { resolveSceneObjectProps } from '../../src/core/properties/resolve';
import {
  nextPropertyName,
  SHADER_KEY,
  SHADER_PROPERTY,
  shaderParamDefs,
} from '../../src/core/properties/shaders';
import { inlineInputs } from '../../src/core/graph/NodeRegistry';
import { compileShader } from '../../src/core/shader/compile';
import {
  SHADER_NODES,
  SHADER_NUMBER,
  SHADER_OUTPUT,
  SHADER_TEXTURE_ALPHA,
  SHADER_TEXTURE_COLOR,
  ShaderOutput,
} from '../../src/core/shader/nodes';
import { deserializeProject } from '../../src/core/store/serialize';
import { paletteEntries } from '../../src/editor/graph/paletteEntries';
import { makeProject, makeSpriteScene, place, type ProjectFixture, roundTrip } from '../helpers';

/** Command-driven editing of a shader asset's graph, the way the tab does it. */
function shaderEditor(f: ProjectFixture, name = 'Adjust colors') {
  const create = new CreateAssetCommand('shader', name);
  f.stack.execute(create);
  const id = create.asset.id;
  const asset = () => findAssetOf(f.project, id, 'shader')!;
  const add = (type: string, params = {}) => {
    const command = new AddNodeCommand(id, type, 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 }));
  return { id, asset, add, link };
}

const property = (name: string, value = 1, min = 0, max = 4) => ({ name, default: value, min, max });

/**
 * rgb += brightness - 1; rgb = (rgb - 0.5) * contrast + 0.5; rgb = mix(luma, rgb, saturation) –
 * twelve nodes, the recipe `tools/lib/shaders.mjs` ships.
 */
function buildAdjustColors(s: ReturnType<typeof shaderEditor>) {
  const texel = s.add(SHADER_TEXTURE_COLOR);
  const brightness = s.add(SHADER_PROPERTY, property('brightness'));
  const contrast = s.add(SHADER_PROPERTY, property('contrast'));
  const saturation = s.add(SHADER_PROPERTY, property('saturation'));
  const offset = s.add('ShaderSubtract', { b: 1 });
  const bright = s.add('ShaderAdd');
  const shifted = s.add('ShaderSubtract', { b: 0.5 });
  const scaled = s.add('ShaderMultiply');
  const centred = s.add('ShaderAdd', { b: 0.5 });
  const luma = s.add('ShaderLuminance');
  const mixed = s.add('ShaderMix');
  const output = s.add(SHADER_OUTPUT);
  s.link(texel, 'color', bright, 'a');
  s.link(brightness, 'value', offset, 'a');
  s.link(offset, 'result', bright, 'b');
  s.link(bright, 'result', shifted, 'a');
  s.link(shifted, 'result', scaled, 'a');
  s.link(contrast, 'value', scaled, 'b');
  s.link(scaled, 'result', centred, 'a');
  s.link(centred, 'result', luma, 'color');
  s.link(luma, 'value', mixed, 'a');
  s.link(centred, 'result', mixed, 'b');
  s.link(saturation, 'value', mixed, 't');
  s.link(mixed, 'result', output, 'color');
  return {
    texel,
    brightness,
    contrast,
    saturation,
    offset,
    bright,
    shifted,
    scaled,
    centred,
    luma,
    mixed,
    output,
  };
}

describe('a shader asset', () => {
  it('is created with an empty graph, is not placeable, and is a graph asset for the commands', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    expect(isPlaceableType('shader')).toBe(false);
    expect(s.asset().graph).toEqual({ nodes: [], edges: [] });
    expect(schemaFor('shader')).toEqual([]);
    expect(compileShader(s.asset().graph).reason).toBe('No Output node');
  });

  it('offers shader nodes and nothing else, the Output only once, properties numbered', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    const offered = () => paletteEntries(f.project, s.asset()).map((e) => e.def.type);
    expect(offered()).toEqual(SHADER_NODES.map((d) => d.type));
    expect(offered()).not.toContain('OnStart');
    s.add(SHADER_OUTPUT);
    expect(offered()).not.toContain(SHADER_OUTPUT);
    const entry = () => paletteEntries(f.project, s.asset()).find((e) => e.def.type === SHADER_PROPERTY)!;
    expect(entry().params).toEqual({ name: 'value' });
    s.add(SHADER_PROPERTY, entry().params);
    expect(nextPropertyName(s.asset().graph)).toBe('value2');
    expect(entry().params).toEqual({ name: 'value2' });
  });

  it('keeps its nodes out of every behavior palette and out of the built-in set', () => {
    const f = makeProject();
    const behavior = new CreateAssetCommand('behavior', 'B');
    f.stack.execute(behavior);
    const asset = findAssetOf(f.project, behavior.asset.id, 'behavior')!;
    const offered = paletteEntries(f.project, asset).map((e) => e.def.type);
    for (const def of SHADER_NODES) expect(offered).not.toContain(def.type);
    expect(BUILT_IN_NODES.some((d) => d.category === 'shader')).toBe(false);
  });
});

describe('shaderParamDefs', () => {
  it('turns every well-named Property node into a number row, in node order', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    s.add(SHADER_PROPERTY, property('brightness'));
    s.add(SHADER_PROPERTY, { name: 'amount', default: 0.25, min: 0, max: 1 });
    s.add(SHADER_PROPERTY, { name: 'not a name', default: 1 });
    s.add(SHADER_PROPERTY, { name: 'brightness', default: 7 });
    const rows = shaderParamDefs(s.asset().graph);
    expect(rows.map((r) => r.key)).toEqual(['brightness', 'amount']);
    expect(rows[0]).toEqual({
      key: 'brightness',
      label: 'brightness',
      kind: 'number',
      default: 1,
      min: 0,
      max: 4,
      step: 0.05,
    });
    expect(rows[1]).toMatchObject({ default: 0.25, min: 0, max: 1, step: 0.05 });
  });

  it('drops a max at or below the min instead of pinning the value', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    s.add(SHADER_PROPERTY, { name: 'wide', default: 50, min: 10, max: 10 });
    const [row] = shaderParamDefs(s.asset().graph);
    expect(row).toMatchObject({ min: 10, default: 50 });
    expect(row).not.toHaveProperty('max');
  });
});

describe('allShaderParamRows', () => {
  it('offers one row per name, naming every shader that declares it', () => {
    const f = makeProject();
    const adjust = shaderEditor(f, 'Adjust colors');
    adjust.add(SHADER_PROPERTY, property('brightness'));
    const neon = shaderEditor(f, 'Neon glow');
    neon.add(SHADER_PROPERTY, property('brightness'));
    neon.add(SHADER_PROPERTY, property('glow'));
    expect(allShaderParamRows(f.project).map((row) => `${row.def.key} (${row.shaders.join(', ')})`)).toEqual([
      'brightness (Adjust colors, Neon glow)',
      'glow (Neon glow)',
    ]);
  });
});

describe('compileShader', () => {
  it('compiles the Adjust colors graph to a fragment with one uniform per property', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    buildAdjustColors(s);
    const build = compileShader(s.asset().graph);
    expect(build.problems).toEqual({});
    expect(build.reason).toBe('');
    const compiled = build.compiled!;
    expect(compiled.uniforms).toEqual([
      { name: 'brightness', default: 1 },
      { name: 'contrast', default: 1 },
      { name: 'saturation', default: 1 },
    ]);
    const { fragment } = compiled;
    for (const name of ['brightness', 'contrast', 'saturation'])
      expect(fragment).toContain(`uniform float u_${name};`);
    // Drawn in the sprite's own pass: the tint goes in before the graph, premultiplied like the texel.
    expect(fragment).toContain('vec4 texel = texture(uTexture, vTextureCoord) * vColor;');
    expect(fragment).toContain('vec3 rgb = texel.a > 0.0 ? texel.rgb / texel.a : texel.rgb;');
    expect(fragment).toMatch(/float n\d+_result = \(n\d+_value - 1\.0\);/);
    expect(fragment).toMatch(/vec3 n\d+_result = \(n\d+_color \+ vec3\(n\d+_result\)\);/);
    expect(fragment).toMatch(/vec3 n\d+_result = \(n\d+_result - vec3\(0\.5\)\);/);
    expect(fragment).toMatch(/float n\d+_value = dot\(n\d+_result, vec3\(0\.2126, 0\.7152, 0\.0722\)\);/);
    expect(fragment).toMatch(/vec3 n\d+_result = mix\(vec3\(n\d+_value\), n\d+_result, n\d+_value\);/);
    expect(fragment).toContain('float alpha = clamp(texel.a, 0.0, 1.0);');
    expect(fragment).toMatch(/finalColor = vec4\(clamp\(n\d+_result, 0\.0, 1\.0\) \* alpha, alpha\);/);
    // Each local is declared before it is used.
    const declared = new Set<string>();
    for (const line of fragment.split('\n')) {
      const m = /^\s+(?:float|vec3) (n\d+_\w+) = (.*);$/.exec(line);
      if (!m) continue;
      for (const used of m[2]!.match(/n\d+_\w+/g) ?? []) expect(declared.has(used)).toBe(true);
      declared.add(m[1]!);
    }
  });

  it('passes the texel through an Output with nothing connected, and types literals as floats', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    s.add(SHADER_OUTPUT);
    expect(compileShader(s.asset().graph).compiled!.fragment).toContain('clamp(rgb, 0.0, 1.0) * alpha');

    const g = shaderEditor(f, 'Scaled');
    const times = g.add('ShaderMultiply', { a: 2, b: 0.75 });
    const output = g.add(SHADER_OUTPUT);
    g.link(times, 'result', output, 'color');
    const { fragment } = compileShader(g.asset().graph).compiled!;
    // Two numbers make a number; feeding it to the colour broadcasts it.
    expect(fragment).toMatch(/float n0_result = \(2\.0 \* 0\.75\);/);
    expect(fragment).toContain('clamp(vec3(n0_result), 0.0, 1.0)');

    const c = shaderEditor(f, 'Tinted');
    const tint = c.add('ShaderColor', { value: 0xff8000 });
    const times2 = c.add('ShaderMultiply');
    const out2 = c.add(SHADER_OUTPUT);
    c.link(tint, 'color', times2, 'b');
    c.link(times2, 'result', out2, 'color');
    const tinted = compileShader(c.asset().graph).compiled!.fragment;
    expect(tinted).toContain('vec3 n0_color = vec3(1.0000, 0.5020, 0.0000);');
    // An unconnected `a` is its typed-in default, 1 for a Multiply, lifted to a colour.
    expect(tinted).toMatch(/vec3 n1_result = \(vec3\(1\.0\) \* n0_color\);/);
  });

  it('turns a Number node into one float literal: decimal point kept, sign kept, NaN made 0', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    const half = s.add(SHADER_NUMBER, { value: 0.5 });
    const two = s.add(SHADER_NUMBER, { value: 2 });
    const minus = s.add(SHADER_NUMBER, { value: -0.25 });
    const broken = s.add(SHADER_NUMBER, { value: Number.NaN });
    const sum = s.add('ShaderAdd');
    const times = s.add('ShaderMultiply');
    const diff = s.add('ShaderSubtract');
    const output = s.add(SHADER_OUTPUT);
    s.link(half, 'value', sum, 'a');
    s.link(two, 'value', sum, 'b');
    s.link(sum, 'result', times, 'a');
    s.link(minus, 'value', times, 'b');
    s.link(times, 'result', diff, 'a');
    s.link(broken, 'value', diff, 'b');
    s.link(diff, 'result', output, 'color');
    const { fragment } = compileShader(s.asset().graph).compiled!;
    // Locals are numbered in visiting order; the literal's form is what matters here.
    expect(fragment).toMatch(/float n\d+_value = 0\.5;/);
    expect(fragment).toMatch(/float n\d+_value = 2\.0;/);
    expect(fragment).toMatch(/float n\d+_value = -0\.25;/);
    expect(fragment).toMatch(/float n\d+_value = 0\.0;/);
    expect(fragment).toMatch(/float n\d+_result = \(n\d+_value \+ n\d+_value\);/);
    expect(fragment).toMatch(/clamp\(vec3\(n\d+_result\), 0\.0, 1\.0\)/);
    expect(fragment).not.toContain('uniform float u_');
  });

  it('keeps the texture’s alpha until the Output’s alpha is fed, then clamps what feeds it', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    const alpha = s.add(SHADER_TEXTURE_ALPHA);
    const opacity = s.add(SHADER_PROPERTY, { name: 'opacity', default: 1, min: 0, max: 1 });
    const times = s.add('ShaderMultiply');
    const output = s.add(SHADER_OUTPUT);
    // Nothing to type on the alpha port: a typed-in 0 would hide every sprite by default.
    expect(inlineInputs(ShaderOutput).map((p) => p.key)).toEqual([]);
    let { fragment } = compileShader(s.asset().graph).compiled!;
    expect(fragment).toContain('float alpha = clamp(texel.a, 0.0, 1.0);');
    expect(fragment).toContain('finalColor = vec4(clamp(rgb, 0.0, 1.0) * alpha, alpha);');

    s.link(alpha, 'value', times, 'a');
    s.link(opacity, 'value', times, 'b');
    s.link(times, 'result', output, 'alpha');
    ({ fragment } = compileShader(s.asset().graph).compiled!);
    expect(fragment).toMatch(/float n0_value = texel\.a;/);
    expect(fragment).toMatch(/float n1_value = u_opacity;/);
    expect(fragment).toMatch(/float n2_result = \(n0_value \* n1_value\);/);
    expect(fragment).toContain('float alpha = clamp(n2_result, 0.0, 1.0);');
    expect(fragment).toContain('finalColor = vec4(clamp(rgb, 0.0, 1.0) * alpha, alpha);');

    // A colour reaching the alpha port through an `any` result is the Output's problem.
    const wrong = shaderEditor(f, 'Wrong');
    const texel = wrong.add(SHADER_TEXTURE_COLOR);
    const doubled = wrong.add('ShaderMultiply', { b: 2 });
    const out = wrong.add(SHADER_OUTPUT);
    wrong.link(texel, 'color', doubled, 'a');
    wrong.link(doubled, 'result', out, 'alpha');
    const build = compileShader(wrong.asset().graph);
    expect(build.compiled).toBeNull();
    expect(build.problems).toEqual({ [out]: '"alpha" needs a number, not a colour' });
  });

  it('names the problem on the node, and compiles nothing while there is one', () => {
    const f = makeProject();
    const s = shaderEditor(f);
    // A colour cannot be wired to a number port; a Multiply's `any` result can, and carries one.
    const texel = s.add(SHADER_TEXTURE_COLOR);
    const doubled = s.add('ShaderMultiply', { b: 2 });
    const mixed = s.add('ShaderMix');
    const output = s.add(SHADER_OUTPUT);
    s.link(texel, 'color', doubled, 'a');
    s.link(doubled, 'result', mixed, 't');
    s.link(mixed, 'result', output, 'color');
    let build = compileShader(s.asset().graph);
    expect(build.compiled).toBeNull();
    expect(build.problems).toEqual({ [mixed]: '"t" needs a number, not a colour' });
    expect(build.reason).toBe('"t" needs a number, not a colour');

    const dup = shaderEditor(f, 'Dup');
    dup.add(SHADER_OUTPUT);
    const first = dup.add(SHADER_PROPERTY, property('amount'));
    const second = dup.add(SHADER_PROPERTY, property('amount'));
    const odd = dup.add(SHADER_PROPERTY, property('2fast'));
    build = compileShader(dup.asset().graph);
    expect(build.compiled).toBeNull();
    expect(build.problems).toEqual({
      [second]: 'Another property is named "amount"',
      [odd]: 'Name must be one word: letters, digits, _',
    });
    expect(build.problems[first]).toBeUndefined();

    // A node cannot be wired to itself, but two can be wired into a ring.
    const loop = shaderEditor(f, 'Loop');
    const one = loop.add('ShaderAdd');
    const two = loop.add('ShaderAdd');
    const out = loop.add(SHADER_OUTPUT);
    loop.link(one, 'result', two, 'a');
    loop.link(two, 'result', one, 'b');
    loop.link(two, 'result', out, 'color');
    build = compileShader(loop.asset().graph);
    expect(build.compiled).toBeNull();
    expect(build.problems).toEqual({ [two]: 'Feeds itself' });
  });
});

describe('a sprite wearing a shader', () => {
  function wearing() {
    const f = makeSpriteScene();
    const s = shaderEditor(f);
    buildAdjustColors(s);
    const sprite = () => findAssetOf(f.project, f.spriteId, 'sprite')!;
    return { ...f, s, sprite };
  }

  it('has the shader picker and, once one is picked, that shader’s rows', () => {
    const f = wearing();
    const keys = () => schemaForAsset(f.project, f.sprite()).map((d) => d.key);
    expect(keys()).toEqual(['enabled', 'width', 'height', 'color', 'texture', 'shader']);
    expect(schemaFor('sprite').find((d) => d.key === SHADER_KEY)).toMatchObject({
      kind: 'asset',
      assetTypes: ['shader'],
      default: '',
    });
    f.stack.execute(new SetAssetPropertyCommand(f.spriteId, SHADER_KEY, f.s.id));
    expect(keys()).toEqual([
      'enabled',
      'width',
      'height',
      'color',
      'texture',
      'shader',
      'brightness',
      'contrast',
      'saturation',
    ]);
    expect(assetDeviates(f.project, f.sprite())).toBe(true);
    // A dangling id adds no rows and draws nothing – the picker still shows it.
    f.sprite().props[SHADER_KEY] = 'gone';
    expect(keys()).toHaveLength(6);
  });

  it('gives a placed object the rows of the shader it overrides, filled with the node defaults', () => {
    const f = wearing();
    const id = place(f, f.spriteId, { x: 0, y: 0 });
    const object = () => f.object(id);
    expect(schemaForObject(f.project, f.sprite(), object())).toHaveLength(6);
    f.stack.execute(new SetOverrideCommand(f.sceneId, id, SHADER_KEY, f.s.id));
    expect(schemaForObject(f.project, f.sprite(), object()).map((d) => d.key)).toContain('saturation');
    expect(resolveSceneObjectProps(f.project, f.sprite(), object())).toMatchObject({
      shader: f.s.id,
      brightness: 1,
      contrast: 1,
      saturation: 1,
    });
    // A parameter override passes through its row: clamped to the node's range.
    f.stack.execute(new SetOverrideCommand(f.sceneId, id, 'brightness', 9));
    expect(object().overrides.brightness).toBe(4);
    expect(objectDeviates(f.project, f.sceneOf(), object())).toBe(true);
    // Back to no shader the parameter row is gone, so the stored value marks nothing.
    f.stack.execute(new SetOverrideCommand(f.sceneId, id, SHADER_KEY, ''));
    expect(resolveSceneObjectProps(f.project, f.sprite(), object())).not.toHaveProperty('brightness');
    expect(objectDeviates(f.project, f.sceneOf(), object())).toBe(false);
  });

  it('rejects a parameter no shader of the sprite declares', () => {
    const f = wearing();
    const id = place(f, f.spriteId, { x: 0, y: 0 });
    expect(() => f.stack.execute(new SetOverrideCommand(f.sceneId, id, 'brightness', 0.5))).toThrow(
      /not defined/,
    );
  });

  it('survives a round trip: the graph, the sprite’s values and an override-only shader’s values', () => {
    const f = wearing();
    const other = shaderEditor(f, 'Fade');
    other.add(SHADER_OUTPUT);
    other.add(SHADER_PROPERTY, { name: 'amount', default: 0.5, min: 0, max: 1 });
    f.stack.execute(new SetAssetPropertyCommand(f.spriteId, SHADER_KEY, f.s.id));
    f.stack.execute(new SetAssetPropertyCommand(f.spriteId, 'contrast', 2));
    const id = place(f, f.spriteId, { x: 0, y: 0 });
    f.stack.execute(new SetOverrideCommand(f.sceneId, id, SHADER_KEY, other.id));
    f.stack.execute(new SetOverrideCommand(f.sceneId, id, 'amount', 0.25));

    const loaded = roundTrip(f.project);
    const shader = findAssetOf(loaded, f.s.id, 'shader')!;
    expect(shader.graph).toEqual(f.s.asset().graph);
    expect(compileShader(shader.graph).compiled!.fragment).toBe(
      compileShader(f.s.asset().graph).compiled!.fragment,
    );
    const sprite = findAssetOf(loaded, f.spriteId, 'sprite')!;
    expect(sprite.props).toMatchObject({ shader: f.s.id, contrast: 2, brightness: 1 });
    const placed = findAssetOf(loaded, f.sceneId, 'scene')!.objects.find((o) => o.id === id)!;
    expect(placed.overrides).toEqual({ shader: other.id, amount: 0.25 });
  });

  it('drops anything but shader nodes from a shader graph on load', () => {
    const graph: Graph = {
      nodes: [
        { id: 'o', type: SHADER_OUTPUT, x: 0, y: 0, params: {} },
        { id: 'e', type: 'OnStart', x: 0, y: 0, params: {} },
      ],
      edges: [],
    };
    const loaded = deserializeProject(
      { id: 'p', name: 'P', assets: [{ id: 'sh', type: 'shader', name: 'S', props: {}, graph }] },
      'p',
    );
    const shader = loaded.assets[0] as ShaderAsset;
    expect(shader.graph.nodes.map((n) => n.id)).toEqual(['o']);
  });
});