wesl 0.6.0-pre10

package/src/LinkedWesl.ts

@@ -0,0 +1,184 @@
+import { SrcMap } from "mini-parse";
+import { assertThat } from "../../mini-parse/src/Assertions";
+import { errorHighlight, offsetToLineNumber } from "./Util";
+import type { WeslDevice } from "./WeslDevice";
+
+/** Results of shader compilation. Has {@link WeslGPUCompilationMessage}
+ * which are aware of the WESL module that an error was thrown from. */
+export interface WeslGPUCompilationInfo extends GPUCompilationInfo {
+  messages: WeslGPUCompilationMessage[];
+}
+
+export interface WeslGPUCompilationMessage extends GPUCompilationMessage {
+  module: {
+    // LATER this should be a qualified module path.
+    // And something else should map it to a URL that is relative to the correct place.
+    url: string;
+    // LATER: I don't think that the text should be a part of the compilation message.
+    // Instead the module url should be usable as a key.
+    text?: string;
+  };
+}
+
+/**
+ * A {@link GPUValidationError} with an inner error (for a stack trace).
+ * Can also point at a WESL source file.
+ */
+export interface ExtendedGPUValidationError extends GPUValidationError {
+  cause?: Error;
+  compilationInfo?: WeslGPUCompilationInfo;
+}
+
+/**
+ * Multiple WESL files that have been linked together to produce WGSL code.
+ *
+ * Call {@link LinkedWesl.createShaderModule} on a {@link WeslDevice}
+ * to make the error reporting aware of the WESL code.
+ */
+export class LinkedWesl {
+  constructor(public sourceMap: SrcMap) {}
+
+  /**
+   * Creates a {@link GPUShaderModule}.
+   * When errors occur, they will point at the original WESL source code.
+   *
+   * The compilation info {@link GPUShaderModule.getCompilationInfo}
+   * can be remapped with {@link mapGPUCompilationInfo}
+   * @param device GPUDevice. Preferably a {@link WeslDevice} for better error reporting.
+   * @param descriptor - Description of the {@link GPUShaderModule} to create.
+   */
+  createShaderModule(
+    device: GPUDevice | WeslDevice,
+    descriptor: Omit<GPUShaderModuleDescriptor, "code">,
+  ): GPUShaderModule {
+    // Skip the custom behaviour if we do not have a WESL device.
+    if (!("injectError" in device)) {
+      return device.createShaderModule({
+        ...descriptor,
+        code: this.dest,
+      });
+    }
+
+    device.pushErrorScope("validation"); // Suppress the normal error
+    const module = device.createShaderModule({
+      ...descriptor,
+      code: this.dest,
+    });
+    device.popErrorScope();
+    // And report the error!
+    let { promise, resolve } = Promise.withResolvers<GPUError | null>();
+    device.injectError("validation", promise); // Inject our custom error
+    module.getCompilationInfo().then(compilationInfo => {
+      if (compilationInfo.messages.length === 0) {
+        resolve(null);
+        return;
+      }
+
+      const mappedCompilationInfo = this.mapGPUCompilationInfo(compilationInfo);
+      const errorMessage = compilationInfoToErrorMessage(
+        mappedCompilationInfo,
+        module,
+      );
+      // Error message cannot be null, since we're passing at least one message to it.
+      assertThat(errorMessage !== null);
+      const error: ExtendedGPUValidationError = new GPUValidationError(
+        errorMessage,
+      );
+      error.cause = new Error("createShaderModule failed");
+      error.compilationInfo = mappedCompilationInfo;
+      resolve(error);
+    });
+    return module;
+  }
+
+  /**
+   * Use {@link LinkedWesl.createShaderModule} for a
+   * better error reporting experience.
+   */
+  get dest() {
+    return this.sourceMap.dest.text;
+  }
+
+  /** Turns raw compilation info into compilation info
+   * that points at the WESL sources. */
+  public mapGPUCompilationInfo(
+    compilationInfo: GPUCompilationInfo,
+  ): WeslGPUCompilationInfo {
+    return {
+      __brand: compilationInfo.__brand,
+      messages: compilationInfo.messages.map(v =>
+        this.mapGPUCompilationMessage(v),
+      ),
+    };
+  }
+
+  private mapGPUCompilationMessage(
+    message: GPUCompilationMessage,
+  ): WeslGPUCompilationMessage {
+    const srcMap = this.sourceMap;
+    const srcPosition = srcMap.destToSrc(message.offset);
+    // LATER what if this gets mapped to a completely different place?
+    const srcEndPosition =
+      message.length > 0 ?
+        srcMap.destToSrc(message.offset + message.length)
+      : srcPosition;
+    const length = srcEndPosition.position - srcPosition.position;
+
+    let [lineNum, linePos] = offsetToLineNumber(
+      srcPosition.position,
+      srcPosition.src.text,
+    );
+
+    return {
+      __brand: message.__brand,
+      type: message.type,
+      message: message.message,
+      offset: srcPosition.position,
+      length,
+      lineNum,
+      linePos,
+      module: {
+        url: srcPosition.src.path ?? "",
+        text: srcPosition.src.text,
+      },
+    };
+  }
+}
+
+/**
+ * Tries to imitate the way the browser logs the compilation info.
+ * Does not do the remapping.
+ * @returns A string with errors, or `null` if there were no compilation messages.
+ */
+function compilationInfoToErrorMessage(
+  compilationInfo: WeslGPUCompilationInfo,
+  shaderModule: GPUShaderModule,
+): string | null {
+  if (compilationInfo.messages.length === 0) return null;
+
+  let result = `Compilation log for [Invalid ShaderModule (${
+    shaderModule.label || "unlabled"
+  })]:\n`;
+  let errorCount = compilationInfo.messages.filter(
+    v => v.type === "error",
+  ).length;
+  if (errorCount > 0) {
+    result += `${errorCount} error(s) generated while compiling the shader:\n`;
+  }
+  for (const message of compilationInfo.messages) {
+    const { lineNum, linePos } = message;
+
+    result += `${message.module.url}:${lineNum}:${linePos}`;
+    result += ` ${message.type}: ${message.message}\n`;
+    // LATER unmangle code snippets in the message
+
+    const source = message.module.text;
+    if (source) {
+      result += errorHighlight(source, [
+        message.offset,
+        message.offset + message.length,
+      ]).join("\n");
+    }
+  }
+  return result;
+}

package/src/Linker.ts

@@ -0,0 +1,284 @@
+import { SrcMap, SrcMapBuilder, tracing } from "mini-parse";
+import { AbstractElem, ModuleElem } from "./AbstractElems.ts";
+import { bindIdents, EmittableElem } from "./BindIdents.ts";
+import { LinkedWesl } from "./LinkedWesl.ts";
+import { lowerAndEmit } from "./LowerAndEmit.ts";
+import { ManglerFn } from "./Mangler.ts";
+import {
+  parsedRegistry,
+  ParsedRegistry,
+  parseIntoRegistry,
+  parseLibsIntoRegistry,
+  selectModule,
+} from "./ParsedRegistry.ts";
+import { WeslAST } from "./ParseWESL.ts";
+import { Conditions, DeclIdent, SrcModule } from "./Scope.ts";
+import { filterMap, mapValues } from "./Util.ts";
+import { WeslBundle } from "./WeslBundle.ts";
+
+type LinkerTransform = (boundAST: TransformedAST) => TransformedAST;
+
+export interface WeslJsPlugin {
+  transform?: LinkerTransform;
+}
+
+export interface TransformedAST
+  extends Pick<WeslAST, "srcModule" | "moduleElem"> {
+  globalNames: Set<string>;
+  notableElems: Record<string, AbstractElem[]>;
+}
+
+export interface LinkConfig {
+  plugins?: WeslJsPlugin[];
+}
+
+export interface LinkParams {
+  /** record of file paths and wesl text for modules.
+   *   key is module path or file path
+   *     `package::foo::bar`, or './foo/bar.wesl', or './foo/bar'
+   *   value is wesl src
+   *
+   *
+   * Only accepts unix-style, relative filesystem paths that are valid WGSL identifiers
+   * - Unix-style: Slashes as separators.
+   * - Valid WGSL identifiers: No backslashes, no `..`, or other non-identifier symbols.
+   * - Relative paths: They have to be relative to the wesl root.
+   */
+  weslSrc: Record<string, string>;
+
+  /** name of root wesl module
+   *    for an app, the root module normally contains the '@compute', '@vertex' or '@fragment' entry points
+   *    for a library, the root module defines the public api fo the library
+   *  can be specified as file path (./main.wesl), a module path (package::main), or just a module name (main)
+   */
+  rootModuleName?: string;
+
+  /** For debug logging. Will be prepended to file paths. */
+  debugWeslRoot?: string;
+
+  /** runtime conditions for conditional compiling with @if and friends */
+  conditions?: Conditions;
+
+  /** libraries available for the link */
+  libs?: WeslBundle[];
+
+  /** generate wesl from code at runtime */
+  virtualLibs?: Record<string, VirtualLibraryFn>;
+
+  /** plugins and other configuration to use while linking */
+  config?: LinkConfig;
+
+  /** Host (ts/js) provided wgsl constants.
+   * Users can import the values from wesl code via the `constants' virtual library:
+   *  `import constants::num_lights;` */
+  constants?: Record<string, string | number>;
+
+  /** function to construct globally unique wgsl identifiers */
+  mangler?: ManglerFn;
+}
+
+/** Generate a virtual WESL module based on a set of conditions. */
+export type VirtualLibraryFn = (conditions: Conditions) => string;
+
+/**
+ * Link a set of WESL source modules (typically the text from .wesl files) into a single WGSL string.
+ * Linking starts with a specified 'root' source module, and recursively incorporates code
+ * referenced from other modules (in local files or libraries).
+ *
+ * Unreferenced (dead) code outside the root module is not included in the output WGSL.
+ * Additionally the caller can specify conditions for to control conditional compilation.
+ * Only code that is valid with the current conditions is included in the output.
+ */
+export async function link(params: LinkParams): Promise<LinkedWesl> {
+  const { weslSrc, debugWeslRoot, libs = [] } = params;
+  const registry = parsedRegistry();
+  parseIntoRegistry(weslSrc, registry, "package", debugWeslRoot);
+  parseLibsIntoRegistry(libs, registry);
+  return new LinkedWesl(linkRegistry({ registry, ...params }));
+}
+
+export interface LinkRegistryParams
+  extends Pick<
+    LinkParams,
+    | "rootModuleName"
+    | "conditions"
+    | "virtualLibs"
+    | "config"
+    | "constants"
+    | "mangler"
+  > {
+  registry: ParsedRegistry;
+}
+
+/** Link wesl from a registry of already parsed modules.
+ *
+ * This entry point is intended for users who want to link multiple times
+ * from the same sources. (e.g. linking with different conditions
+ * each time, or perhaps to produce multiple wgsl shaders
+ * that share some sources.)
+ */
+export function linkRegistry(params: LinkRegistryParams): SrcMap {
+  const bound = bindAndTransform(params);
+  const { transformedAst, newDecls, newStatements } = bound;
+
+  return SrcMapBuilder.build(
+    emitWgsl(
+      transformedAst.moduleElem,
+      transformedAst.srcModule,
+      newDecls,
+      newStatements,
+      params.conditions,
+    ),
+  );
+}
+
+export interface BoundAndTransformed {
+  transformedAst: TransformedAST;
+  newDecls: DeclIdent[];
+  newStatements: EmittableElem[];
+}
+
+/** bind identifers and apply any transform plugins */
+export function bindAndTransform(
+  params: LinkRegistryParams,
+): BoundAndTransformed {
+  const { registry, mangler } = params;
+  const { rootModuleName = "main", conditions = {} } = params;
+  const rootAst = getRootModule(registry, rootModuleName);
+
+  // setup virtual modules from code generation or host constants provided by the user
+  const { constants, config } = params;
+  let { virtualLibs } = params;
+  if (constants) {
+    virtualLibs = { ...virtualLibs, constants: constantsGenerator(constants) };
+  }
+  let virtuals = virtualLibs && mapValues(virtualLibs, fn => ({ fn }));
+
+  /* --- Step #2   Binding Idents --- */
+  // link active Ident references to declarations, and uniquify global declarations
+  const bindParams = { rootAst, registry, conditions, virtuals, mangler };
+  const bindResults = bindIdents(bindParams);
+  const { globalNames, decls: newDecls, newStatements } = bindResults;
+
+  const transformedAst = applyTransformPlugins(rootAst, globalNames, config);
+  return { transformedAst, newDecls, newStatements };
+}
+
+function constantsGenerator(
+  constants: Record<string, string | number>,
+): () => string {
+  return () =>
+    Object.entries(constants)
+      .map(([name, value]) => `const ${name} = ${value};`)
+      .join("\n");
+}
+
+/** get a reference to the root module, selecting by module name */
+function getRootModule(
+  parsed: ParsedRegistry,
+  rootModuleName: string,
+): WeslAST {
+  const rootModule = selectModule(parsed, rootModuleName);
+  if (!rootModule) {
+    if (tracing) {
+      console.log(`parsed modules: ${Object.keys(parsed.modules)}`);
+      console.log(`root module not found: ${rootModuleName}`);
+    }
+    throw new Error(`Root module not found: ${rootModuleName}`);
+  }
+  return rootModule;
+}
+
+/** run any plugins that transform the AST */
+function applyTransformPlugins(
+  rootModule: WeslAST,
+  globalNames: Set<string>,
+  config?: LinkConfig,
+): TransformedAST {
+  const { moduleElem, srcModule } = rootModule;
+
+  // for now only transform the root module
+  const startAst = { moduleElem, srcModule, globalNames, notableElems: {} };
+  const plugins = config?.plugins ?? [];
+  const transforms = filterMap(plugins, plugin => plugin.transform);
+  const transformedAst = transforms.reduce(
+    (ast, transform) => transform(ast),
+    startAst,
+  );
+
+  return transformedAst;
+}
+
+/** traverse the AST and emit WGSL */
+function emitWgsl(
+  rootModuleElem: ModuleElem,
+  srcModule: SrcModule,
+  newDecls: DeclIdent[],
+  newStatements: EmittableElem[],
+  conditions: Conditions = {},
+): SrcMapBuilder[] {
+  /* --- Step #3   Writing WGSL --- */ // note doesn't require the scope tree anymore
+
+  // emit any new statements (module level const asserts)
+  const prologueBuilders = newStatements.map(s => {
+    const { elem, srcModule } = s;
+    const { src: text, debugFilePath: path } = srcModule;
+    const builder = new SrcMapBuilder({ text, path });
+    lowerAndEmit(builder, [elem], conditions);
+    builder.addNl();
+    return builder;
+  });
+
+  const rootBuilder = new SrcMapBuilder({
+    text: srcModule.src,
+    path: srcModule.debugFilePath,
+  });
+  lowerAndEmit(rootBuilder, [rootModuleElem], conditions, false); // emit the entire root module
+
+  const declBuilders = newDecls.map(decl => {
+    const builder = new SrcMapBuilder({
+      text: decl.srcModule.src,
+      path: decl.srcModule.debugFilePath,
+    });
+    lowerAndEmit(builder, [decl.declElem!], conditions); // emit referenced declarations from other modules
+    return builder;
+  });
+
+  return [...prologueBuilders, rootBuilder, ...declBuilders];
+}
+
+/* ---- Commentary on present and future features ---- */
+/*
+
+TODO 
+- distinguish between global and local declaration idents (only global ones need be uniquified)
+
+Conditions
+- conditions are attached to the AST elements where they are defined
+  - only conditionally valid elements are emitted
+- consolidated conditions are attached to Idents
+  - only conditionally valid ref Idents are bound, and only to conditionaly valid declarations
+  - a condition stack (akin to the scope stack) is maintained while parsing to attach consolidated conditions to Idents
+- re-linking with new conditions, conservatively 
+  - clear all mutated Ident fields (refersTo and mangled links) 
+  - re-bind Idents, re-emit 
+
+Generics & specialization
+- attach generic parameters to ref and decl Idents, effectively creating a new Ident for each specialization
+- generate specialized elements at emit time, by checking the generic parameters of the decl ident
+
+Incrementally rebuilding
+- unchanged files don't need to be reparsed, only reparse dirty files.
+- support reflection only mode? no need to bind idents or emit for e.g. vite/IDE plugin generating reflection types 
+
+Parallel Processing (coarse grained via webworkers)
+- Parsing each module can be done in parallel
+- binding could be done partially in parallel? (esbuild doesn't parallelize here though)
+  - finding the declaration for each local ident could be done in parallel by module
+  - matching 
+- Emitting could be easily modified to be done in partially in parallel
+  - traversing the AST to list the top level elements to emit could be done serially
+  - the text for each top level element could be emitted in parallel (presumably the bulk of the work)
+  - the merged text can be assembled serially
+
+*/

package/src/LinkerUtil.ts

@@ -0,0 +1,29 @@
+import { srcLog } from "mini-parse";
+import {
+  AbstractElem,
+  ContainerElem,
+  DeclIdentElem,
+  RefIdentElem,
+} from "./AbstractElems.ts";
+
+export function visitAst(
+  elem: AbstractElem,
+  visitor: (elem: AbstractElem) => void,
+) {
+  visitor(elem);
+  if ((elem as ContainerElem).contents) {
+    const container = elem as ContainerElem;
+    container.contents.forEach(child => visitAst(child, visitor));
+  }
+}
+
+export function identElemLog(
+  identElem: DeclIdentElem | RefIdentElem,
+  ...messages: any[]
+): void {
+  srcLog(
+    identElem.srcModule.src,
+    [identElem.start, identElem.end],
+    ...messages,
+  );
+}

package/src/LiveDeclarations.ts

@@ -0,0 +1,31 @@
+import { identToString } from "./debug/ScopeToString.ts";
+import { DeclIdent } from "./Scope.ts";
+
+/** decls currently visible in this scope */
+export interface LiveDecls {
+  /** decls currently visible in this scope */
+  decls: Map<string, DeclIdent>;
+
+  /** live decls in the parent scope. null for the modue root scope */
+  parent?: LiveDecls | null;
+}
+
+/** create a LiveDecls */
+export function makeLiveDecls(parent: LiveDecls | null = null): LiveDecls {
+  return { decls: new Map<string, DeclIdent>(), parent };
+}
+
+/** debug routine for logging LiveDecls */
+export function liveDeclsToString(liveDecls: LiveDecls): string {
+  const { decls, parent } = liveDecls;
+  const declsStr = Array.from(decls.entries())
+    .map(([name, decl]) => `${name}:${identToString(decl)}`)
+    .join(", ");
+  const parentStr = parent ? liveDeclsToString(parent) : "null";
+  return `decls: { ${declsStr} }, parent: ${parentStr}`;
+}
+
+/*
+LATER try not creating a map for small scopes. 
+Instead just track the current live index in the scope array.
+*/