Add moduleSearchPaths to preprocessor config (#43)

This update adds a new configuration property,
`wgsl.preprocessor.moduleSearchPaths`. This represents an optional array
of additional directories to consider when trying to resolve module
imports.

Currently, adding an external directory has the unfortunate side effect
of reporting errors and warnings from every WGSL document in that
directory, regardless of whether they're pertinent to documents in the
user's workspace. In VS Code, those diagnostics are eventually pruned,
but it takes some time. I don't see this as a huge issue for now, but
this can be investigated in the future if it proves to be annoying or
problematic for users.

packages/client/package.json

@@ -99,6 +99,15 @@
 						}
 					},
 					"default": {}
+				},
+				"wgsl.preprocessor.moduleSearchPaths": {
+					"title": "Module Search Paths",
+					"description": "A list of additional directories to search for WGSL modules.",
+					"type": "array",
+					"items": {
+						"type": "string"
+					},
+					"default": []
 				}
 			}
 		}

packages/client/src/app/configuration.ts

@@ -5,4 +5,5 @@ export interface Configuration {
 export interface Preprocessor {
 	globalShaderDefs: Record<string, string>;
 	shaderDefs: Record<string, Record<string, string>>;
+	moduleSearchPaths: string[];
 }

packages/server/src/config.rs

@@ -31,6 +31,7 @@ pub struct Config {
 pub struct PreprocessorConfig {
 	pub global_shader_defs: HashMap<String, String>,
 	pub shader_defs: HashMap<String, HashMap<String, String>>,
+	pub module_search_paths: Vec<String>,
 }
 
 pub struct ConfigPlugin;

packages/server/src/documents/ecs_types.rs

@@ -26,6 +26,17 @@ pub struct PendingReads(pub HashSet<Url>);
 #[derive(Resource, Clone, Debug, Default, Deref, DerefMut)]
 pub struct PendingDocs(pub HashMap<Url, SourceStr>);
 
+/// Documents that are currently "open" in the client.
+///
+/// Note that this does not _necessarily_ represent documents that the user
+/// actually has open in their editor, because we request that the client
+/// synthetically "open" all documents in the workspace (and configured module
+/// search paths) so that the server can process them. VS Code eventually sends
+/// a "close" notification for these not-actually-open documents, but I'm not
+/// sure about the behavior of other clients.
+#[derive(Resource, Clone, Debug, Default, Deref, DerefMut)]
+pub struct OpenDocuments(pub HashSet<Url>);
+
 // --- Components --------------------------------------------------------------
 
 #[derive(Component, Clone, Debug, Deref, DerefMut, Hash)]

packages/server/src/documents/mod.rs

@@ -13,7 +13,10 @@ use bevy_ecs::{
 use bevy_utils::HashMap;
 use gramatika::{ParseStreamer, Substr, Token as _};
 use lsp_server::{Message, Notification};
-use lsp_types::{notification::Notification as _, Url};
+use lsp_types::{
+	notification::{Notification as _, ShowMessage},
+	MessageType, ShowMessageParams, Url,
+};
 use parser::{
 	decl::{Decl, ImportPathDecl},
 	ParseResult, ParseStream, SyntaxTree, TokenKind,
@@ -43,16 +46,20 @@ impl Plugin for DocumentsPlugin {
 		app.insert_resource(DocumentsMap::default())
 			.insert_resource(ModulePathsMap::default())
 			.insert_resource(PendingReads::default())
-			.insert_resource(PendingDocs::default());
+			.insert_resource(PendingDocs::default())
+			.insert_resource(OpenDocuments::default());
 
 		app.add_systems(Startup, request_all_docs_in_workspace.map(utils::log_error));
 		app.add_systems(
 			Update,
 			(
 				(
+					handle_document_closes,
 					read_new_documents,
 					update_existing_documents.map(utils::log_error),
-					handle_config_changes.run_if(resource_exists_and_changed::<Config>),
+					handle_config_changes
+						.map(utils::log_error)
+						.run_if(resource_exists_and_changed::<Config>),
 				),
 				process_pending_documents.run_if(resource_exists::<Config>),
 			)
@@ -98,6 +105,7 @@ fn request_all_docs_in_workspace(
 
 fn read_new_documents(
 	mut e_doc_reads: ResMut<Events<notify::DocumentOpen>>,
+	mut r_open_documents: ResMut<OpenDocuments>,
 	mut r_pending_reads: ResMut<PendingReads>,
 	mut r_pending_docs: ResMut<PendingDocs>,
 ) {
@@ -106,6 +114,8 @@ fn read_new_documents(
 		let uri = event.0.text_document.uri;
 		log::info!("  reading document \"{uri}\"");
 
+		r_open_documents.insert(uri.clone());
+
 		let canon_uri = uri.clone().into_canonical();
 		if r_pending_reads.contains(&canon_uri) {
 			r_pending_reads.remove(&canon_uri);
@@ -153,16 +163,42 @@ fn update_existing_documents(
 }
 
 fn handle_config_changes(
+	r_config: Res<Config>,
+	r_ipc: Res<Ipc>,
+	mut r_pending_reads: ResMut<PendingReads>,
 	mut r_pending_docs: ResMut<PendingDocs>,
 	q_docs: Query<(&DocumentUri, &WgslSource)>,
-) {
+) -> anyhow::Result<()> {
 	log::info!("  Handling configuration change");
 
 	for (uri, source) in q_docs.iter() {
 		if !r_pending_docs.contains_key(&**uri) {
 			r_pending_docs.insert(uri.0.clone(), source.0.clone());
 		}
 	}
+
+	if !r_config.preprocessor.module_search_paths.is_empty() {
+		for path_str in r_config.preprocessor.module_search_paths.iter() {
+			let path = Path::new(path_str);
+			if !path.is_absolute() {
+				r_ipc.send(Message::Notification(Notification {
+					method: ShowMessage::METHOD.into(),
+					params: serde_json::to_value(ShowMessageParams {
+						message: format!(
+							"Module search paths must be absolute. Failed to process path: \"{path_str}\"",
+						),
+						typ: MessageType::ERROR,
+					})?,
+				}))?;
+
+				continue;
+			}
+
+			request_docs_in_folder(&r_ipc, &mut r_pending_reads, path)?;
+		}
+	}
+
+	Ok(())
 }
 
 fn process_pending_documents(
@@ -283,6 +319,18 @@ fn process_pending_documents(
 	}
 }
 
+fn handle_document_closes(
+	mut r_open_documents: ResMut<OpenDocuments>,
+	mut e_document_close: ResMut<Events<notify::DocumentClose>>,
+) {
+	for event in e_document_close.drain() {
+		let uri = &event.text_document.uri;
+
+		log::info!("  removing from open documents: \"{uri}\"");
+		r_open_documents.remove(uri);
+	}
+}
+
 // --- Helpers -----------------------------------------------------------------
 
 fn request_docs_in_folder<P: AsRef<Path>>(

packages/server/src/main.rs

@@ -53,7 +53,7 @@ fn start() -> App {
 		TypeRegistrationPlugin,
 		LogPlugin {
 			level: Level::WARN,
-			filter: "warn,server=trace".into(),
+			filter: "warn,server=debug,server::ipc=trace,server::pre=warn".into(),
 		},
 		ConfigPlugin,
 		DocumentsPlugin,

packages/server/src/testing.rs

@@ -147,6 +147,7 @@ fn handle_server_requests(
 								"VERTEX_POSITIONS": "",
 							},
 							"shaderDefs": {},
+							"moduleSearchPaths": [],
 						}
 					}]);
 

packages/server/src/validation/mod.rs

@@ -1,4 +1,4 @@
-use std::sync::Arc;
+use std::{path::PathBuf, str::FromStr, sync::Arc};
 
 use bevy_app::{App, Plugin, PostUpdate};
 use bevy_derive::{Deref, DerefMut};
@@ -21,11 +21,12 @@ use parser::utils::ToRange;
 
 use crate::{
 	documents::{
-		DocumentUri, InactiveRanges, ParseErrors, PruneErrors, WgslAst, WgslScopes, WgslSource,
-		WgslSourceMap, WgslTokens,
+		DocumentUri, InactiveRanges, OpenDocuments, ParseErrors, PruneErrors, WgslAst, WgslScopes,
+		WgslSource, WgslSourceMap, WgslTokens,
 	},
 	ipc::Ipc,
 	utils,
+	workspace::Workspace,
 };
 
 use self::{
@@ -46,6 +47,7 @@ impl Plugin for DiagnosticsPlugin {
 				PostUpdate,
 				(
 					validate_documents.in_set(ValidationSystem),
+					prune_closed_external_documents.in_set(ValidationSystem),
 					publish_diagnostics
 						.map(utils::log_error)
 						.run_if(resource_changed::<DiagnosticsCollection>),
@@ -168,6 +170,50 @@ fn validate_documents(
 	}
 }
 
+fn prune_closed_external_documents(
+	mut r_collection: ResMut<DiagnosticsCollection>,
+	r_workspace: Res<Workspace>,
+	r_open_documents: Res<OpenDocuments>,
+) {
+	let workspace_paths = match &*r_workspace {
+		Workspace::RootPath(path) => match PathBuf::from_str(path) {
+			Ok(path) => vec![path],
+			Err(_) => vec![],
+		},
+		Workspace::RootUri(uri) => match uri.to_file_path() {
+			Ok(path) => vec![path],
+			Err(_) => vec![],
+		},
+		Workspace::Folders(folders) => folders
+			.iter()
+			.filter_map(|folder| folder.uri.to_file_path().ok())
+			.collect(),
+	};
+
+	for (uri, entry) in r_collection.iter_mut() {
+		if entry.is_empty() {
+			continue;
+		}
+
+		if r_open_documents.contains(uri) {
+			continue;
+		}
+
+		let Ok(path) = uri.to_file_path() else {
+			log::error!("failed to convert URI to file path: \"{uri}\"");
+			continue;
+		};
+
+		if !workspace_paths
+			.iter()
+			.any(|folder| path.starts_with(folder))
+		{
+			log::debug!("pruning diagnostics for URI: \"{uri}\"");
+			entry.clear();
+		}
+	}
+}
+
 fn publish_diagnostics(
 	r_ipc: Res<Ipc>,
 	r_collection: Res<DiagnosticsCollection>,