naga-oil directive evaluation (#38)

This update significantly expands support for [`naga-oil`] modules.

The language server now accepts a `preprocessor` key, with this
interface:

```ts
export interface Preprocessor {
  globalShaderDefs: Record<string, string>;
  shaderDefs: Record<string, Record<string, string>>;
}
```

An example configuration looks like this:

```json
{
  "wgsl.preprocessor.globalShaderDefs": {
    "AVAILABLE_STORAGE_BUFFER_BINDINGS": "8",
    "PER_OBJECT_BUFFER_BATCH_SIZE": "64",
    "MAX_CASCADES_PER_LIGHT": "8",
    "MAX_DIRECTIONAL_LIGHTS": "4",
    "MORPH_TARGETS": "",
    "SKINNED": "",
    "VERTEX_POSITIONS": "",
    "SLICE_COUNT": "3",
    "SAMPLES_PER_SLICE_SIDE": "3"
  },
  "wgsl.preprocessor.shaderDefs": {
    "bevy_core_pipeline::tonemapping": {
      "TONEMAP_METHOD_AGX": ""
    }
  }
}
```

When WGSL documents are processed, conditional directives like the
following are parsed and evaluated using the provided configuration:

```wgsl
#ifdef SKINNED
    // ...
#else
    // ...
#endif
```
```wgsl
#if AVAILABLE_STORAGE_BUFFER_BINDINGS >= 3
    // ...
#else
    // ...
#endif
```

Conditional branches that evaluate to `false` are marked inactive (with
[`DiagnosticTag.Unnecessary`]) and removed from the input that's passed
to the WGSL parser, along with all lines containing conditional
directives.

Additionally, shader-def interpolations like the following are replaced
with their configured values, with a helpful diagnostic if the server
couldn't find a definition for the identifier:

```wgsl
#ifdef PER_OBJECT_BUFFER_BATCH_SIZE
// The WGSL parser will read this type as `array<Mesh, 64u>`
@group(1) @binding(0) var<uniform> mesh: array<Mesh, #{PER_OBJECT_BUFFER_BATCH_SIZE}u>;
#else
@group(1) @binding(0) var<storage> mesh: array<Mesh>;
#endif // PER_OBJECT_BUFFER_BATCH_SIZE
```
```wgsl
@compute
@workgroup_size(8, 8, 1)
fn gtao(@Builtin(global_invocation_id) global_id: vec3<u32>) {
    // The parser will read this value as `f32(3)`
    let slice_count = f32(#SLICE_COUNT);
    // ...
}
```

Hovering over the `#{PER_OBJECT_BUFFER_BATCH_SIZE}` range in the editor
displays a tooltip like the following, reflecting the value provided in
the LSP configuration:

```wgsl
#define PER_OBJECT_BUFFER_BATCH_SIZE 64
```

Additionally, in-source `#define` statements like the example above are
handled by the pre-parser in the same way.

### Remaining Work

This update does not make any changes to the handling of module imports
-- while the server still provides features like hover tooltips,
go-to-definition and reference lookups for imported/exported symbols,
these elements are not yet expanded to valid WGSL source code before
passing to the Naga validator, so Naga will still emit a parse error and
exit early in any modules that include a `#define_import_path` or
`#import` statement.

[`naga-oil`]: https://github.com/bevyengine/naga_oil
[`DiagnosticTag.Unnecessary`]:
https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#diagnosticTag

Cargo.toml

@@ -10,7 +10,8 @@ members = [
 [workspace.dependencies]
 anyhow = "1.0"
 bitflags = "2.5"
-gramatika = "0.5"
+gramatika = { version = "0.6.0", features = ["substr-source"] }
+parking_lot = "0.12"
 
 [profile.dev-custom]
 inherits = "dev"

packages/client/package.json

@@ -79,36 +79,26 @@
 			"title": "WGSL",
 			"type": "object",
 			"properties": {
-				"wgsl.preprocessor.includeKeyword": {
-					"title": "Include keyword",
-					"type": "string",
-					"default": ""
-				},
-				"wgsl.preprocessor.includePaths": {
-					"title": "Include paths",
-					"type": "array",
-					"items": {
+				"wgsl.preprocessor.globalShaderDefs": {
+					"title": "Global Shader Defs",
+					"type": "object",
+					"additionalProperties": {
 						"type": "string"
-					}
+					},
+					"default": {}
 				},
-				"wgsl.preprocessor.includeModuleAliases": {
-					"title": "Module aliases",
-					"type": "array",
-					"items": {
-						"type": "object",
-						"properties": {
-							"name": {
-								"type": "string"
-							},
-							"path": {
-								"type": "string"
+				"wgsl.preprocessor.shaderDefs": {
+					"title": "Shader Defs",
+					"type": "object",
+					"patternProperties": {
+						"([a-zA-Z_][a-zA-Z0-9_]*)(::[a-zA-Z_][a-zA-Z0-9_]*)*": {
+							"type": "object",
+							"patternProperties": {
+								"[a-zA-Z_][a-zA-Z0-9_]*": { "type": "string" }
 							}
-						},
-						"required": [
-							"name",
-							"path"
-						]
-					}
+						}
+					},
+					"default": {}
 				}
 			}
 		}

packages/client/src/app/configuration.ts

@@ -3,12 +3,6 @@ export interface Configuration {
 }
 
 export interface Preprocessor {
-	includeKeyword: string;
-	includePaths: string[];
-	includeModuleAliases: ModuleAlias[];
-}
-
-export interface ModuleAlias {
-	name: string;
-	path: string;
+	globalShaderDefs: Record<string, string>;
+	shaderDefs: Record<string, Record<string, string>>;
 }

packages/grammar/src/lib/repository/identifiers.ts

@@ -29,15 +29,15 @@ export const identifier: TMGrammarScope = {
 			},
 		},
 		{
-			match: /\b([A-Z][_a-zA-Z0-9]*)\b/,
+			match: /\b((?:[a-z][_a-zA-Z0-9]*)|(?:[A-Z][_A-Z0-9]*))\b/,
 			captures: {
-				1: { name: "entity.name.type.struct.wgsl" },
+				1: { name: "variable.other.wgsl" },
 			},
 		},
 		{
-			match: /\b([a-z][_a-zA-Z0-9]*)\b/,
+			match: /\b([A-Z][_a-zA-Z0-9]*)\b/,
 			captures: {
-				1: { name: "variable.other.wgsl" },
+				1: { name: "entity.name.type.struct.wgsl" },
 			},
 		},
 	],

packages/grammar/src/lib/repository/preprocessor.ts

@@ -5,11 +5,20 @@ export const preprocessor: TMGrammarScope = {
 	patterns: [
 		{
 			name: "meta.preprocessor.conditional.wgsl",
-			match: regex`/(#(ifn?def))\s+(${IDENT})/`,
-			captures: {
-				1: { name: "keyword.control.directive.$2.wgsl" },
-				3: { name: "entity.name.function.preprocessor.wgsl" },
+			begin: /#(ifn?def)\b/,
+			beginCaptures: {
+				0: { name: "keyword.control.directive.$1.wgsl" },
 			},
+			end: /(?=$)/,
+			patterns: [
+				{
+					match: IDENT,
+					name: "entity.name.function.preprocessor.wgsl",
+				},
+				{
+					include: "source.wgsl",
+				},
+			],
 		},
 		{
 			name: "meta.preprocessor.conditional.wgsl",
@@ -18,6 +27,23 @@ export const preprocessor: TMGrammarScope = {
 				0: { name: "keyword.control.directive.if.wgsl" },
 			},
 			end: /(?=$)/,
+			patterns: [
+				{
+					match: IDENT,
+					name: "entity.name.function.preprocessor.wgsl",
+				},
+				{
+					include: "source.wgsl",
+				},
+			],
+		},
+		{
+			name: "meta.preprocessor.define.wgsl",
+			begin: /#define\b/,
+			beginCaptures: {
+				0: { name: "keyword.control.directive.define.wgsl" },
+			},
+			end: /(?=$)/,
 			patterns: [
 				{
 					match: /[_A-Z][_A-Z0-9]*/,
@@ -90,5 +116,13 @@ export const preprocessor: TMGrammarScope = {
 				},
 			],
 		},
+		{
+			name: "meta.preprocessor.interpolation.wgsl",
+			match: regex`/(#)(${IDENT})/`,
+			captures: {
+				1: { name: "punctuation.definition.template-expression.wgsl" },
+				2: { name: "entity.name.function.preprocessor.wgsl" },
+			},
+		},
 	],
 };

packages/parser/Cargo.toml

@@ -12,13 +12,17 @@ exclude = [
 	"test-files",
 ]
 
+[features]
+default = []
+preprocessing = []
+
 [dependencies]
 bitflags = { workspace = true }
 gramatika = { workspace = true }
 itertools = "0.12"
 lazy_static = "1.4"
 lsp-types = "0.95"
-parking_lot = "0.12"
+parking_lot = { workspace = true }
 
 [dev-dependencies]
 parking_lot = "0.12"

packages/parser/project.json

@@ -8,7 +8,8 @@
 		"test": {
 			"executor": "@wgsl/tools:cargo-snapshot-test",
 			"options": {
-				"package": "wgsl-parser"
+				"package": "wgsl-parser",
+				"features": "preprocessing"
 			},
 			"configurations": {
 				"ci": {

packages/parser/snapshots/tests/e2e/bevy_mesh_expanded-001.wgsl-ast

@@ -4324,7 +4324,7 @@ tests::e2e::bevy_mesh_expanded / #1
 (Comment::Line (LineComment
   start: `//` (CommentStart (174:1..174:3)),
   text: (Text
-    `In an ideal world, the offsets are established dynamically and passed as` (174:4..174:76)
+    `In an ideal world, the offsets are established dynamically and passed as #defines` (174:4..174:85)
   ),
 ))
 (Comment::Line (LineComment
@@ -4336,7 +4336,7 @@ tests::e2e::bevy_mesh_expanded / #1
 (Comment::Line (LineComment
   start: `//` (CommentStart (199:5..199:7)),
   text: (Text
-    `Due to https://gpuweb.github.io/gpuweb/wgsl/` (199:8..199:52)
+    `Due to https://gpuweb.github.io/gpuweb/wgsl/#texel-formats` (199:8..199:66)
   ),
 ))
 (Comment::Line (LineComment
@@ -5391,7 +5391,8 @@ tests::e2e::bevy_mesh_expanded / #1
 `or` (Plain (138:31..138:33))
 `1.0` (Plain (138:34..138:37))
 `to` (Plain (138:38..138:40))
-`-1.0` (Plain (138:41..138:45))
+`-` (Plain (138:41..138:42))
+`1.0` (Plain (138:42..138:45))
 `or` (Plain (138:46..138:48))
 `1.0` (Plain (138:49..138:52))
 `respectively` (Plain (138:53..138:65))
@@ -5732,6 +5733,7 @@ tests::e2e::bevy_mesh_expanded / #1
 `and` (Plain (174:63..174:66))
 `passed` (Plain (174:67..174:73))
 `as` (Plain (174:74..174:76))
+`#defines` (Plain (174:77..174:85))
 `//` (CommentStart (175:1..175:3))
 `to` (Plain (175:4..175:6))
 `the` (Plain (175:7..175:10))
@@ -5933,6 +5935,9 @@ tests::e2e::bevy_mesh_expanded / #1
 `/` (Plain (199:46..199:47))
 `wgsl` (Plain (199:47..199:51))
 `/` (Plain (199:51..199:52))
+`#texel` (Plain (199:52..199:58))
+`-` (Plain (199:58..199:59))
+`formats` (Plain (199:59..199:66))
 `//` (CommentStart (200:5..200:7))
 `While` (Plain (200:8..200:13))
 `the` (Plain (200:14..200:17))