From dfd0a4f52954ba02449c71b2bcb74447e13bf736 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Tue, 7 Apr 2026 22:02:20 -0400 Subject: [PATCH 1/4] ide: semantic syntax highlighting We're not really doing anything semantic yet, that will come in a follow up PR. With this change, we now properly highlight: ```sql select 1 and; select 2 select; ``` `and` and `select` are both column labels in these cases, not keywords. --- crates/squawk_ide/src/lib.rs | 1 + crates/squawk_ide/src/semantic_tokens.rs | 174 ++++++++++++++++++ crates/squawk_parser/src/grammar.rs | 2 + .../tests/snapshots/tests__select_ok.snap | 20 +- crates/squawk_server/src/global_state.rs | 8 +- crates/squawk_server/src/handlers.rs | 2 + .../src/handlers/semantic_tokens.rs | 44 +++++ crates/squawk_server/src/lib.rs | 1 + crates/squawk_server/src/lsp_utils.rs | 96 +++++++++- crates/squawk_server/src/semantic_tokens.rs | 26 +++ crates/squawk_server/src/server.rs | 22 ++- crates/xtask/src/codegen.rs | 2 +- squawk-vscode/syntaxes/pgsql.tmLanguage.json | 2 +- 13 files changed, 383 insertions(+), 17 deletions(-) create mode 100644 crates/squawk_ide/src/semantic_tokens.rs create mode 100644 crates/squawk_server/src/handlers/semantic_tokens.rs create mode 100644 crates/squawk_server/src/semantic_tokens.rs diff --git a/crates/squawk_ide/src/lib.rs b/crates/squawk_ide/src/lib.rs index 2c2302ad..dc6287bd 100644 --- a/crates/squawk_ide/src/lib.rs +++ b/crates/squawk_ide/src/lib.rs @@ -18,6 +18,7 @@ mod offsets; mod quote; mod resolve; mod scope; +pub mod semantic_tokens; mod symbols; #[cfg(test)] pub mod test_utils; diff --git a/crates/squawk_ide/src/semantic_tokens.rs b/crates/squawk_ide/src/semantic_tokens.rs new file mode 100644 index 00000000..c7698ea5 --- /dev/null +++ b/crates/squawk_ide/src/semantic_tokens.rs @@ -0,0 +1,174 @@ +use rowan::{NodeOrToken, TextRange}; +use salsa::Database as Db; +use squawk_syntax::{ + SyntaxKind, + ast::{self, AstNode}, +}; + +use crate::db::{File, parse}; + +/// A semantic token with its position and classification. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SemanticToken { + pub range: TextRange, + pub token_type: SemanticTokenType, + pub modifiers: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[repr(u8)] +pub enum SemanticTokenModifier { + Definition = 0, + Readonly, + Documentation, +} + +/// Semantic token types supported by the language server. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum SemanticTokenType { + Keyword, + String, + Bool, + Number, + Function, + Operator, + Punctuation, + Name, + NameRef, + Comment, + Type, + Parameter, + PositionalParam, +} + +#[salsa::tracked] +pub fn semantic_tokens( + db: &dyn Db, + file: File, + range_to_highlight: Option, +) -> Vec { + let parse = parse(db, file); + let tree = parse.tree(); + let root = tree.syntax(); + + // Determine the root based on the given range. + let (root, range_to_highlight) = { + let source_file = root; + match range_to_highlight { + Some(range) => { + let node = match source_file.covering_element(range) { + NodeOrToken::Node(it) => it, + NodeOrToken::Token(it) => it.parent().unwrap_or_else(|| source_file.clone()), + }; + (node, range) + } + None => (source_file.clone(), source_file.text_range()), + } + }; + + let mut out = vec![]; + + // Taken from: https://github.com/rust-lang/rust-analyzer/blob/2efc80078029894eec0699f62ec8d5c1a56af763/crates/ide/src/syntax_highlighting.rs#L267C21-L267C21 + let preorder = root.preorder_with_tokens(); + for event in preorder { + use rowan::WalkEvent::{Enter, Leave}; + + let range = match &event { + Enter(it) | Leave(it) => it.text_range(), + }; + + // Element outside of the viewport, no need to highlight + if range_to_highlight.intersect(range).is_none() { + continue; + } + + match event { + Enter(NodeOrToken::Node(node)) => { + if let Some(target) = ast::Target::cast(node) { + if let Some(as_name) = target.as_name() { + if let Some(name) = as_name.name() { + let range = name.syntax().text_range(); + out.push(SemanticToken { + range, + token_type: SemanticTokenType::Name, + modifiers: None, + }); + } + } + }; + } + Enter(NodeOrToken::Token(token)) => { + if token.kind() == SyntaxKind::WHITESPACE { + continue; + } + if token.kind() == SyntaxKind::POSITIONAL_PARAM { + out.push(SemanticToken { + range: token.text_range(), + token_type: SemanticTokenType::PositionalParam, + modifiers: None, + }) + } + } + Leave(_) => {} + } + } + out +} + +#[cfg(test)] +mod test { + use crate::db::{Database, File}; + use insta::assert_snapshot; + use std::fmt::Write; + + fn semantic_tokens(sql: &str) -> String { + let db = Database::default(); + let file = File::new(&db, sql.to_string().into()); + let tokens = super::semantic_tokens(&db, file, None); + + let mut result = String::new(); + for token in tokens { + let start: usize = token.range.start().into(); + let end: usize = token.range.end().into(); + let token_text = &sql[start..end]; + // TODO: + let modifiers_text = ""; + writeln!( + result, + "{:?} @ {}..{}: {:?}{}", + token_text, start, end, token.token_type, modifiers_text + ) + .unwrap(); + } + result + } + + #[test] + fn create_function() { + assert_snapshot!(semantic_tokens(" +create function add(a int, b int) returns int + as 'select $1 + $2' + language sql; +"), @""); + } + + #[test] + fn select_keywords() { + assert_snapshot!(semantic_tokens(" +select 1 and, 2 select; +"), @r#" + "and" @ 10..13: Name + "select" @ 17..23: Name + "#) + } + + #[test] + fn positional_param() { + assert_snapshot!(semantic_tokens(" +select $1, $2; +"), @r#" + "$1" @ 8..10: PositionalParam + "$2" @ 12..14: PositionalParam + "#) + } +} diff --git a/crates/squawk_parser/src/grammar.rs b/crates/squawk_parser/src/grammar.rs index ea3f65a2..e821c5b2 100644 --- a/crates/squawk_parser/src/grammar.rs +++ b/crates/squawk_parser/src/grammar.rs @@ -2511,7 +2511,9 @@ fn expr_bp(p: &mut Parser<'_>, bp: u8, r: &Restrictions) -> Option; @@ -230,6 +230,8 @@ impl GlobalState { .on::(handle_syntax_tree) .on::(handle_tokens) .on::(handle_references) + .on::(handle_semantic_tokens_full) + .on::(handle_semantic_tokens_range) .finish(); } } diff --git a/crates/squawk_server/src/handlers.rs b/crates/squawk_server/src/handlers.rs index 3f0f8998..9c4af5e6 100644 --- a/crates/squawk_server/src/handlers.rs +++ b/crates/squawk_server/src/handlers.rs @@ -9,6 +9,7 @@ mod inlay_hints; mod notifications; mod references; mod selection_range; +mod semantic_tokens; mod shutdown; mod syntax_tree; mod tokens; @@ -26,6 +27,7 @@ pub(crate) use notifications::{ }; pub(crate) use references::handle_references; pub(crate) use selection_range::handle_selection_range; +pub(crate) use semantic_tokens::{handle_semantic_tokens_full, handle_semantic_tokens_range}; pub(crate) use shutdown::handle_shutdown; pub(crate) use syntax_tree::{SyntaxTreeRequest, handle_syntax_tree}; pub(crate) use tokens::{TokensRequest, handle_tokens}; diff --git a/crates/squawk_server/src/handlers/semantic_tokens.rs b/crates/squawk_server/src/handlers/semantic_tokens.rs new file mode 100644 index 00000000..c61c0d63 --- /dev/null +++ b/crates/squawk_server/src/handlers/semantic_tokens.rs @@ -0,0 +1,44 @@ +use anyhow::Result; +use lsp_types::{ + SemanticTokens, SemanticTokensParams, SemanticTokensRangeParams, SemanticTokensRangeResult, + SemanticTokensResult, +}; + +use squawk_ide::db::line_index; +use squawk_ide::semantic_tokens::semantic_tokens; + +use crate::global_state::Snapshot; +use crate::lsp_utils; + +pub(crate) fn handle_semantic_tokens_full( + snapshot: &Snapshot, + params: SemanticTokensParams, +) -> Result> { + let uri = params.text_document.uri; + let db = snapshot.db(); + let file = snapshot.file(&uri).unwrap(); + let line_index = line_index(db, file); + let text = file.content(db); + let tokens = semantic_tokens(db, file, None); + Ok(Some(SemanticTokensResult::Tokens(SemanticTokens { + result_id: None, + data: lsp_utils::to_semantic_tokens(text, line_index, tokens), + }))) +} + +pub(crate) fn handle_semantic_tokens_range( + snapshot: &Snapshot, + params: SemanticTokensRangeParams, +) -> Result> { + let uri = params.text_document.uri; + let db = snapshot.db(); + let file = snapshot.file(&uri).unwrap(); + let line_index = line_index(db, file); + let range_to_highlight = lsp_utils::text_range(&line_index, params.range); + let tokens = semantic_tokens(db, file, range_to_highlight); + let text = file.content(db); + Ok(Some(SemanticTokensRangeResult::Tokens(SemanticTokens { + result_id: None, + data: lsp_utils::to_semantic_tokens(text, line_index, tokens), + }))) +} diff --git a/crates/squawk_server/src/lib.rs b/crates/squawk_server/src/lib.rs index a28fbc47..e57a3da4 100644 --- a/crates/squawk_server/src/lib.rs +++ b/crates/squawk_server/src/lib.rs @@ -6,6 +6,7 @@ mod ignore; mod lint; mod lsp_utils; mod panic; +mod semantic_tokens; mod server; pub use server::run; diff --git a/crates/squawk_server/src/lsp_utils.rs b/crates/squawk_server/src/lsp_utils.rs index 6c8a42ca..88226300 100644 --- a/crates/squawk_server/src/lsp_utils.rs +++ b/crates/squawk_server/src/lsp_utils.rs @@ -6,16 +6,18 @@ use ::line_index::{LineIndex, TextRange, TextSize}; use log::warn; use lsp_types::{ CodeAction, CodeActionKind, FoldingRange, FoldingRangeKind as LspFoldingRangeKind, Location, - Url, WorkspaceEdit, + SemanticToken, Url, WorkspaceEdit, }; use squawk_ide::builtins::{builtins_line_index, builtins_url}; use squawk_ide::code_actions::ActionKind; use squawk_ide::db::line_index; use squawk_ide::folding_ranges::{Fold, FoldKind}; +use squawk_ide::semantic_tokens::{SemanticTokenModifier, SemanticTokenType}; use crate::global_state::Snapshot; +use crate::semantic_tokens; -fn text_range(index: &LineIndex, range: lsp_types::Range) -> Option { +pub(crate) fn text_range(index: &LineIndex, range: lsp_types::Range) -> Option { let start = offset(index, range.start)?; let end = offset(index, range.end)?; if end >= start { @@ -229,6 +231,96 @@ pub(crate) fn to_location( Some(Location { uri, range }) } +pub(crate) fn to_semantic_tokens( + text: &str, + line_index: LineIndex, + semantic_tokens: Vec, +) -> Vec { + let mut encoder = Encoder { + tokens: Vec::with_capacity(semantic_tokens.len()), + prev_line: 0, + prev_start: 0, + }; + + for token in &*semantic_tokens { + // Taken from rust-analyzer, this solves the case where we have a multi + // line semantic token which isn't supported by the LSP spec. + // see: https://github.com/rust-lang/rust-analyzer/blob/2efc80078029894eec0699f62ec8d5c1a56af763/crates/rust-analyzer/src/lsp/to_proto.rs#L781C28-L781C28 + for mut text_range in line_index.lines(token.range) { + if text[text_range].ends_with('\n') { + text_range = + TextRange::new(text_range.start(), text_range.end() - TextSize::of('\n')); + } + let lsp_range = range(&line_index, text_range); + let len = lsp_range.end.character - lsp_range.start.character; + encoder.push_token_at(lsp_range.start, len, token.token_type, token.modifiers); + } + } + + encoder.tokens +} + +// Taken from Ty +// see: https://github.com/charliermarsh/ruff/blob/5011b253c1aca2a1906762cf45414d0fda1a088a/crates/ty_server/src/server/api/semantic_tokens.rs#L23C25-L23C25 +struct Encoder { + tokens: Vec, + prev_line: u32, + prev_start: u32, +} + +impl Encoder { + fn push_token_at( + &mut self, + start: lsp_types::Position, + length: u32, + ty: SemanticTokenType, + _modifiers: Option, + ) { + // LSP semantic tokens are encoded as deltas + let delta_line = start.line - self.prev_line; + let delta_start = if delta_line == 0 { + start.character - self.prev_start + } else { + start.character + }; + + let token_type = to_token_type(ty); + + let token_index = semantic_tokens::type_index(token_type); + + self.tokens.push(SemanticToken { + delta_line, + delta_start, + length, + token_type: token_index, + // TODO: once we get modifiers going, we'll need to update this + token_modifiers_bitset: 0, + }); + + self.prev_line = start.line; + self.prev_start = start.character; + } +} + +fn to_token_type(ty: SemanticTokenType) -> lsp_types::SemanticTokenType { + match ty { + SemanticTokenType::Keyword => lsp_types::SemanticTokenType::KEYWORD, + SemanticTokenType::String => lsp_types::SemanticTokenType::STRING, + SemanticTokenType::Bool => lsp_types::SemanticTokenType::KEYWORD, + SemanticTokenType::Number => lsp_types::SemanticTokenType::NUMBER, + SemanticTokenType::Function => lsp_types::SemanticTokenType::FUNCTION, + SemanticTokenType::Operator => lsp_types::SemanticTokenType::OPERATOR, + SemanticTokenType::Punctuation => lsp_types::SemanticTokenType::OPERATOR, + SemanticTokenType::Name => lsp_types::SemanticTokenType::VARIABLE, + SemanticTokenType::NameRef => lsp_types::SemanticTokenType::VARIABLE, + SemanticTokenType::Comment => lsp_types::SemanticTokenType::COMMENT, + SemanticTokenType::Type => lsp_types::SemanticTokenType::TYPE, + SemanticTokenType::PositionalParam | SemanticTokenType::Parameter => { + lsp_types::SemanticTokenType::PARAMETER + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/squawk_server/src/semantic_tokens.rs b/crates/squawk_server/src/semantic_tokens.rs new file mode 100644 index 00000000..4e888d24 --- /dev/null +++ b/crates/squawk_server/src/semantic_tokens.rs @@ -0,0 +1,26 @@ +use lsp_types::{SemanticTokenModifier, SemanticTokenType}; + +pub(crate) const SUPPORTED_TYPES: &[SemanticTokenType] = &[ + SemanticTokenType::COMMENT, + SemanticTokenType::FUNCTION, + SemanticTokenType::KEYWORD, + SemanticTokenType::NAMESPACE, + SemanticTokenType::NUMBER, + SemanticTokenType::OPERATOR, + SemanticTokenType::PARAMETER, + SemanticTokenType::PROPERTY, + SemanticTokenType::STRING, + SemanticTokenType::STRUCT, + SemanticTokenType::TYPE, + SemanticTokenType::VARIABLE, +]; + +pub(crate) const SUPPORTED_MODIFIERS: &[SemanticTokenModifier] = &[ + SemanticTokenModifier::DECLARATION, + SemanticTokenModifier::DEFINITION, + SemanticTokenModifier::READONLY, +]; + +pub(crate) fn type_index(ty: SemanticTokenType) -> u32 { + SUPPORTED_TYPES.iter().position(|it| *it == ty).unwrap() as u32 +} diff --git a/crates/squawk_server/src/server.rs b/crates/squawk_server/src/server.rs index 4fcb1ccf..1f846aee 100644 --- a/crates/squawk_server/src/server.rs +++ b/crates/squawk_server/src/server.rs @@ -5,10 +5,15 @@ use lsp_types::{ CodeActionKind, CodeActionOptions, CodeActionProviderCapability, CompletionOptions, DiagnosticOptions, DiagnosticServerCapabilities, FoldingRangeProviderCapability, HoverProviderCapability, InitializeParams, OneOf, SelectionRangeProviderCapability, - ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind, WorkDoneProgressOptions, + SemanticTokensFullOptions, SemanticTokensLegend, SemanticTokensOptions, + SemanticTokensServerCapabilities, ServerCapabilities, TextDocumentSyncCapability, + TextDocumentSyncKind, WorkDoneProgressOptions, }; -use crate::global_state::GlobalState; +use crate::{ + global_state::GlobalState, + semantic_tokens::{SUPPORTED_MODIFIERS, SUPPORTED_TYPES}, +}; pub fn run() -> Result<()> { info!("Starting Squawk LSP server"); @@ -53,6 +58,19 @@ pub fn run() -> Result<()> { }, completion_item: None, }), + semantic_tokens_provider: Some(SemanticTokensServerCapabilities::SemanticTokensOptions( + SemanticTokensOptions { + work_done_progress_options: WorkDoneProgressOptions { + work_done_progress: None, + }, + legend: SemanticTokensLegend { + token_types: SUPPORTED_TYPES.to_vec(), + token_modifiers: SUPPORTED_MODIFIERS.to_vec(), + }, + range: Some(true), + full: Some(SemanticTokensFullOptions::Bool(true)), + }, + )), ..Default::default() }) .unwrap(); diff --git a/crates/xtask/src/codegen.rs b/crates/xtask/src/codegen.rs index 5ad31154..0e806000 100644 --- a/crates/xtask/src/codegen.rs +++ b/crates/xtask/src/codegen.rs @@ -902,7 +902,7 @@ fn generate_nodes(nodes: &[AstNodeSrc], enums: &[AstEnumSrc]) -> String { // Multi-word keyword phrases that should be highlighted as keywords, not // operators. -const KEYWORD_PHRASES: &[&str] = &["if not exists", "if exists"]; +const KEYWORD_PHRASES: &[&str] = &["if not exists", "if exists", "or replace"]; // Multi-word entries must come before their single-word components so the // regex engine matches the longest form first. diff --git a/squawk-vscode/syntaxes/pgsql.tmLanguage.json b/squawk-vscode/syntaxes/pgsql.tmLanguage.json index 2c1ec922..3ba3bd19 100644 --- a/squawk-vscode/syntaxes/pgsql.tmLanguage.json +++ b/squawk-vscode/syntaxes/pgsql.tmLanguage.json @@ -146,7 +146,7 @@ "keywords": { "patterns": [ { - "match": "(?i)\\b(if\\s+not\\s+exists|if\\s+exists)\\b", + "match": "(?i)\\b(if\\s+not\\s+exists|if\\s+exists|or\\s+replace)\\b", "name": "keyword.other.pgsql" }, { From ff1404cb378c58b2592481fdefc7a58ad2d1e23f Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Tue, 7 Apr 2026 22:03:38 -0400 Subject: [PATCH 2/4] comment --- crates/squawk_ide/src/semantic_tokens.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/squawk_ide/src/semantic_tokens.rs b/crates/squawk_ide/src/semantic_tokens.rs index c7698ea5..fcbf932b 100644 --- a/crates/squawk_ide/src/semantic_tokens.rs +++ b/crates/squawk_ide/src/semantic_tokens.rs @@ -131,7 +131,7 @@ mod test { let start: usize = token.range.start().into(); let end: usize = token.range.end().into(); let token_text = &sql[start..end]; - // TODO: + // TODO: once we get modfifiers, we'll need to update this let modifiers_text = ""; writeln!( result, From 371b7f091583065b6207c2b219c9ee83d835bf3a Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Tue, 7 Apr 2026 22:08:07 -0400 Subject: [PATCH 3/4] fix lint --- crates/xtask/src/codegen.rs | 2 +- squawk-vscode/syntaxes/pgsql.tmLanguage.json | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/xtask/src/codegen.rs b/crates/xtask/src/codegen.rs index 0e806000..2f78d3bc 100644 --- a/crates/xtask/src/codegen.rs +++ b/crates/xtask/src/codegen.rs @@ -902,7 +902,7 @@ fn generate_nodes(nodes: &[AstNodeSrc], enums: &[AstEnumSrc]) -> String { // Multi-word keyword phrases that should be highlighted as keywords, not // operators. -const KEYWORD_PHRASES: &[&str] = &["if not exists", "if exists", "or replace"]; +const KEYWORD_PHRASES: &[&str] = &["if not exists", "or replace", "if exists"]; // Multi-word entries must come before their single-word components so the // regex engine matches the longest form first. diff --git a/squawk-vscode/syntaxes/pgsql.tmLanguage.json b/squawk-vscode/syntaxes/pgsql.tmLanguage.json index 3ba3bd19..3661d799 100644 --- a/squawk-vscode/syntaxes/pgsql.tmLanguage.json +++ b/squawk-vscode/syntaxes/pgsql.tmLanguage.json @@ -146,7 +146,7 @@ "keywords": { "patterns": [ { - "match": "(?i)\\b(if\\s+not\\s+exists|if\\s+exists|or\\s+replace)\\b", + "match": "(?i)\\b(if\\s+not\\s+exists|or\\s+replace|if\\s+exists)\\b", "name": "keyword.other.pgsql" }, { From e2a97eec3f21d6e7deecabbbb9ec24262ea5865e Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Tue, 7 Apr 2026 22:14:49 -0400 Subject: [PATCH 4/4] cleanup --- crates/squawk_ide/src/semantic_tokens.rs | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/crates/squawk_ide/src/semantic_tokens.rs b/crates/squawk_ide/src/semantic_tokens.rs index fcbf932b..63ad123f 100644 --- a/crates/squawk_ide/src/semantic_tokens.rs +++ b/crates/squawk_ide/src/semantic_tokens.rs @@ -84,17 +84,16 @@ pub fn semantic_tokens( match event { Enter(NodeOrToken::Node(node)) => { - if let Some(target) = ast::Target::cast(node) { - if let Some(as_name) = target.as_name() { - if let Some(name) = as_name.name() { - let range = name.syntax().text_range(); - out.push(SemanticToken { - range, - token_type: SemanticTokenType::Name, - modifiers: None, - }); - } - } + if let Some(target) = ast::Target::cast(node) + && let Some(as_name) = target.as_name() + && let Some(name) = as_name.name() + { + let range = name.syntax().text_range(); + out.push(SemanticToken { + range, + token_type: SemanticTokenType::Name, + modifiers: None, + }); }; } Enter(NodeOrToken::Token(token)) => {