Complete syntax tree

tools/packages/bench/bin/bench.ts

@@ -1,7 +1,7 @@
 import { WGSLLinker } from "@use-gpu/shader";
 import fs from "fs/promises";
 import path from "node:path";
-import { link, parseWESL } from "wesl";
+import { link, ModulePath, parseWESL } from "wesl";
 import { WgslReflect } from "wgsl_reflect";
 import yargs from "yargs";
 
@@ -105,7 +105,7 @@ function runNTimes(n: number, variant: ParserVariant, file: BenchTest): void {
 interface BenchTest {
   name: string;
   /** Path to the main file */
-  mainFile: string;
+  rootModule: ModulePath;
   /** All relevant files (file paths and their contents) */
   files: Map<string, string>;
 }
@@ -115,34 +115,44 @@ async function loadAllFiles(): Promise<BenchTest[]> {
   const reduceBuffer = await loadBench(
     "reduceBuffer",
     examplesDir,
-    "./reduceBuffer.wgsl",
+    ["./reduceBuffer.wgsl"],
+    ["package", "reduceBuffer"],
+  );
+  const particle = await loadBench(
+    "particle",
+    examplesDir,
+    ["./particle.wgsl"],
+    ["package", "particle"],
   );
-  const particle = await loadBench("particle", examplesDir, "./particle.wgsl");
   const rasterize = await loadBench(
     "rasterize",
     examplesDir,
-    "./rasterize_05_fine.wgsl",
+    ["./rasterize_05_fine.wgsl"],
+    ["package", "rasterize_05_fine"],
   );
   const boat = await loadBench(
     "unity_webgpu_0000026E5689B260",
     examplesDir,
-    "./unity_webgpu_000002B8376A5020.fs.wgsl",
+    ["./unity_webgpu_000002B8376A5020.fs.wgsl"],
+    ["package", "unity_webgpu_000002B8376A5020"],
   );
   const imports_only = await loadBench(
     "imports_only",
     examplesDir,
-    "./imports_only.wgsl",
+    ["./imports_only.wgsl"],
+    ["package", "imports_only"],
   );
   const bevy_deferred_lighting = await loadBench(
     "bevy_deferred_lighting",
     "./src/examples/bevy",
-    "./bevy_generated_deferred_lighting.wgsl",
+    ["./bevy_generated_deferred_lighting.wgsl"],
+    ["package", "bevy_generated_deferred_lighting"],
   );
   const bevy_linking = await loadBench(
     "bevy_linking",
     "./src/examples/naga_oil_example",
-    "./pbr.wgsl",
     [
+      "./pbr.wgsl",
       "./clustered_forward.wgsl",
       "./mesh_bindings.wgsl",
       "./mesh_types.wgsl",
@@ -156,6 +166,7 @@ async function loadAllFiles(): Promise<BenchTest[]> {
       "./shadows.wgsl",
       "./utils.wgsl",
     ],
+    ["package", "pbr"],
   );
   return [
     reduceBuffer,
@@ -171,18 +182,17 @@ async function loadAllFiles(): Promise<BenchTest[]> {
 async function loadBench(
   name: string,
   cwd: string,
-  mainFile: string,
-  extraFiles: string[] = [],
+  filePaths: string[] = [],
+  rootModule: ModulePath,
 ): Promise<BenchTest> {
   const files = new Map<string, string>();
   const addFile = async (p: string) =>
     files.set(p, await fs.readFile(path.join(cwd, p), { encoding: "utf8" }));
 
-  await addFile(mainFile);
-  for (const path of extraFiles) {
+  for (const path of filePaths) {
     await addFile(path);
   }
-  return { name, mainFile, files };
+  return { name, rootModule, files };
 }
 
 function runOnce(parserVariant: ParserVariant, test: BenchTest): void {
@@ -193,7 +203,7 @@ function runOnce(parserVariant: ParserVariant, test: BenchTest): void {
   } else if (parserVariant === "wesl-link") {
     link({
       weslSrc: Object.fromEntries(test.files.entries()),
-      rootModuleName: test.mainFile,
+      rootModulePath: test.rootModule,
     });
   } else if (parserVariant === "wgsl_reflect") {
     for (const [path, text] of test.files) {

tools/packages/bulk-test/src/stripWgsl.ts

@@ -1,34 +1,14 @@
 import { WeslStream } from "wesl";
 
-/** Remove extra bits from WGSL for test comparisons.
- *
- * removes:
- *  . extra whitespace,
- *  . comments,
- *  . trailing commas in brackets, paren, and array containers
- */
+/** Removes extra whitespace and comments from WGSL */
 export function stripWesl(text: string): string {
   const stream = new WeslStream(text);
   const firstToken = stream.nextToken();
   if (firstToken === null) return "";
-
   let result = firstToken.text;
   while (true) {
     const token = stream.nextToken();
     if (token === null) return result;
-
-    if (token.text === ",") {
-      const nextToken = stream.nextToken();
-      const nextText = nextToken?.text;
-      if (nextText === "}" || nextText === "]" || nextText === ")") {
-        // Ignore trailing comma
-        result += " ";
-        result += nextText;
-      } else {
-        result += ", " + (nextToken?.text ?? "");
-      }
-    } else {
-      result += " " + token.text;
-    }
+    result += " " + token.text;
   }
 }

tools/packages/bulk-test/src/testWgslFiles.ts

@@ -19,13 +19,13 @@ export async function testWgslFiles(namedPaths: NamedPath[]) {
   const config = { plugins: [bindingStructsPlugin()] };
 
   namedPaths.forEach(({ name, filePath }) => {
-    const shortPath = "./" + name;
     test(name, async () => {
       const orig = await fs.readFile(filePath, { encoding: "utf8" });
       const result = await expectNoLogAsync(() => {
-        const weslSrc = { [shortPath]: orig };
-        const rootModuleName = noSuffix(name);
-        return link({ weslSrc, rootModuleName, config });
+        const shortPath = "./" + name;
+        const weslSrc = { [shortPath]: text };
+        const rootModulePath = ["package", ...noSuffix(name).split("/")];
+        return link({ weslSrc, rootModulePath, config });
       });
       expect(stripWesl(result.dest)).toBe(stripWesl(orig));
     });

tools/packages/mini-parse/src/Parser.ts

@@ -18,7 +18,8 @@ import {
   tracing,
   withTraceLogging,
 } from "./ParserTracing.js";
-import { Stream, Token, TypedToken } from "./Stream.js";
+import { Span } from "./Span.js";
+import { peekToken, Stream, Token, TypedToken } from "./Stream.js";
 
 export interface AppState<C, S> {
   /**
@@ -215,6 +216,18 @@ export class Parser<I, T> {
     return map(this, fn);
   }
 
+  /** map results to a new value.
+   */
+  mapSpanned<U>(fn: (value: T, span: Span) => U): Parser<I, U> {
+    return mapSpanned(this, fn);
+  }
+
+  /** map results to a new value, or backtracks
+   */
+  verifyMap<U>(fn: (value: T) => OptParserResult<U>): Parser<I, U> {
+    return verifyMap(this, fn);
+  }
+
   /** Queue a function that runs later, typically to collect AST elements from the parse.
    * when a commit() is parsed.
    * Collection functions are dropped with parser backtracking, so
@@ -363,6 +376,43 @@ function map<I, T, U>(p: Parser<I, T>, fn: (value: T) => U): Parser<I, U> {
   return mapParser;
 }
 
+/** return a parser that maps the current results */
+function mapSpanned<I, T, U>(
+  p: Parser<I, T>,
+  fn: (value: T, span: Span) => U,
+): Parser<I, U> {
+  const mapSpannedParser = parser(
+    `mapSpanned`,
+    function _mapSpanned(ctx: ParserContext): OptParserResult<U> {
+      const start = peekToken(ctx.stream)?.span?.[0] ?? null;
+      const result = p._run(ctx);
+      if (result === null) return null;
+      const end = ctx.stream.checkpoint();
+      return { value: fn(result.value, [start ?? end, end]) };
+    },
+  );
+
+  trackChildren(mapSpannedParser, p);
+  return mapSpannedParser;
+}
+
+function verifyMap<I, T, U>(
+  p: Parser<I, T>,
+  fn: (value: T) => OptParserResult<U>,
+): Parser<I, U> {
+  const verifyMapParser = parser(
+    `verifyMap`,
+    function _verifyMap(ctx: ParserContext): OptParserResult<U> {
+      const result = p._run(ctx);
+      if (result === null) return null;
+      return fn(result.value);
+    },
+  );
+
+  trackChildren(verifyMapParser, p);
+  return verifyMapParser;
+}
+
 type ToParserFn<I, T, X> = (results: ParserResult<T>) => Parser<I, X> | null;
 
 function toParser<I, T, O>(

tools/packages/mini-parse/src/ParserLogging.ts

@@ -1,17 +1,13 @@
 import { ParserContext } from "./Parser.js";
 import { parserLog, tracePos, tracing } from "./ParserTracing.js";
-import { SrcMap, SrcWithPath } from "./SrcMap.js";
+import { Span } from "./Span.js";
 import { log } from "./WrappedLog.js";
 
 /** log an message along with the source line and a caret indicating the error position in the line
  * @param pos is the position the source string, or if src is a SrcMap, then
  *  pos is the position in the dest (e.g. preprocessed) text
  */
-export function srcLog(
-  src: string | SrcMap,
-  pos: number | [number, number],
-  ...msgs: any[]
-): void {
+export function srcLog(src: string, pos: number | Span, ...msgs: any[]): void {
   logInternal(log, src, pos, ...msgs);
 }
 
@@ -21,7 +17,7 @@ export function quotedText(text?: string): string {
 
 /** log a message along with src line, but only if tracing is active in the current parser */
 export function srcTrace(
-  src: string | SrcMap,
+  src: string,
   pos: number | [number, number],
   ...msgs: any[]
 ): void {
@@ -38,45 +34,17 @@ export function ctxLog(ctx: ParserContext, ...msgs: any[]): void {
  */
 function logInternal(
   log: typeof console.log,
-  srcOrSrcMap: string | SrcMap,
-  destPos: number | [number, number],
+  src: string,
+  destPos: number | Span,
   ...msgs: any[]
 ): void {
-  if (typeof srcOrSrcMap === "string") {
-    logInternalSrc(log, srcOrSrcMap, destPos, ...msgs);
-    return;
-  }
-  const { src, positions } = mapSrcPositions(srcOrSrcMap, destPos);
-
-  logInternalSrc(log, src.text, positions, ...msgs);
-}
-
-interface SrcPositions {
-  positions: number | [number, number];
-  src: SrcWithPath;
-}
-
-function mapSrcPositions(
-  srcMap: SrcMap,
-  destPos: number | [number, number],
-): SrcPositions {
-  const srcPos = srcMap.mapPositions(...[destPos].flat());
-  const { src } = srcPos[0];
-
-  let positions: [number, number] | number;
-  if (srcPos[1]?.src?.path === src.path && srcPos[1]?.src?.text === src.text) {
-    positions = srcPos.map(p => p.position) as [number, number];
-  } else {
-    positions = srcPos[0].position;
-  }
-
-  return { src, positions };
+  logInternalSrc(log, src, destPos, ...msgs);
 }
 
 function logInternalSrc(
   log: typeof console.log,
   src: string,
-  pos: number | [number, number],
+  pos: number | Span,
   ...msgs: any[]
 ): void {
   log(...msgs);
@@ -120,10 +88,7 @@ interface SrcLine {
 }
 
 /** return the line in the src containing a given character postion */
-export function srcLine(
-  src: string,
-  position: number | [number, number],
-): SrcLine {
+export function srcLine(src: string, position: number | Span): SrcLine {
   let pos: number;
   let pos2: number | undefined;
   if (typeof position === "number") {

tools/packages/mini-parse/src/Span.ts

@@ -2,3 +2,7 @@
  * An range, from start (inclusive) to end (exclusive).
  */
 export type Span = readonly [number, number];
+
+export function isSpan(span: any): span is Span {
+  return Array.isArray(span) && span.length === 2;
+}