image-studio/renderer/brush-preview.ts

import type { ImageDocument } from "@core/document";
import type { Vec2D } from "@core/geometry";
import type { Layer } from "@core/layer";
import { getLayerMask } from "@core/layer-mask-utils";
import type { EditorState } from "@editor/state";
import type { RgbaColor, WebGlRendererContext } from "./types";

export type BrushPreviewRenderer = {
  render(document: ImageDocument, editor: EditorState): void;
  dispose(): void;
};

const previewHaloColor: RgbaColor = [0, 0, 0, 0.55];
const previewColor: RgbaColor = [1, 1, 1, 0.92];
const previewFillOpacity = 0.08;

export function createBrushPreviewRenderer(context: WebGlRendererContext): BrushPreviewRenderer {
  const { gl } = context;
  const program = createProgram(gl);
  const positionLocation = gl.getAttribLocation(program, "a_position");
  const canvasSizeLocation = gl.getUniformLocation(program, "u_canvasSize");
  const centerLocation = gl.getUniformLocation(program, "u_center");
  const radiusLocation = gl.getUniformLocation(program, "u_radius");
  const hardnessLocation = gl.getUniformLocation(program, "u_hardness");
  const colorLocation = gl.getUniformLocation(program, "u_color");
  const ringWidthLocation = gl.getUniformLocation(program, "u_ringWidth");
  const fillOpacityLocation = gl.getUniformLocation(program, "u_fillOpacity");
  const positionBuffer = gl.createBuffer();

  if (!canvasSizeLocation || !centerLocation || !radiusLocation || !hardnessLocation || !colorLocation || !ringWidthLocation || !fillOpacityLocation || !positionBuffer) {
    throw new Error("Failed to create brush preview renderer");
  }

  return {
    render(document, editor) {
      const preview = resolveBrushPreview(document, editor, context.canvas);
      if (!preview) return;

      gl.disable(gl.SCISSOR_TEST);
      gl.enable(gl.BLEND);
      gl.blendFunc(gl.ONE, gl.ONE_MINUS_SRC_ALPHA);
      gl.useProgram(program);

      gl.uniform2f(canvasSizeLocation, context.canvas.width, context.canvas.height);
      gl.uniform2f(centerLocation, preview.center.x, preview.center.y);
      gl.uniform2f(radiusLocation, preview.radius.x, preview.radius.y);
      gl.uniform1f(hardnessLocation, preview.hardness);

      gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer);
      gl.bufferData(gl.ARRAY_BUFFER, previewVertices(preview.center, preview.radius), gl.DYNAMIC_DRAW);
      gl.enableVertexAttribArray(positionLocation);
      gl.vertexAttribPointer(positionLocation, 2, gl.FLOAT, false, 0, 0);

      drawPreviewPass(gl, colorLocation, ringWidthLocation, fillOpacityLocation, previewHaloColor, 3, 0);
      drawPreviewPass(gl, colorLocation, ringWidthLocation, fillOpacityLocation, previewColor, 1.35, previewFillOpacity);

      gl.disable(gl.BLEND);
    },
    dispose() {
      gl.deleteBuffer(positionBuffer);
      gl.deleteProgram(program);
    },
  };
}

function drawPreviewPass(
  gl: WebGL2RenderingContext,
  colorLocation: WebGLUniformLocation,
  ringWidthLocation: WebGLUniformLocation,
  fillOpacityLocation: WebGLUniformLocation,
  color: RgbaColor,
  ringWidth: number,
  fillOpacity: number,
) {
  gl.uniform4fv(colorLocation, color);
  gl.uniform1f(ringWidthLocation, ringWidth);
  gl.uniform1f(fillOpacityLocation, fillOpacity);
  gl.drawArrays(gl.TRIANGLES, 0, 6);
}

function resolveBrushPreview(document: ImageDocument, editor: EditorState, canvas: HTMLCanvasElement) {
  if (!editor.brushPreview || (editor.tools.activeTool !== "brush" && editor.tools.activeTool !== "eraser" && editor.tools.activeTool !== "feather")) return undefined;
  if (editor.tools.interactionMode.type === "temporary-pan" || (editor.tools.interactionMode.type === "tool" && editor.tools.interactionMode.tool === "pan")) return undefined;

  const layer = resolveBrushTargetLayer(document, editor);
  if (!layer) return undefined;

  const size = Math.max(1, editor.tools.activeTool === "feather" ? editor.tools.feather.size : editor.tools.brush.size);
  const zoom = editor.viewport.zoom;
  return {
    center: documentPointToScreenPoint(canvas, editor.brushPreview.position, editor),
    radius: {
      x: Math.max(1, Math.abs(layer.transform.scale.x) * size * zoom * 0.5),
      y: Math.max(1, Math.abs(layer.transform.scale.y) * size * zoom * 0.5),
    },
    hardness: editor.tools.activeTool === "feather" ? 0.72 : Math.max(0, Math.min(1, editor.tools.brush.hardness / 100)),
  };
}

function resolveBrushTargetLayer(document: ImageDocument, editor: EditorState): Extract<Layer, { type: "image" | "raster" }> | undefined {
  const editingMask = Boolean(editor.maskEdit);
  const layerId = editor.maskEdit?.kind === "inpaintRegion" ? editor.maskEdit.targetLayerId : editor.maskEdit?.maskLayerId ?? editor.selection.layerIds[0];
  if (!layerId) return undefined;

  const layers = document.artboards.flatMap((artboard) => artboard.layers);
  const selectedLayer = findPaintableLayer(layers, layerId);
  const maskLayerId = editor.tools.activeTool === "feather" && !editingMask && selectedLayer ? getLayerMask(selectedLayer)?.maskLayerId : undefined;
  const layer = maskLayerId ? findPaintableLayer(layers, maskLayerId) : selectedLayer;
  if (!layer || layer.locked || (!editingMask && !layer.visible)) return undefined;
  return layer;
}

function documentPointToScreenPoint(canvas: HTMLCanvasElement, point: Vec2D, editor: EditorState): Vec2D {
  return {
    x: canvas.width / 2 + (point.x - editor.viewport.center.x) * editor.viewport.zoom,
    y: canvas.height / 2 + (point.y - editor.viewport.center.y) * editor.viewport.zoom,
  };
}

function previewVertices(center: Vec2D, radius: Vec2D) {
  const padding = 5;
  const x1 = center.x - radius.x - padding;
  const x2 = center.x + radius.x + padding;
  const y1 = center.y - radius.y - padding;
  const y2 = center.y + radius.y + padding;

  return new Float32Array([x1, y1, x2, y1, x1, y2, x1, y2, x2, y1, x2, y2]);
}

function findPaintableLayer(layers: readonly Layer[], layerId: string): Extract<Layer, { type: "image" | "raster" }> | undefined {
  for (const layer of layers) {
    if (layer.id === layerId && (layer.type === "image" || layer.type === "raster")) return layer;
    if (layer.type === "group") {
      const child = findPaintableLayer(layer.children, layerId);
      if (child) return child;
    }
  }
  return undefined;
}

function createProgram(gl: WebGL2RenderingContext) {
  const vertexShader = compileShader(
    gl,
    gl.VERTEX_SHADER,
    `#version 300 es
    in vec2 a_position;
    uniform vec2 u_canvasSize;
    out vec2 v_position;
    void main() {
      vec2 clip = vec2((a_position.x / u_canvasSize.x) * 2.0 - 1.0, 1.0 - (a_position.y / u_canvasSize.y) * 2.0);
      gl_Position = vec4(clip, 0.0, 1.0);
      v_position = a_position;
    }`,
  );
  const fragmentShader = compileShader(
    gl,
    gl.FRAGMENT_SHADER,
    `#version 300 es
    precision mediump float;
    uniform vec2 u_center;
    uniform vec2 u_radius;
    uniform float u_hardness;
    uniform vec4 u_color;
    uniform float u_ringWidth;
    uniform float u_fillOpacity;
    in vec2 v_position;
    out vec4 outColor;
    void main() {
      vec2 radius = max(u_radius, vec2(1.0));
      float minimumRadius = max(1.0, min(radius.x, radius.y));
      float normalizedDistance = length((v_position - u_center) / radius);
      float outerDistance = abs(normalizedDistance - 1.0) * minimumRadius;
      float outerRing = 1.0 - smoothstep(max(0.0, u_ringWidth - 1.0), u_ringWidth + 1.0, outerDistance);

      float innerRing = 0.0;
      if (u_hardness > 0.05 && u_hardness < 0.98) {
        float innerDistance = abs(normalizedDistance - u_hardness) * minimumRadius;
        innerRing = 0.45 * (1.0 - smoothstep(max(0.0, u_ringWidth - 1.0), u_ringWidth + 1.0, innerDistance));
      }

      float fillStart = min(u_hardness, 0.98);
      float fill = normalizedDistance <= 1.0 ? 1.0 - smoothstep(fillStart, 1.0, normalizedDistance) : 0.0;
      float alpha = max(max(outerRing, innerRing) * u_color.a, fill * u_fillOpacity);
      if (alpha <= 0.001) discard;
      outColor = vec4(u_color.rgb * alpha, alpha);
    }`,
  );
  const program = gl.createProgram();
  if (!program) throw new Error("Failed to create brush preview shader program");

  gl.attachShader(program, vertexShader);
  gl.attachShader(program, fragmentShader);
  gl.linkProgram(program);
  gl.deleteShader(vertexShader);
  gl.deleteShader(fragmentShader);

  if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
    const message = gl.getProgramInfoLog(program) ?? "Unknown brush preview program link error";
    gl.deleteProgram(program);
    throw new Error(message);
  }

  return program;
}

function compileShader(gl: WebGL2RenderingContext, type: number, source: string) {
  const shader = gl.createShader(type);
  if (!shader) throw new Error("Failed to create shader");

  gl.shaderSource(shader, source);
  gl.compileShader(shader);
  if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
    const message = gl.getShaderInfoLog(shader) ?? "Unknown shader compile error";
    gl.deleteShader(shader);
    throw new Error(message);
  }

  return shader;
}