wesl 0.7.26 → 0.7.28

package/src/Linker.ts

@@ -14,7 +14,7 @@ import {
   type ModuleResolver,
   RecordResolver,
 } from "./ModuleResolver.ts";
-import type { WeslAST } from "./ParseWESL.ts";
+import type { WeslAST, WeslExtensions } from "./ParseWESL.ts";
 import type { Conditions, DeclIdent, SrcModule } from "./Scope.ts";
 import { type SrcMap, SrcMapBuilder } from "./SrcMap.ts";
 import { filterMap, mapValues } from "./Util.ts";
@@ -37,29 +37,17 @@ export interface LinkConfig {
 }
 
 export interface LinkParams {
-  /** Module resolver for lazy loading (NEW API - preferred)
-   * Replaces weslSrc for lazy module loading.
-   * If provided, weslSrc will be ignored. */
+  /** Module resolver for lazy loading. If provided, weslSrc is ignored. */
   resolver?: ModuleResolver;
 
-  /** record of file paths and wesl text for modules (LEGACY API - still supported).
-   *   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.
-   */
+  /** Record of module sources keyed by module path (`package::foo::bar`)
+   * or relative file path (`./foo/bar.wesl`). Paths must be unix-style,
+   * relative to the wesl root, and valid WGSL identifiers. */
   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)
-   */
+  /** Root module name: file path (`./main.wesl`), module path (`package::main`), or bare name (`main`).
+   * For apps, the root module contains entry points (`@compute`, `@vertex`, `@fragment`).
+   * For libraries, it defines the public API. */
   rootModuleName?: string;
 
   /** For debug logging. Will be prepended to file paths. */
@@ -89,6 +77,11 @@ export interface LinkParams {
 
   /** function to construct globally unique wgsl identifiers */
   mangler?: ManglerFn;
+
+  /** opt-in parsing of experimental, not-yet-spec'd syntax extensions.
+   * Only applied to the local source resolver built from `weslSrc`; callers
+   * that supply their own `resolver` must set it on that resolver. */
+  weslExtensions?: WeslExtensions;
 }
 
 /** Project config for web components and tools. */
@@ -100,7 +93,11 @@ export type WeslProject = Pick<
   | "constants"
   | "libs"
   | "packageName"
->;
+  | "weslExtensions"
+> & {
+  /** Shader directory relative to project root (set by ?link from wesl.toml). */
+  shaderRoot?: string;
+};
 
 /** Context passed to virtual library generators. */
 export interface VirtualLibContext {
@@ -127,26 +124,26 @@ export async function link(params: LinkParams): Promise<LinkedWesl> {
 
 /** linker api for benchmarking */
 export function _linkSync(params: LinkParams): SrcMap {
-  const { weslSrc, libs = [], packageName, debugWeslRoot } = params;
-  const { resolver } = params;
-
-  const resolvers: ModuleResolver[] = [];
+  const { weslSrc, libs = [], packageName, debugWeslRoot, resolver } = params;
+  const { weslExtensions } = params;
 
-  if (resolver) {
-    resolvers.push(resolver);
-  } else if (weslSrc) {
-    resolvers.push(new RecordResolver(weslSrc, { packageName, debugWeslRoot }));
-  } else {
+  if (!resolver && !weslSrc) {
     throw new Error("Either resolver or weslSrc must be provided");
   }
-
-  if (libs.length > 0) {
-    const libResolvers = createLibraryResolvers(libs, debugWeslRoot);
-    resolvers.push(...libResolvers);
-  }
-
+  const primaryResolver =
+    resolver ??
+    new RecordResolver(weslSrc!, {
+      packageName,
+      debugWeslRoot,
+      weslExtensions,
+    });
+
+  const libResolvers = createLibraryResolvers(libs, debugWeslRoot);
+  const allResolvers = [primaryResolver, ...libResolvers];
   const finalResolver =
-    resolvers.length === 1 ? resolvers[0] : new CompositeResolver(resolvers);
+    allResolvers.length === 1
+      ? allResolvers[0]
+      : new CompositeResolver(allResolvers);
 
   return linkRegistry({ ...params, resolver: finalResolver });
 }
@@ -200,11 +197,8 @@ export interface LinkRegistryParams
  * that share some sources.)
  */
 export function linkRegistry(params: LinkRegistryParams): SrcMap {
-  const {
-    transformedAst: ast,
-    newDecls,
-    newStatements,
-  } = bindAndTransform(params);
+  const bound = bindAndTransform(params);
+  const { transformedAst: ast, newDecls, newStatements } = bound;
   const builders = emitWgsl(
     ast.moduleElem,
     ast.srcModule,
@@ -230,7 +224,6 @@ export function bindAndTransform(
 
   const modulePath = normalizeModuleName(rootModuleName);
   const rootAst = getRootModule(resolver, modulePath, rootModuleName);
-
   const virtuals = setupVirtualLibs(params.virtualLibs, constants);
 
   const bound = bindIdents({
@@ -291,14 +284,14 @@ function setupVirtualLibs(
   return libs && mapValues(libs, fn => ({ fn }));
 }
 
+/** Run registered transform plugins over the bound AST. */
 function applyTransformPlugins(
   rootModule: WeslAST,
   globalNames: Set<string>,
   config?: LinkConfig,
 ): TransformedAST {
-  const { moduleElem, srcModule } = rootModule;
-
   // for now only transform the root module
+  const { moduleElem, srcModule } = rootModule;
   const startAst = { moduleElem, srcModule, globalNames, notableElems: {} };
   const plugins = config?.plugins ?? [];
   const transforms = filterMap(plugins, plugin => plugin.transform);
@@ -313,12 +306,11 @@ function emitWgsl(
   newStatements: EmittableElem[],
   conditions: Conditions = {},
 ): SrcMapBuilder[] {
-  const prologueBuilders = newStatements.map(s => emitStatement(s, conditions));
+  const prologueBuilders = newStatements.map(s =>
+    emitElem(s.srcModule, s.elem, conditions, { addNl: true }),
+  );
 
-  const rootBuilder = new SrcMapBuilder({
-    text: srcModule.src,
-    path: srcModule.debugFilePath,
-  });
+  const rootBuilder = builderFromModule(srcModule);
   lowerAndEmit({
     srcBuilder: rootBuilder,
     rootElems: [rootModuleElem],
@@ -326,36 +318,40 @@ function emitWgsl(
     extracting: false,
   });
 
-  const declBuilders = newDecls.map(decl => emitDecl(decl, conditions));
+  const declBuilders = newDecls.map(decl =>
+    emitElem(decl.srcModule, decl.declElem!, conditions, {
+      skipConditionalFiltering: true,
+    }),
+  );
 
   return [...prologueBuilders, rootBuilder, ...declBuilders];
 }
 
-function emitStatement(
-  s: EmittableElem,
+/** Emit a single element (prologue statement or imported declaration) into a SrcMapBuilder. */
+function emitElem(
+  srcModule: SrcModule,
+  elem: AbstractElem,
   conditions: Conditions,
+  opts: { addNl?: boolean; skipConditionalFiltering?: boolean } = {},
 ): SrcMapBuilder {
-  const { elem, srcModule } = s;
-  const { src: text, debugFilePath: path } = srcModule;
-  const builder = new SrcMapBuilder({ text, path });
-  lowerAndEmit({ srcBuilder: builder, rootElems: [elem], conditions });
-  builder.addNl();
-  return builder;
-}
-
-/** Skip conditional filtering because findValidRootDecls already validated these declarations */
-function emitDecl(decl: DeclIdent, conditions: Conditions): SrcMapBuilder {
-  const { src: text, debugFilePath: path } = decl.srcModule;
-  const builder = new SrcMapBuilder({ text, path });
+  const builder = builderFromModule(srcModule);
   lowerAndEmit({
     srcBuilder: builder,
-    rootElems: [decl.declElem!],
+    rootElems: [elem],
     conditions,
-    skipConditionalFiltering: true,
+    skipConditionalFiltering: opts.skipConditionalFiltering,
   });
+  if (opts.addNl) builder.addNl();
   return builder;
 }
 
+function builderFromModule(srcModule: SrcModule): SrcMapBuilder {
+  return new SrcMapBuilder({
+    text: srcModule.src,
+    path: srcModule.debugFilePath,
+  });
+}
+
 /*
 
 LATER

package/src/LinkerUtil.ts

@@ -1,22 +1,40 @@
 import type {
   AbstractElem,
-  ContainerElem,
   DeclIdentElem,
+  ExpressionElem,
+  HasAttributes,
+  ModuleElem,
   RefIdentElem,
 } from "./AbstractElems.ts";
+import { assertUnreachableSilent } from "./Assertions.ts";
 import { srcLog } from "./Logging.ts";
 
+/** A module's top-level declarations of a given kind, narrowed to that elem type. */
+export function declsOfKind<K extends AbstractElem["kind"]>(
+  moduleElem: ModuleElem,
+  kind: K,
+): Extract<AbstractElem, { kind: K }>[] {
+  return moduleElem.decls.filter(
+    (e): e is Extract<AbstractElem, { kind: K }> => e.kind === kind,
+  );
+}
+
+/** Visit an elem and all its descendants, parent before children (pre-order). */
 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);
-    });
-  }
+  for (const child of childElems(elem)) visitAst(child, visitor);
+}
+
+/**
+ * Child elems of any AST node: the typed structural fields (attributes plus
+ * body / condition / decls / ...) in source order. Returns [] for leaf elems.
+ */
+export function childElems(elem: AbstractElem): readonly AbstractElem[] {
+  const attributes = (elem as HasAttributes).attributes ?? [];
+  return [...attributes, ...structuralFields(elem)];
 }
 
 export function identElemLog(
@@ -29,3 +47,119 @@ export function identElemLog(
     ...messages,
   );
 }
+
+/** The child elems held in a node's typed fields, in source order (attributes
+ *  are added separately by {@link childElems}). Empty for leaf kinds. The
+ *  switch is exhaustive: a new elem kind without a case fails to compile in the
+ *  default, so children can't be silently dropped. */
+function structuralFields(elem: AbstractElem): AbstractElem[] {
+  switch (elem.kind) {
+    case "module":
+      return elem.decls;
+    case "var":
+    case "gvar":
+      return [
+        ...(elem.template ?? []),
+        elem.name,
+        ...(elem.init ? [elem.init] : []),
+      ];
+    case "let":
+    case "const":
+    case "override":
+      return [elem.name, ...(elem.init ? [elem.init] : [])];
+    case "alias":
+      return [elem.name, elem.typeRef];
+    case "assert":
+      return [elem.expression];
+    case "struct":
+      return [elem.name, ...elem.members];
+    case "member":
+      return [elem.name, elem.typeRef];
+    case "type":
+      return [elem.name.refIdentElem, ...(elem.templateParams ?? [])];
+    case "expression":
+      return [elem.expression];
+    case "binary-expression":
+      return [elem.left, elem.right];
+    case "unary-expression":
+    case "parenthesized-expression":
+      return [elem.expression];
+    case "component-expression":
+      return [elem.base, elem.access];
+    case "component-member-expression":
+      return [elem.base, elem.access];
+    case "call-expression":
+      return [elem.function, ...(elem.templateArgs ?? []), ...elem.arguments];
+    case "param":
+      return [elem.name];
+    case "typeDecl":
+      return elem.typeRef ? [elem.decl, elem.typeRef] : [elem.decl];
+    case "fn":
+      return [
+        elem.name,
+        ...elem.params,
+        ...(elem.returnType ? [elem.returnType] : []),
+        elem.body,
+      ];
+    case "block":
+      return elem.body;
+    case "if":
+      return [elem.condition, elem.body, ...(elem.else ? [elem.else] : [])];
+    case "for":
+      return [elem.init, elem.condition, elem.update, elem.body].filter(
+        isDefined,
+      );
+    case "while":
+      return [elem.condition, elem.body];
+    case "loop":
+    case "continuing":
+      return [elem.body];
+    case "switch":
+      return [elem.selector, ...(elem.bodyAttributes ?? []), ...elem.clauses];
+    case "switch-clause":
+      return [...exprSelectors(elem.selectors), elem.body];
+    case "return":
+      return elem.value ? [elem.value] : [];
+    case "break":
+      return elem.condition ? [elem.condition] : [];
+    case "assign":
+      return elem.lhs.kind === "phony" ? [elem.rhs] : [elem.lhs, elem.rhs];
+    case "increment":
+    case "decrement":
+      return [elem.target];
+    case "call":
+      return [elem.call];
+    case "continue":
+    case "discard":
+    case "empty":
+      return [];
+    case "do":
+      return [elem.name, ...elem.params, elem.body];
+    // import/directive hold no child elems in typed fields (their payload isn't
+    // an AbstractElem), but they carry attributes that childElems prepends.
+    case "import":
+    case "directive":
+    // leaf kinds: no child elems and no attributes.
+    case "attribute":
+    case "name":
+    case "literal":
+    case "ref":
+    case "decl":
+    case "synthetic":
+      return [];
+    default:
+      assertUnreachableSilent(elem);
+      return [];
+  }
+}
+
+function isDefined<T>(value: T | undefined): value is T {
+  return value !== undefined;
+}
+
+/** Drop the `"default"` sentinel, keeping only real case-selector expressions. */
+function exprSelectors(
+  selectors: (ExpressionElem | "default")[],
+): ExpressionElem[] {
+  return selectors.filter((s): s is ExpressionElem => s !== "default");
+}