@basementstudio/shader-lab 0.1.0

package/packages/shader-lab-react/src/renderer/custom-shader-pass.ts

@@ -0,0 +1,108 @@
+import { clamp, float, type TSLNode, uniform, vec3, vec4 } from "three/tsl"
+import { CUSTOM_SHADER_ENTRY_EXPORT } from "../lib/editor/custom-shader/shared"
+import { compileCustomShaderModule } from "./custom-shader-runtime"
+import { PassNode } from "./pass-node"
+import type { LayerParameterValues } from "../types/editor"
+
+type Node = TSLNode
+
+export class CustomShaderPass extends PassNode {
+  private readonly onRuntimeError: ((message: string | null) => void) | undefined
+  private compiledSketch: (() => Node) | null = null
+  private compileRequestId = 0
+  private lastCompileSignature = ""
+  private readonly timeUniform: Node
+
+  constructor(layerId: string, onRuntimeError?: (message: string | null) => void) {
+    super(layerId)
+    this.onRuntimeError = onRuntimeError
+    this.timeUniform = uniform(0)
+    this.rebuildEffectNode()
+  }
+
+  override needsContinuousRender(): boolean {
+    return true
+  }
+
+  override updateParams(params: LayerParameterValues): void {
+    const sourceCode =
+      typeof params.sourceCode === "string" ? params.sourceCode : ""
+    const entryExport =
+      typeof params.entryExport === "string" && params.entryExport.trim()
+        ? params.entryExport.trim()
+        : CUSTOM_SHADER_ENTRY_EXPORT
+    const sourceFileName =
+      typeof params.sourceFileName === "string" ? params.sourceFileName : ""
+    const sourceRevision =
+      typeof params.sourceRevision === "number" ? params.sourceRevision : 0
+    const compileSignature = [
+      entryExport,
+      sourceCode,
+      sourceFileName,
+      sourceRevision,
+    ].join("\n")
+
+    if (compileSignature === this.lastCompileSignature) {
+      return
+    }
+
+    this.lastCompileSignature = compileSignature
+    this.compileRequestId += 1
+    const requestId = this.compileRequestId
+
+    void compileCustomShaderModule({
+      entryExport,
+      extraScope: {
+        time: this.timeUniform,
+      },
+      fileName: sourceFileName || "custom-shader.ts",
+      sourceCode,
+    })
+      .then((compiled) => {
+        if (requestId !== this.compileRequestId) {
+          return
+        }
+
+        this.compiledSketch = compiled.buildNode
+        this.onRuntimeError?.(null)
+        this.rebuildEffectNode()
+      })
+      .catch((error) => {
+        if (requestId !== this.compileRequestId) {
+          return
+        }
+
+        this.compiledSketch = null
+        this.onRuntimeError?.(
+          error instanceof Error ? error.message : "Custom shader compilation failed.",
+        )
+        this.rebuildEffectNode()
+      })
+  }
+
+  protected override beforeRender(time: number): void {
+    this.timeUniform.value = time
+  }
+
+  protected override buildEffectNode(): Node {
+    if (!this.compiledSketch) {
+      return vec4(vec3(float(0), float(0), float(0)), float(1))
+    }
+
+    try {
+      return vec4(
+        clamp(
+          this.compiledSketch(),
+          vec3(float(0), float(0), float(0)),
+          vec3(float(1), float(1), float(1))
+        ),
+        float(1)
+      )
+    } catch (error) {
+      this.onRuntimeError?.(
+        error instanceof Error ? error.message : "Custom shader execution failed.",
+      )
+      return vec4(vec3(float(0), float(0), float(0)), float(1))
+    }
+  }
+}

package/packages/shader-lab-react/src/renderer/custom-shader-runtime.ts

@@ -0,0 +1,309 @@
+import type { TSLNode } from "three/tsl"
+import * as tsl from "three/tsl"
+import * as shaderUtils from "./shaders/tsl/utils/index"
+
+type CompiledShaderModule = {
+  buildNode: () => TSLNode
+}
+
+const PRELUDE = {
+  ...tsl,
+  ...shaderUtils,
+}
+
+const TRANSPILED_CACHE = new Map<string, string>()
+let typescriptPromise: Promise<typeof import("typescript")> | null = null
+
+function isNodeLike(value: unknown): value is TSLNode {
+  return Boolean(
+    value &&
+      typeof value === "object" &&
+      "mul" in value &&
+      "add" in value &&
+      "sub" in value
+  )
+}
+
+function formatDiagnostics(
+  compiler: typeof import("typescript"),
+  diagnostics: readonly import("typescript").Diagnostic[] | undefined
+): string | null {
+  if (!(diagnostics && diagnostics.length > 0)) {
+    return null
+  }
+
+  return diagnostics
+    .map((diagnostic) =>
+      compiler.flattenDiagnosticMessageText(diagnostic.messageText, "\n")
+    )
+    .join("\n\n")
+}
+
+function createSourceFile(
+  compiler: typeof import("typescript"),
+  fileName: string,
+  sourceCode: string
+) {
+  return compiler.createSourceFile(
+    fileName,
+    sourceCode,
+    compiler.ScriptTarget.ES2020,
+    true,
+    getScriptKind(compiler, fileName)
+  )
+}
+
+function getSourceFileDiagnostics(sourceFile: import("typescript").SourceFile) {
+  return (
+    sourceFile as import("typescript").SourceFile & {
+      parseDiagnostics?: readonly import("typescript").Diagnostic[]
+    }
+  ).parseDiagnostics
+}
+
+function isDirectiveStatement(
+  compiler: typeof import("typescript"),
+  statement: import("typescript").Statement
+) {
+  return (
+    compiler.isExpressionStatement(statement) &&
+    compiler.isStringLiteral(statement.expression) &&
+    (statement.expression.text === "use client" ||
+      statement.expression.text === "use server")
+  )
+}
+
+function statementContainsJsx(
+  compiler: typeof import("typescript"),
+  statement: import("typescript").Statement
+): boolean {
+  let containsJsx = false
+
+  const visit = (node: import("typescript").Node) => {
+    if (
+      compiler.isJsxElement(node) ||
+      compiler.isJsxSelfClosingElement(node) ||
+      compiler.isJsxFragment(node)
+    ) {
+      containsJsx = true
+      return
+    }
+
+    compiler.forEachChild(node, visit)
+  }
+
+  compiler.forEachChild(statement, visit)
+
+  return containsJsx
+}
+
+async function sanitizeCustomShaderSource({
+  fileName,
+  sourceCode,
+}: {
+  fileName: string
+  sourceCode: string
+}) {
+  const compiler = await getTypeScript()
+  const sourceFile = createSourceFile(compiler, fileName, sourceCode)
+  const diagnosticsMessage = formatDiagnostics(
+    compiler,
+    getSourceFileDiagnostics(sourceFile)
+  )
+
+  if (diagnosticsMessage) {
+    throw new Error(diagnosticsMessage)
+  }
+
+  const statements = sourceFile.statements.filter((statement) => {
+    if (isDirectiveStatement(compiler, statement)) {
+      return false
+    }
+
+    if (
+      compiler.isImportDeclaration(statement) ||
+      compiler.isImportEqualsDeclaration(statement) ||
+      compiler.isExportAssignment(statement)
+    ) {
+      return false
+    }
+
+    if (
+      compiler.isExportDeclaration(statement) &&
+      statement.moduleSpecifier !== undefined
+    ) {
+      return false
+    }
+
+    if (statementContainsJsx(compiler, statement)) {
+      return false
+    }
+
+    return true
+  })
+
+  const sanitizedFile = compiler.factory.updateSourceFile(
+    sourceFile,
+    statements
+  )
+  const printer = compiler.createPrinter({
+    newLine: compiler.NewLineKind.LineFeed,
+  })
+
+  return `${printer.printFile(sanitizedFile).trim()}\n`
+}
+
+function assertNoExplicitImports(sourceCode: string) {
+  if (/^\s*import[\s{*]/m.test(sourceCode)) {
+    throw new Error(
+      "Custom shader imports are resolved through the injected prelude. Remove the imports or paste the whole sketch file and let the custom shader layer strip them."
+    )
+  }
+}
+
+async function getTypeScript() {
+  if (!typescriptPromise) {
+    typescriptPromise = import("typescript")
+  }
+
+  return typescriptPromise
+}
+
+function getScriptKind(
+  compiler: typeof import("typescript"),
+  fileName: string
+): import("typescript").ScriptKind {
+  return fileName.endsWith(".tsx")
+    ? compiler.ScriptKind.TSX
+    : compiler.ScriptKind.TS
+}
+
+export async function formatCustomShaderSource({
+  fileName,
+  sourceCode,
+}: {
+  fileName?: string
+  sourceCode: string
+}): Promise<string> {
+  const compiler = await getTypeScript()
+  const resolvedFileName = fileName ?? "custom-shader.ts"
+  const sourceFile = createSourceFile(compiler, resolvedFileName, sourceCode)
+  const diagnosticsMessage = formatDiagnostics(
+    compiler,
+    getSourceFileDiagnostics(sourceFile)
+  )
+
+  if (diagnosticsMessage) {
+    throw new Error(diagnosticsMessage)
+  }
+
+  const printer = compiler.createPrinter({
+    newLine: compiler.NewLineKind.LineFeed,
+  })
+
+  return `${printer.printFile(sourceFile).trim()}\n`
+}
+
+export async function compileCustomShaderModule({
+  entryExport,
+  extraScope,
+  fileName,
+  force = false,
+  sourceCode,
+}: {
+  entryExport: string
+  extraScope?: Record<string, unknown>
+  fileName?: string
+  force?: boolean
+  sourceCode: string
+}): Promise<CompiledShaderModule> {
+  const resolvedFileName = fileName ?? "custom-shader.ts"
+  const sanitizedSourceCode = await sanitizeCustomShaderSource({
+    fileName: resolvedFileName,
+    sourceCode,
+  })
+  assertNoExplicitImports(sanitizedSourceCode)
+
+  const cacheKey = `${entryExport}\n${resolvedFileName}\n${sanitizedSourceCode}`
+  let outputText = !force ? (TRANSPILED_CACHE.get(cacheKey) ?? null) : null
+
+  const compiler = await getTypeScript()
+  if (!outputText) {
+    const transpiled = compiler.transpileModule(sanitizedSourceCode, {
+      compilerOptions: {
+        esModuleInterop: true,
+        jsx: compiler.JsxEmit.ReactJSX,
+        module: compiler.ModuleKind.CommonJS,
+        target: compiler.ScriptTarget.ES2020,
+      },
+      fileName: resolvedFileName,
+      reportDiagnostics: true,
+    })
+    const diagnosticsMessage = formatDiagnostics(
+      compiler,
+      transpiled.diagnostics
+    )
+
+    if (diagnosticsMessage) {
+      throw new Error(diagnosticsMessage)
+    }
+
+    outputText = transpiled.outputText
+    TRANSPILED_CACHE.set(cacheKey, outputText)
+  }
+
+  const runtimeScope = {
+    ...PRELUDE,
+    ...(extraScope ?? {}),
+  }
+  const scopeNames = Object.keys(runtimeScope)
+  const scopeValues = scopeNames.map(
+    (key) => runtimeScope[key as keyof typeof runtimeScope]
+  )
+  const module = { exports: {} as Record<string, unknown> }
+  const exportsObject = module.exports
+  const evaluator = new Function(
+    "exports",
+    "module",
+    ...scopeNames,
+    `${outputText}\nreturn module.exports;`
+  )
+
+  let resolvedExports: Record<string, unknown>
+
+  try {
+    resolvedExports = evaluator(
+      exportsObject,
+      module,
+      ...scopeValues
+    ) as Record<string, unknown>
+  } catch (error) {
+    throw new Error(
+      error instanceof Error
+        ? error.message
+        : "Custom shader evaluation failed."
+    )
+  }
+
+  const exported = resolvedExports[entryExport] ?? module.exports[entryExport]
+
+  if (typeof exported !== "function") {
+    throw new Error(
+      `Expected a named export \`${entryExport}\` that resolves to a TSL sketch function.`
+    )
+  }
+
+  return {
+    buildNode: () => {
+      const result = (exported as () => unknown)()
+
+      if (!isNodeLike(result)) {
+        throw new Error(
+          `The export \`${entryExport}\` did not return a valid TSL node.`
+        )
+      }
+
+      return result
+    },
+  }
+}

package/packages/shader-lab-react/src/renderer/dither-textures.ts

@@ -0,0 +1,99 @@
+import * as THREE from "three/webgpu"
+
+const BAYER_2X2 = [0, 2, 3, 1] as const
+
+const BAYER_4X4 = [
+  0, 8, 2, 10,
+  12, 4, 14, 6,
+  3, 11, 1, 9,
+  15, 7, 13, 5,
+] as const
+
+const BAYER_8X8 = [
+  0, 32, 8, 40, 2, 34, 10, 42,
+  48, 16, 56, 24, 50, 18, 58, 26,
+  12, 44, 4, 36, 14, 46, 6, 38,
+  60, 28, 52, 20, 62, 30, 54, 22,
+  3, 35, 11, 43, 1, 33, 9, 41,
+  51, 19, 59, 27, 49, 17, 57, 25,
+  15, 47, 7, 39, 13, 45, 5, 37,
+  63, 31, 55, 23, 61, 29, 53, 21,
+] as const
+
+function buildBayerTexture(values: readonly number[], size: number): THREE.DataTexture {
+  const normalizer = size * size
+  const data = new Uint8Array(size * size * 4)
+
+  for (let index = 0; index < size * size; index += 1) {
+    const channel = Math.round(((values[index] ?? 0) / normalizer) * 255)
+    const offset = index * 4
+    data[offset] = channel
+    data[offset + 1] = channel
+    data[offset + 2] = channel
+    data[offset + 3] = 255
+  }
+
+  const texture = new THREE.DataTexture(
+    data,
+    size,
+    size,
+    THREE.RGBAFormat,
+    THREE.UnsignedByteType,
+  )
+  texture.magFilter = THREE.NearestFilter
+  texture.minFilter = THREE.NearestFilter
+  texture.wrapS = THREE.RepeatWrapping
+  texture.wrapT = THREE.RepeatWrapping
+  texture.needsUpdate = true
+
+  return texture
+}
+
+function buildBlueNoiseTexture(size = 64): THREE.DataTexture {
+  const data = new Uint8Array(size * size * 4)
+  const fract = (value: number) => value - Math.floor(value)
+
+  for (let y = 0; y < size; y += 1) {
+    for (let x = 0; x < size; x += 1) {
+      const noise = Math.round(
+        fract(52.9829189 * fract(0.06711056 * x + 0.00583715 * y)) * 255,
+      )
+      const offset = (y * size + x) * 4
+      data[offset] = noise
+      data[offset + 1] = noise
+      data[offset + 2] = noise
+      data[offset + 3] = 255
+    }
+  }
+
+  const texture = new THREE.DataTexture(
+    data,
+    size,
+    size,
+    THREE.RGBAFormat,
+    THREE.UnsignedByteType,
+  )
+  texture.magFilter = THREE.NearestFilter
+  texture.minFilter = THREE.NearestFilter
+  texture.wrapS = THREE.RepeatWrapping
+  texture.wrapT = THREE.RepeatWrapping
+  texture.needsUpdate = true
+
+  return texture
+}
+
+export interface DitherTextures {
+  bayer2: THREE.DataTexture
+  bayer4: THREE.DataTexture
+  bayer8: THREE.DataTexture
+  noise: THREE.DataTexture
+}
+
+export function buildDitherTextures(): DitherTextures {
+  return {
+    bayer2: buildBayerTexture(BAYER_2X2, 2),
+    bayer4: buildBayerTexture(BAYER_4X4, 4),
+    bayer8: buildBayerTexture(BAYER_8X8, 8),
+    noise: buildBlueNoiseTexture(),
+  }
+}