Source maps

packages/parser/src/token.rs

@@ -104,6 +104,35 @@ pub enum Token {
 	Plain(Substr, Span),
 }
 
+impl From<(TokenKind, Substr, Span)> for Token {
+	fn from((kind, substr, span): (TokenKind, Substr, Span)) -> Self {
+		match kind {
+			TokenKind::Brace => Token::Brace(substr, span),
+			TokenKind::CommentStart => Token::CommentStart(substr, span),
+			TokenKind::CommentEnd => Token::CommentEnd(substr, span),
+			TokenKind::Type => Token::Type(substr, span),
+			TokenKind::Directive => Token::Directive(substr, span),
+			TokenKind::Pragma => Token::Pragma(substr, span),
+			TokenKind::PreprocessorInsertion => Token::PreprocessorInsertion(substr, span),
+			TokenKind::Keyword => Token::Keyword(substr, span),
+			TokenKind::Ident => Token::Ident(substr, span),
+			TokenKind::FloatLiteral => Token::FloatLiteral(substr, span),
+			TokenKind::IntLiteral => Token::IntLiteral(substr, span),
+			TokenKind::Path => Token::Path(substr, span),
+			TokenKind::Operator => Token::Operator(substr, span),
+			TokenKind::Punct => Token::Punct(substr, span),
+			TokenKind::Unrecognized => Token::Unrecognized(substr, span),
+			TokenKind::Attribute => Token::Attribute(substr, span),
+			TokenKind::Function => Token::Function(substr, span),
+			TokenKind::Param => Token::Param(substr, span),
+			TokenKind::Struct => Token::Struct(substr, span),
+			TokenKind::Field => Token::Field(substr, span),
+			TokenKind::Module => Token::Module(substr, span),
+			TokenKind::Plain => Token::Plain(substr, span),
+		}
+	}
+}
+
 impl fmt::Debug for Token {
 	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 		<Self as gramatika::DebugLisp>::fmt(self, f, 0)

packages/server/src/code_lens.rs

@@ -1,4 +1,3 @@
-use anyhow::bail;
 use bevy_app::{App, Plugin, PostUpdate};
 use bevy_ecs::{
 	event::EventReader,
@@ -14,7 +13,7 @@ use lsp_types::{
 use parser::{decl::Decl, utils::ToRange, Spanned};
 
 use crate::{
-	documents::{DocumentsMap, WgslAst},
+	documents::{DocumentsMap, WgslAst, WgslSourceMap},
 	ipc::{request, HasRequestId, Ipc},
 	references::TokenReferences,
 	utils,
@@ -34,6 +33,7 @@ fn handle_requests(
 	r_docs: Res<DocumentsMap>,
 	mut er_requests: EventReader<request::CodeLens>,
 	q_docs: Query<(&WgslAst, &TokenReferences)>,
+	q_source_maps: Query<&WgslSourceMap>,
 ) -> anyhow::Result<()> {
 	for request in er_requests.read() {
 		let uri = &request.text_document.uri;
@@ -72,17 +72,12 @@ fn handle_requests(
 				}
 
 				let ident_span = decl.name().span();
-				let pos = Position {
-					line: ident_span.start.line as u32,
-					character: ident_span.start.character as u32,
-				};
-
 				let refs = token_refs
 					.get(decl.name())
 					.map(|refs| {
 						refs.iter()
 							.filter_map(|src_loc| {
-								if src_loc.uri() == uri && src_loc.span() == decl.name().span() {
+								if src_loc.uri() == uri && src_loc.span() == ident_span {
 									None
 								} else {
 									Some(Location::from(src_loc))
@@ -98,10 +93,19 @@ fn handle_requests(
 					format!("{} References", refs.len())
 				};
 
+				let mapped_ident_span = q_source_maps
+					.get(entity)
+					.ok()
+					.and_then(|source_map| source_map.to_src(ident_span))
+					.unwrap_or(ident_span);
+
 				let cmd_params = ReferenceParams {
 					text_document_position: TextDocumentPositionParams::new(
 						request.text_document.clone(),
-						pos,
+						Position {
+							line: mapped_ident_span.start.line as u32,
+							character: mapped_ident_span.start.character as u32,
+						},
 					),
 					context: ReferenceContext {
 						include_declaration: false,
@@ -110,20 +114,26 @@ fn handle_requests(
 					work_done_progress_params: Default::default(),
 				};
 
-				let range = if let Some(attributes) = decl.attributes() {
-					attributes.span().through(decl.span()).to_range()
+				let decl_span = if let Some(attributes) = decl.attributes() {
+					attributes.span().through(decl.span())
 				} else {
-					decl.span().to_range()
+					decl.span()
 				};
 
+				let mapped_decl_span = q_source_maps
+					.get(entity)
+					.ok()
+					.and_then(|source_map| source_map.to_src(decl_span))
+					.unwrap_or(decl_span);
+
 				Some(CodeLens {
 					command: Some(Command {
 						title,
 						command: "wgsl.resolveLensReferences".into(),
 						arguments: Some(vec![serde_json::to_value(cmd_params).unwrap()]),
 					}),
 					data: Some(serde_json::to_value(&refs).unwrap()),
-					range,
+					range: mapped_decl_span.to_range(),
 				})
 			})
 			.collect::<Vec<_>>();

packages/server/src/config.rs

@@ -1,9 +1,9 @@
 use anyhow::{anyhow, bail};
-use bevy_app::{App, Plugin, PreUpdate, Startup, Update};
+use bevy_app::{App, Plugin, PreUpdate, Startup};
 use bevy_ecs::{
 	event::EventReader,
 	schedule::{
-		common_conditions::{not, resource_exists, resource_exists_and_changed},
+		common_conditions::{not, resource_exists},
 		IntoSystemConfigs,
 	},
 	system::{Commands, IntoSystem, Res, ResMut, Resource},
@@ -49,10 +49,12 @@ impl Plugin for ConfigPlugin {
 					.run_if(resource_exists::<Config>),
 			),
 		);
-		app.add_systems(
-			Update,
-			log_config_changes.run_if(resource_exists_and_changed::<Config>),
-		);
+
+		// TODO: Figure out how to do proper logging with filters and such
+		// app.add_systems(
+		// 	Update,
+		// 	log_config_changes.run_if(resource_exists_and_changed::<Config>),
+		// );
 	}
 }
 
@@ -103,9 +105,9 @@ fn update_config(
 	Ok(())
 }
 
-fn log_config_changes(r_config: Res<Config>) {
-	eprintln!("  Config changed: {:#?}", &*r_config);
-}
+// fn log_config_changes(r_config: Res<Config>) {
+// 	eprintln!("  Config changed: {:#?}", &*r_config);
+// }
 
 impl TryFrom<&serde_json::Value> for Config {
 	type Error = anyhow::Error;