fleeting-garden/src/pipelines/eraser/eraser-texture-pipeline.ts

import { vec2 } from 'gl-matrix';

import { appConfig } from '../../config';
import {
  createCachedFloat32BufferWrite,
  writeFloat32BufferIfChanged,
} from '../../utils/graphics/cached-buffer-write';
import { smartCompile } from '../../utils/graphics/smart-compile';
import { CommonState } from '../common-state/common-state';
import shader from './eraser-texture.wgsl?raw';

interface LineSegment {
  from: vec2;
  to: vec2;
}

export class EraserTexturePipeline {
  private static readonly UNIFORM_COUNT = 4;
  private static readonly MAX_LINE_COUNT = appConfig.pipelines.eraser.maxTextureLineCount;
  private static readonly VERTICES_PER_LINE_SEGMENT = 6;
  private static readonly ATTRIBUTES_PER_LINE_SEGMENT = 6;

  private readonly bindGroupLayout: GPUBindGroupLayout;
  private readonly bindGroup: GPUBindGroup;
  private readonly pipeline: GPURenderPipeline;
  private readonly multiTargetPipeline: GPURenderPipeline;
  private readonly uniforms: GPUBuffer;
  private readonly uniformValues = new Float32Array(EraserTexturePipeline.UNIFORM_COUNT);
  private readonly uniformCache = createCachedFloat32BufferWrite(
    EraserTexturePipeline.UNIFORM_COUNT
  );
  private readonly vertexBuffer: GPUBuffer;
  private readonly vertexUploadData = new Float32Array(
    EraserTexturePipeline.MAX_LINE_COUNT *
      EraserTexturePipeline.VERTICES_PER_LINE_SEGMENT *
      EraserTexturePipeline.ATTRIBUTES_PER_LINE_SEGMENT
  );

  private lineSegments: Array<LineSegment> = [];
  private actualSegments: Array<LineSegment> = [];

  public constructor(
    private readonly device: GPUDevice,
    private readonly commonState: CommonState
  ) {
    this.bindGroupLayout = device.createBindGroupLayout({
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.FRAGMENT,
          buffer: {
            type: 'uniform',
          },
        },
      ],
    });

    this.vertexBuffer = device.createBuffer({
      size:
        EraserTexturePipeline.MAX_LINE_COUNT *
        EraserTexturePipeline.VERTICES_PER_LINE_SEGMENT *
        EraserTexturePipeline.ATTRIBUTES_PER_LINE_SEGMENT *
        Float32Array.BYTES_PER_ELEMENT,
      usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
    });

    const shaderModule = smartCompile(device, CommonState.shaderCode, shader);
    this.pipeline = this.createPipeline(shaderModule, 'fragment', 1);
    this.multiTargetPipeline = this.createPipeline(shaderModule, 'fragmentMrt', 3);

    this.uniforms = this.device.createBuffer({
      size: EraserTexturePipeline.UNIFORM_COUNT * Float32Array.BYTES_PER_ELEMENT,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });

    this.bindGroup = this.device.createBindGroup({
      layout: this.bindGroupLayout,
      entries: [
        {
          binding: 0,
          resource: {
            buffer: this.uniforms,
          },
        },
      ],
    });
  }

  public addSwipeSegment(from: vec2, to: vec2): void {
    this.lineSegments.push({
      from: vec2.clone(from),
      to: vec2.clone(to),
    });
  }

  public clearSwipes(): void {
    this.lineSegments.length = 0;
    this.actualSegments.length = 0;
  }

  public setParameters({ eraserSize }: { eraserSize: number }): void {
    const eraserRadius = eraserSize / 2;

    this.uniformValues[0] = eraserRadius * eraserRadius;
    this.uniformValues[1] = 0;
    this.uniformValues[2] = 0;
    this.uniformValues[3] = 0;
    writeFloat32BufferIfChanged(
      this.device,
      this.uniforms,
      this.uniformValues,
      this.uniformCache
    );

    this.actualSegments = this.lineSegments.slice();
    this.lineSegments.length = 0;

    if (this.actualSegments.length === 0) {
      return;
    }

    if (this.actualSegments.length > EraserTexturePipeline.MAX_LINE_COUNT) {
      this.actualSegments = EraserTexturePipeline.subsampleSegments(this.actualSegments);
    }

    const lineCount = this.lineCount;
    let floatOffset = 0;
    for (let i = 0; i < lineCount; i++) {
      const segment = this.actualSegments[i];
      floatOffset = this.writeSegmentVertices(
        this.vertexUploadData,
        floatOffset,
        segment.from,
        segment.to,
        eraserRadius
      );
    }

    this.device.queue.writeBuffer(
      this.vertexBuffer,
      0,
      this.vertexUploadData,
      0,
      floatOffset
    );
  }

  public execute(commandEncoder: GPUCommandEncoder, textureOut: GPUTextureView): void {
    this.executeWithPipeline(commandEncoder, this.pipeline, [textureOut]);
  }

  public executeMultiTarget(
    commandEncoder: GPUCommandEncoder,
    sourceMapOut: GPUTextureView,
    influenceMapOut: GPUTextureView,
    trailMapOut: GPUTextureView
  ): void {
    this.executeWithPipeline(commandEncoder, this.multiTargetPipeline, [
      sourceMapOut,
      influenceMapOut,
      trailMapOut,
    ]);
  }

  private executeWithPipeline(
    commandEncoder: GPUCommandEncoder,
    pipeline: GPURenderPipeline,
    textureViews: Array<GPUTextureView>
  ): void {
    if (this.lineCount === 0) {
      return;
    }

    const renderPassDescriptor: GPURenderPassDescriptor = {
      colorAttachments: textureViews.map<GPURenderPassColorAttachment>((view) => ({
        view,
        loadOp: 'load',
        storeOp: 'store',
      })),
    };

    const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
    passEncoder.setPipeline(pipeline);
    this.commonState.execute(passEncoder);
    passEncoder.setBindGroup(1, this.bindGroup);
    passEncoder.setVertexBuffer(0, this.vertexBuffer);
    passEncoder.draw(EraserTexturePipeline.VERTICES_PER_LINE_SEGMENT * this.lineCount, 1);
    passEncoder.end();
  }

  public destroy(): void {
    this.vertexBuffer.destroy();
    this.uniforms.destroy();
  }

  private createPipeline(
    shaderModule: GPUShaderModule,
    fragmentEntryPoint: string,
    colorTargetCount: number
  ): GPURenderPipeline {
    return this.device.createRenderPipeline({
      layout: this.device.createPipelineLayout({
        bindGroupLayouts: [this.commonState.bindGroupLayout, this.bindGroupLayout],
      }),
      vertex: {
        module: shaderModule,
        entryPoint: 'vertex',
        buffers: [
          {
            arrayStride: Float32Array.BYTES_PER_ELEMENT * 6,
            attributes: [
              {
                shaderLocation: 0,
                format: 'float32x2',
                offset: 0,
              },
              {
                shaderLocation: 1,
                format: 'float32x2',
                offset: Float32Array.BYTES_PER_ELEMENT * 2,
              },
              {
                shaderLocation: 2,
                format: 'float32x2',
                offset: Float32Array.BYTES_PER_ELEMENT * 4,
              },
            ],
          },
        ],
      },
      fragment: {
        module: shaderModule,
        entryPoint: fragmentEntryPoint,
        targets: Array.from({ length: colorTargetCount }, () => ({
          format: 'rgba16float' as const,
        })),
      },
      primitive: {
        topology: 'triangle-list',
      },
    });
  }

  private static subsampleSegments(segments: Array<LineSegment>): Array<LineSegment> {
    if (segments.length <= EraserTexturePipeline.MAX_LINE_COUNT) {
      return segments;
    }

    const result: Array<LineSegment> = [];
    for (let i = 0; i < EraserTexturePipeline.MAX_LINE_COUNT; i++) {
      const index = Math.round(
        (i * (segments.length - 1)) / (EraserTexturePipeline.MAX_LINE_COUNT - 1)
      );
      result.push(segments[index]);
    }

    return result;
  }

  private writeSegmentVertices(
    target: Float32Array,
    offset: number,
    from: vec2,
    to: vec2,
    width: number
  ): number {
    const dx = to[0] - from[0];
    const dy = to[1] - from[1];
    const length = Math.hypot(dx, dy);
    const directionX = length > 0 ? dx / length : 1;
    const directionY = length > 0 ? dy / length : 0;
    const scaledDirectionX = directionX * width;
    const scaledDirectionY = directionY * width;
    const perpendicularX = directionY * width;
    const perpendicularY = -directionX * width;

    const startX = from[0] - scaledDirectionX;
    const startY = from[1] - scaledDirectionY;
    const endX = to[0] + scaledDirectionX;
    const endY = to[1] + scaledDirectionY;

    offset = this.writeVertex(
      target,
      offset,
      startX + perpendicularX,
      startY + perpendicularY,
      from,
      to
    );
    offset = this.writeVertex(
      target,
      offset,
      startX - perpendicularX,
      startY - perpendicularY,
      from,
      to
    );
    offset = this.writeVertex(
      target,
      offset,
      endX + perpendicularX,
      endY + perpendicularY,
      from,
      to
    );
    offset = this.writeVertex(
      target,
      offset,
      startX - perpendicularX,
      startY - perpendicularY,
      from,
      to
    );
    offset = this.writeVertex(
      target,
      offset,
      endX + perpendicularX,
      endY + perpendicularY,
      from,
      to
    );
    return this.writeVertex(
      target,
      offset,
      endX - perpendicularX,
      endY - perpendicularY,
      from,
      to
    );
  }

  private writeVertex(
    target: Float32Array,
    offset: number,
    screenX: number,
    screenY: number,
    from: vec2,
    to: vec2
  ): number {
    target[offset++] = screenX;
    target[offset++] = screenY;
    target[offset++] = from[0];
    target[offset++] = from[1];
    target[offset++] = to[0];
    target[offset++] = to[1];
    return offset;
  }

  private get lineCount(): number {
    return this.actualSegments.length;
  }
}