add a new flag preserveCodeReference (#5)

Author: fs-eire

bin/cli.ts

@@ -295,6 +295,7 @@ Options:
                                  Available: "dynamic", "static-cpp", "static-cpp-literal"
   --ext <extension>              Template file extension (default: ".wgsl.template")
   --include-prefix, -I <prefix>  Include path prefix for generated headers
+  --preserve-code-ref            Preserve code references in generated output
   --clean, -c                    Clean output directory before building
   --debug, -d                    Enable debug output for detailed error information
   --help, -h                     Show this help message
@@ -335,7 +336,7 @@ async function main() {
       d: "debug",
     },
     string: ["input", "output", "generator", "ext", "include-prefix"],
-    boolean: ["help", "version", "clean", "debug"],
+    boolean: ["help", "version", "clean", "debug", "preserve-code-ref"],
   });
 
   const options = {
@@ -346,6 +347,7 @@ async function main() {
     generator: argv.generator,
     ext: argv.ext,
     includePrefix: argv["include-prefix"],
+    preserveCodeRef: argv["preserve-code-ref"],
     clean: argv.clean,
     debug: argv.debug,
   };
@@ -422,6 +424,9 @@ async function main() {
     if (options.includePrefix) {
       console.log(`  Include prefix: ${options.includePrefix}`);
     }
+    if (options.preserveCodeRef) {
+      console.log(`  Preserve code references: enabled`);
+    }
     if (options.clean) {
       console.log(`  Clean directory: enabled`);
     }
@@ -436,6 +441,7 @@ async function main() {
       templateExt,
       generator,
       includePathPrefix: options.includePrefix,
+      preserveCodeReference: options.preserveCodeRef,
     });
 
     // Calculate generation time

src/generator-impl.ts

@@ -12,6 +12,7 @@ import type {
   TemplateRepository,
   TemplatePass2,
   TemplatePass1,
+  GenerateOptions,
 } from "./types.js";
 
 interface FunctionCallState {
@@ -113,7 +114,7 @@ function matchNextPattern(
   return earliestMatch;
 }
 
-function generateImpl(generatorState: GeneratorState) {
+function generateImpl(generatorState: GeneratorState, options: GenerateOptions) {
   const { pass1, codeGenerator, preprocessIfStack, patterns, currentFunctionCall } = generatorState;
 
   let currentLine = generatorState.currentLine;
@@ -382,6 +383,13 @@ function generateImpl(generatorState: GeneratorState) {
       output("code", "\n");
     };
 
+    if (options.preserveCodeReference) {
+      const maxLineNumber = pass1.length;
+      const lineNumberWidth = String(maxLineNumber).length;
+      const paddedLineNumber = String(currentLine + 1).padStart(lineNumberWidth, " ");
+      output("raw", `// ${paddedLineNumber} | ${line}\n`);
+    }
+
     if (line.startsWith("#")) {
       if (line.startsWith("#use ")) {
         const uses = line
@@ -663,7 +671,8 @@ function generateImpl(generatorState: GeneratorState) {
 const generate = (
   filePath: string,
   repo: TemplateRepository<TemplatePass1>,
-  codeGenerator: CodeGenerator
+  codeGenerator: CodeGenerator,
+  options?: GenerateOptions
 ): GenerateResult => {
   const pass1 = repo.templates.get(filePath)?.pass1;
   if (!pass1) {
@@ -689,7 +698,7 @@ const generate = (
     usedVariables: new Map(),
   };
 
-  generateImpl(generatorState);
+  generateImpl(generatorState, options || {});
 
   if (generatorState.preprocessIfStack.length > 0) {
     throw new WgslTemplateGenerateError(
@@ -738,11 +747,11 @@ const generate = (
 export const generator: Generator = {
   generate,
 
-  generateDirectory(repo, codeGenerator): TemplateRepository<TemplatePass2> {
+  generateDirectory(repo, codeGenerator, options): TemplateRepository<TemplatePass2> {
     const result = new Map<string, TemplatePass2>();
 
     for (const [filePath, template] of repo.templates) {
-      const generateResult = generate(filePath, repo, codeGenerator);
+      const generateResult = generate(filePath, repo, codeGenerator, options);
       result.set(filePath, {
         filePath: template.filePath,
         generateResult,

src/index.ts

@@ -42,6 +42,16 @@ export interface BuildOptions {
    * `#include "myproject/a/b/shader.wgsl"`
    */
   includePathPrefix?: string;
+
+  /**
+   * Whether to preserve code references in the generated output.
+   *
+   * If true, comments will be added to the generated code indicating the original template file and line number.
+   * This can be useful for debugging and tracing back to the source templates.
+   *
+   * Default is false.
+   */
+  preserveCodeReference?: boolean;
 }
 
 /**
@@ -80,7 +90,9 @@ export const build = async (options: BuildOptions): Promise<BuildResult> => {
     pass1 = parser.parse(pass0);
 
     const codeGenerator = resolveCodeGenerator(options.generator);
-    pass2 = generator.generateDirectory(pass1, codeGenerator);
+    pass2 = generator.generateDirectory(pass1, codeGenerator, {
+      preserveCodeReference: options.preserveCodeReference,
+    });
 
     const files = codeGenerator.build(pass2, {
       templateExt: options.templateExt,

src/types.ts

@@ -132,6 +132,10 @@ export interface CodePattern {
 // Generator Types
 // ============================================================================
 
+export interface GenerateOptions {
+  preserveCodeReference?: boolean;
+}
+
 export interface GenerateResult {
   code: string;
   params: Map<string, NonNullable<CodePattern["paramType"]>>; // name -> type
@@ -140,11 +144,17 @@ export interface GenerateResult {
 }
 
 export interface Generator {
-  generate(filePath: string, repo: TemplateRepository<TemplatePass1>, generator: CodeGenerator): GenerateResult;
+  generate(
+    filePath: string,
+    repo: TemplateRepository<TemplatePass1>,
+    generator: CodeGenerator,
+    options?: GenerateOptions
+  ): GenerateResult;
 
   generateDirectory(
     repo: TemplateRepository<TemplatePass1>,
-    generator: CodeGenerator
+    generator: CodeGenerator,
+    options?: GenerateOptions
   ): TemplateRepository<TemplatePass2>;
 }
 

test/test-runner-generator.ts

@@ -88,7 +88,9 @@ export async function runGeneratorTest(testCase: TestCase, debug?: boolean): Pro
 
       try {
         // Generate the actual output
-        const generateResult = generator.generate(entryConfig.templatePath, parsedRepo, codeGenerator);
+        const generateResult = generator.generate(entryConfig.templatePath, parsedRepo, codeGenerator, {
+          preserveCodeReference: config.preserveCodeReference,
+        });
         const actualOutput = generateResult.code;
 
         // If we expected an error but got here without throwing, that's a failure

test/test-types.ts

@@ -46,6 +46,7 @@ export interface BuildTestConfig extends BaseTestConfig {
 // Generator test configuration
 export interface GeneratorTestConfig extends BaseTestConfig {
   type: "generator";
+  preserveCodeReference?: boolean; // Whether to preserve code references in generated output
   // Generator tests use object with template file paths as keys
   entries: Record<
     string,

test/testcases/generator-preserve-code-ref/test-config.json

@@ -0,0 +1,13 @@
+{
+  "type": "generator",
+  "description": "Basic generator test with preserveCodeReference enabled",
+  "preserveCodeReference": true,
+  "entries": {
+    "test.wgsl.template": {
+      "generators": {
+        "dynamic": {},
+        "static-cpp-literal": {}
+      }
+    }
+  }
+}

test/testcases/generator-preserve-code-ref/test.wgsl.template

@@ -0,0 +1,5 @@
+@compute @workgroup_size(128)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+    let buffer_count = 4;
+    let index = global_id.x;
+}

test/testcases/generator-preserve-code-ref/test.wgsl.template.dynamic.gen

@@ -0,0 +1,10 @@
+// 1 | @compute @workgroup_size(128)
+emit("@compute @workgroup_size(128)\n");
+// 2 | fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+emit("fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {\n");
+// 3 |     let buffer_count = 4;
+emit("    let buffer_count = 4;\n");
+// 4 |     let index = global_id.x;
+emit("    let index = global_id.x;\n");
+// 5 | }
+emit("}\n");

test/testcases/generator-preserve-code-ref/test.wgsl.template.static-cpp-literal.gen

@@ -0,0 +1,10 @@
+// 1 | @compute @workgroup_size(128)
+(*ss_ptr) << "@compute @workgroup_size(128)\n";
+// 2 | fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+(*ss_ptr) << "fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {\n";
+// 3 |     let buffer_count = 4;
+(*ss_ptr) << "    let buffer_count = 4;\n";
+// 4 |     let index = global_id.x;
+(*ss_ptr) << "    let index = global_id.x;\n";
+// 5 | }
+(*ss_ptr) << "}\n";