image-studio/renderer/layers.ts

import type { Asset } from "@core/asset";
import type { ImageDocument } from "@core/document";
import type { Rect } from "@core/geometry";
import type { Layer } from "@core/layer";
import { getLayerMask } from "@core/layer-mask-utils";
import { createDocumentReadIndex, resolveIndexedLayerBounds, type DocumentReadIndex } from "@editor/document-indexes";
import type { EditorState, GenerationCandidate, MaskViewMode, ViewportState } from "@editor/state";
import { clearScreenRect } from "./clear-rect";
import type { ImageTextureRenderer } from "./image-textures";
import { documentRectToScreenRect } from "./screen-rect";
import type { RgbaColor, ScreenRect, WebGlRendererContext } from "./types";
import type { AdjustmentPass } from "./adjustment-pass";
import { rasterizedTextLayerRenderAsset } from "./text-asset";

const imageLayerColor: RgbaColor = [0.38, 0.42, 0.5, 1];
const imageLayerInsetColor: RgbaColor = [0.48, 0.54, 0.64, 1];
const hiddenMaskOverlayColor: RgbaColor = [1, 0.08, 0.08, 0.45];
const inpaintRegionOverlayColor: RgbaColor = [1, 0.12, 0.18, 0.48];
const comparisonDividerColor: RgbaColor = [1, 1, 1, 0.9];
const maskRevealPreviewOpacity = 0.28;

export function renderLayers(context: WebGlRendererContext, document: ImageDocument, editor: EditorState, imageTextureRenderer: ImageTextureRenderer, adjustmentPass: AdjustmentPass) {
  const documentIndex = createDocumentReadIndex(document);
  const generationCandidate = selectedGenerationCandidate(editor);

  for (const artboard of document.artboards) {
    if (!artboard.visible) continue;
    const clipRect = documentRectToScreenRect(context.canvas, artboard.bounds, editor.viewport);
    const maskLayerIds = documentIndex.maskLayerIdsByArtboardId.get(artboard.id) ?? emptyLayerIds;
    renderLayerTree(context, documentIndex, editor, artboard.layers, imageTextureRenderer, adjustmentPass, clipRect, maskLayerIds);
    if (generationCandidate?.placement.artboardId === artboard.id) renderGenerationCandidatePreview(context, editor, generationCandidate, imageTextureRenderer, clipRect);
  }
}

export function generationCandidatePreviewAssets(editor: EditorState): Asset[] {
  return editor.generation.candidates.map((candidate) => generationCandidateAsset(candidate));
}

function renderLayerTree(
  context: WebGlRendererContext,
  documentIndex: DocumentReadIndex,
  editor: EditorState,
  layers: readonly Layer[],
  imageTextureRenderer: ImageTextureRenderer,
  adjustmentPass: AdjustmentPass,
  clipRect: ScreenRect,
  maskLayerIds: ReadonlySet<string>,
) {
  const stack: Array<{ layer: Layer; clipRect: ScreenRect; inheritedOpacity: number }> = [];
  for (const layer of layers) stack.push({ layer, clipRect, inheritedOpacity: 1 });
  while (stack.length > 0) {
    const frame = stack.pop();
    if (!frame) continue;
    const { layer } = frame;
    if (!layer.visible || maskLayerIds.has(layer.id)) continue;
    const effectiveOpacity = frame.inheritedOpacity * layer.opacity;
    if (effectiveOpacity <= 0) continue;
    const effectiveClipRect = resolveLayerClipRect(context, documentIndex, editor.viewport, layer, frame.clipRect);
    if (!effectiveClipRect) continue;
    if (layer.type === "group") {
      for (const child of layer.children) stack.push({ layer: child, clipRect: effectiveClipRect, inheritedOpacity: effectiveOpacity });
      continue;
    }
    if (layer.type === "adjustment") {
      adjustmentPass.apply(layer.adjustment, effectiveOpacity, effectiveClipRect);
      continue;
    }
    if (layer.type === "text") {
      renderTextLayer(context, documentIndex, editor, layer, imageTextureRenderer, effectiveClipRect, effectiveOpacity);
    } else {
      renderLeafLayer(context, documentIndex, editor, layer, imageTextureRenderer, effectiveClipRect, effectiveOpacity);
    }
  }
}

function renderTextLayer(context: WebGlRendererContext, documentIndex: DocumentReadIndex, editor: EditorState, layer: Extract<Layer, { type: "text" }>, imageTextureRenderer: ImageTextureRenderer, clipRect: ScreenRect, opacity: number) {
  if (editor.maskEdit?.viewMode && isIsolatedMaskView(editor.maskEdit.viewMode)) return;
  const bounds = resolveIndexedLayerBounds(documentIndex, layer);
  if (!bounds) return;
  const rect = documentRectToScreenRect(context.canvas, bounds, editor.viewport);
  imageTextureRenderer.render(rasterizedTextLayerRenderAsset(layer), rect, clipRect, opacity, layer.transform.rotation);
}

function renderLeafLayer(
  context: WebGlRendererContext,
  documentIndex: DocumentReadIndex,
  editor: EditorState,
  layer: Extract<Layer, { type: "image" | "raster" }>,
  imageTextureRenderer: ImageTextureRenderer,
  effectiveClipRect: ScreenRect,
  effectiveOpacity: number,
) {
  const editingMaskLayer = editor.maskEdit?.maskLayerId === layer.id;
  const maskViewMode = editor.maskEdit?.viewMode ?? "composite";
  const isolatedMaskView = isIsolatedMaskView(maskViewMode);
  if (isolatedMaskView && layer.id !== editor.maskEdit?.targetLayerId) return;

  const bounds = resolveIndexedLayerBounds(documentIndex, layer);
  if (!bounds) return;

  const rect = documentRectToScreenRect(context.canvas, bounds, editor.viewport);
  const asset = assetWithBrushStrokePreview(documentIndex.assetById.get(layer.assetId), editor);
  const activeInpaintRegion = editor.maskEdit?.kind === "inpaintRegion" && editor.maskEdit.targetLayerId === layer.id
    ? documentIndex.assetById.get(editor.maskEdit.maskAssetId)
    : undefined;
  if (asset && activeInpaintRegion) {
    const inpaintMaskAsset = assetWithBrushStrokePreview(activeInpaintRegion, editor);
    const inpaintMaskRect = documentRectToScreenRect(context.canvas, {
      x: layer.transform.position.x,
      y: layer.transform.position.y,
      w: inpaintMaskAsset.intrinsicSize.w * layer.transform.scale.x,
      h: inpaintMaskAsset.intrinsicSize.h * layer.transform.scale.y,
    }, editor.viewport);
    if (maskViewMode === "blackWhite" && imageTextureRenderer.renderMaskVisualization(inpaintMaskAsset, inpaintMaskRect, "blackWhite", undefined, effectiveClipRect)) return;
    if (maskViewMode === "alpha" && imageTextureRenderer.renderMaskVisualization(inpaintMaskAsset, inpaintMaskRect, "alpha", undefined, effectiveClipRect)) return;
    if (imageTextureRenderer.render(asset, rect, effectiveClipRect, effectiveOpacity, layer.transform.rotation, layer.sourceRect)) {
      imageTextureRenderer.renderMaskVisualization(inpaintMaskAsset, inpaintMaskRect, "activeOverlay", inpaintRegionOverlayColor, effectiveClipRect);
      return;
    }
  }
  const layerMask = getLayerMask(layer);
  const maskLayer = !editingMaskLayer && layerMask?.enabled ? documentIndex.layerById.get(layerMask.maskLayerId) : undefined;
  const pendingMaskPreview = editor.brushStrokePreview?.pendingTargetLayerId === layer.id && editor.brushStrokePreview.intrinsicSize
    ? {
        id: editor.brushStrokePreview.assetId,
        name: `${layer.name} Mask Preview`,
        mimeType: "image/png",
        source: editor.brushStrokePreview.source,
        intrinsicSize: editor.brushStrokePreview.intrinsicSize,
      }
    : undefined;
  const maskAsset = assetWithBrushStrokePreview(maskLayer && (maskLayer.type === "image" || maskLayer.type === "raster") ? documentIndex.assetById.get(maskLayer.assetId) : undefined, editor) ?? pendingMaskPreview;
  const maskBounds = maskLayer ? resolveIndexedLayerBounds(documentIndex, maskLayer) : undefined;
  const maskRect = maskBounds ? documentRectToScreenRect(context.canvas, maskBounds, editor.viewport) : pendingMaskPreview ? rect : undefined;
  const activeMaskTarget = Boolean((editor.maskEdit?.targetLayerId === layer.id && editor.maskEdit.maskLayerId === layerMask?.maskLayerId) || pendingMaskPreview);
  const showMaskRevealPreview = editor.tools.activeTool === "brush" && activeMaskTarget && maskViewMode === "composite";

  if (asset && maskAsset && maskRect && activeMaskTarget) {
    if (maskViewMode === "blackWhite" && imageTextureRenderer.renderMaskVisualization(maskAsset, maskRect, "blackWhite", undefined, effectiveClipRect)) return;
    if (maskViewMode === "alpha" && imageTextureRenderer.renderMaskVisualization(maskAsset, maskRect, "alpha", undefined, effectiveClipRect)) return;
    if (maskViewMode === "overlay" && imageTextureRenderer.render(asset, rect, effectiveClipRect, effectiveOpacity, layer.transform.rotation, layer.sourceRect)) {
      const contentClipRect = intersectScreenRects(effectiveClipRect, rect);
      if (contentClipRect) imageTextureRenderer.renderMaskVisualization(maskAsset, maskRect, "hiddenOverlay", hiddenMaskOverlayColor, contentClipRect);
      return;
    }
  }

  if (asset && maskAsset && maskRect && imageTextureRenderer.renderMasked(asset, rect, maskAsset, maskRect, effectiveClipRect, effectiveOpacity, layer.transform.rotation, layer.sourceRect)) {
    if (showMaskRevealPreview) imageTextureRenderer.renderMaskRevealPreview(asset, rect, maskAsset, maskRect, maskRevealPreviewOpacity, effectiveClipRect, effectiveOpacity, layer.transform.rotation, layer.sourceRect);
    return;
  }
  if (asset && imageTextureRenderer.render(asset, rect, effectiveClipRect, effectiveOpacity, layer.transform.rotation, layer.sourceRect)) return;

  const fallbackRect = intersectScreenRects(rect, effectiveClipRect);
  if (!fallbackRect) return;
  clearScreenRect(context, fallbackRect, withOpacity(imageLayerColor, effectiveOpacity));
  const insetRect = intersectScreenRects({ x: rect.x + 4, y: rect.y + 4, w: Math.max(0, rect.w - 8), h: Math.max(0, rect.h - 8) }, effectiveClipRect);
  if (insetRect) clearScreenRect(context, insetRect, withOpacity(imageLayerInsetColor, effectiveOpacity));
}

const emptyLayerIds = new Set<string>();

function resolveLayerClipRect(
  context: WebGlRendererContext,
  documentIndex: DocumentReadIndex,
  viewport: ViewportState,
  layer: Layer,
  clipRect: ScreenRect,
): ScreenRect | undefined {
  const layerMask = getLayerMask(layer);
  if (!layerMask?.enabled) return clipRect;

  const maskBounds = resolveIndexedLayerBounds(documentIndex, layerMask.maskLayerId);
  if (!maskBounds) return clipRect;

  return intersectScreenRects(clipRect, documentRectToScreenRect(context.canvas, maskBounds, viewport));
}

function isIsolatedMaskView(mode: MaskViewMode) {
  return mode === "blackWhite" || mode === "alpha" || mode === "overlay";
}

function renderGenerationCandidatePreview(
  context: WebGlRendererContext,
  editor: EditorState,
  candidate: GenerationCandidate,
  imageTextureRenderer: ImageTextureRenderer,
  clipRect: ScreenRect,
) {
  const rect = documentRectToScreenRect(context.canvas, generationCandidateBounds(candidate), editor.viewport);
  const compareMode = editor.generation.compareMode ?? "result";
  if (compareMode === "before") return;

  if (compareMode === "split") {
    const splitClipRect = intersectScreenRects(clipRect, {
      x: rect.x + rect.w / 2,
      y: rect.y,
      w: rect.w / 2,
      h: rect.h,
    });
    if (!splitClipRect) return;
    imageTextureRenderer.render(generationCandidateAsset(candidate), rect, splitClipRect);

    const divider = intersectScreenRects(clipRect, {
      x: rect.x + rect.w / 2 - 1,
      y: rect.y,
      w: 2,
      h: rect.h,
    });
    if (divider) clearScreenRect(context, divider, comparisonDividerColor);
    return;
  }

  imageTextureRenderer.render(generationCandidateAsset(candidate), rect, clipRect);
}

function selectedGenerationCandidate(editor: EditorState): GenerationCandidate | undefined {
  return editor.generation.candidates.find((candidate) => candidate.id === editor.generation.selectedCandidateId) ?? editor.generation.candidates[0];
}

function generationCandidateAsset(candidate: GenerationCandidate): Asset {
  return {
    id: generationCandidateAssetId(candidate.id),
    name: candidate.placement.layerName,
    mimeType: candidate.mimeType,
    source: candidate.source,
    intrinsicSize: candidate.intrinsicSize,
  };
}

function generationCandidateAssetId(candidateId: string) {
  return `generation-candidate:${candidateId}`;
}

function generationCandidateBounds(candidate: GenerationCandidate): Rect {
  return {
    x: candidate.placement.transform.position.x,
    y: candidate.placement.transform.position.y,
    w: candidate.intrinsicSize.w * candidate.placement.transform.scale.x,
    h: candidate.intrinsicSize.h * candidate.placement.transform.scale.y,
  };
}

function assetWithBrushStrokePreview<TAsset extends ImageDocument["assets"][number] | undefined>(asset: TAsset, editor: EditorState): TAsset {
  if (!asset || editor.brushStrokePreview?.assetId !== asset.id) return asset;
  return { ...asset, source: editor.brushStrokePreview.source } as TAsset;
}

function intersectScreenRects(a: ScreenRect, b: ScreenRect): ScreenRect | undefined {
  const x1 = Math.max(a.x, b.x);
  const y1 = Math.max(a.y, b.y);
  const x2 = Math.min(a.x + a.w, b.x + b.w);
  const y2 = Math.min(a.y + a.h, b.y + b.h);
  if (x2 <= x1 || y2 <= y1) return undefined;

  return { x: x1, y: y1, w: x2 - x1, h: y2 - y1 };
}

function withOpacity(color: RgbaColor, opacity: number): RgbaColor {
  return [color[0], color[1], color[2], color[3] * opacity];
}