From 119ac410fb92feb810f50231137891b7d6891015 Mon Sep 17 00:00:00 2001 From: Steve Dignam Date: Mon, 13 Apr 2026 18:38:56 -0400 Subject: [PATCH] ide: semantic token highlighting for func params & types --- crates/squawk_ide/src/semantic_tokens.rs | 292 +++++++++++++++++-- crates/squawk_parser/src/grammar.rs | 6 +- crates/squawk_syntax/src/lib.rs | 2 +- crates/squawk_syntax/src/postgresql.ungram | 3 +- crates/squawk_syntax/src/syntax_node.rs | 2 +- squawk-vscode/syntaxes/pgsql.tmLanguage.json | 2 +- 6 files changed, 278 insertions(+), 29 deletions(-) diff --git a/crates/squawk_ide/src/semantic_tokens.rs b/crates/squawk_ide/src/semantic_tokens.rs index 63ad123f5..bfc4511d6 100644 --- a/crates/squawk_ide/src/semantic_tokens.rs +++ b/crates/squawk_ide/src/semantic_tokens.rs @@ -1,12 +1,97 @@ use rowan::{NodeOrToken, TextRange}; use salsa::Database as Db; use squawk_syntax::{ - SyntaxKind, + SyntaxElement, SyntaxKind, ast::{self, AstNode}, }; use crate::db::{File, parse}; +fn highlight_param_mode(out: &mut SemanticTokenBuilder, mode: ast::ParamMode) { + match mode { + ast::ParamMode::ParamIn(param_in) => { + if let Some(token) = param_in.in_token() { + out.push_keyword(token.into()); + } + } + ast::ParamMode::ParamInOut(param_in_out) => { + if let Some(token) = param_in_out.in_token() { + out.push_keyword(token.into()); + } + if let Some(token) = param_in_out.inout_token() { + out.push_keyword(token.into()); + } + if let Some(token) = param_in_out.out_token() { + out.push_keyword(token.into()); + } + } + ast::ParamMode::ParamOut(param_out) => { + if let Some(token) = param_out.out_token() { + out.push_keyword(token.into()); + } + } + ast::ParamMode::ParamVariadic(param_variadic) => { + if let Some(token) = param_variadic.variadic_token() { + out.push_keyword(token.into()); + } + } + } +} + +fn highlight_type(out: &mut SemanticTokenBuilder, ty: ast::Type) { + match ty { + ast::Type::ArrayType(array_type) => { + if let Some(ty) = array_type.ty() { + highlight_type(out, ty); + } + } + ast::Type::BitType(bit_type) => { + if let Some(token) = bit_type.bit_token() { + out.push_type(token.into()); + } + } + ast::Type::CharType(char_type) => { + if let Some(token) = char_type + .varchar_token() + .or_else(|| char_type.nchar_token()) + .or_else(|| char_type.character_token()) + .or_else(|| char_type.char_token()) + { + out.push_type(token.into()); + }; + } + ast::Type::DoubleType(double_type) => { + if let Some(token) = double_type.double_token() { + out.push_type(token.into()); + } + } + ast::Type::ExprType(_) => (), + ast::Type::IntervalType(interval_type) => { + if let Some(token) = interval_type.interval_token() { + out.push_type(token.into()); + } + } + ast::Type::PathType(path_type) => { + if let Some(name_ref) = path_type + .path() + .and_then(|path| path.segment()) + .and_then(|ps| ps.name_ref()) + { + out.push_type(name_ref.syntax().clone().into()); + } + } + ast::Type::PercentType(_) => (), + ast::Type::TimeType(time_type) => { + if let Some(token) = time_type + .timestamp_token() + .or_else(|| time_type.time_token()) + { + out.push_type(token.into()); + } + } + } +} + /// A semantic token with its position and classification. #[derive(Debug, Clone, PartialEq, Eq)] pub struct SemanticToken { @@ -41,6 +126,35 @@ pub enum SemanticTokenType { PositionalParam, } +#[derive(Default)] +struct SemanticTokenBuilder { + tokens: Vec, +} + +impl SemanticTokenBuilder { + fn build(mut self) -> Vec { + self.tokens + .sort_by_key(|token| (token.range.start(), token.range.end())); + self.tokens + } + + fn push_keyword(&mut self, syntax_element: SyntaxElement) { + self.push_token(syntax_element, SemanticTokenType::Keyword); + } + + fn push_type(&mut self, syntax_element: SyntaxElement) { + self.push_token(syntax_element, SemanticTokenType::Type); + } + + fn push_token(&mut self, syntax_element: SyntaxElement, token_type: SemanticTokenType) { + self.tokens.push(SemanticToken { + range: syntax_element.text_range(), + token_type, + modifiers: None, + }); + } +} + #[salsa::tracked] pub fn semantic_tokens( db: &dyn Db, @@ -66,7 +180,7 @@ pub fn semantic_tokens( } }; - let mut out = vec![]; + let mut out = SemanticTokenBuilder::default(); // 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(); @@ -84,34 +198,76 @@ pub fn semantic_tokens( match event { Enter(NodeOrToken::Node(node)) => { - if let Some(target) = ast::Target::cast(node) + if let Some(target) = ast::Target::cast(node.clone()) && 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, - }); + out.push_token(name.syntax().clone().into(), SemanticTokenType::Name); }; + + if let Some(alias) = ast::Alias::cast(node.clone()) + && let Some(column_list) = alias.column_list() + { + for column in column_list.columns() { + if let Some(ty) = column.ty() { + highlight_type(&mut out, ty); + } + } + } + + if let Some(cast_expr) = ast::CastExpr::cast(node.clone()) + && let Some(ty) = cast_expr.ty() + { + highlight_type(&mut out, ty); + } + + if let Some(create_function) = ast::CreateFunction::cast(node) { + if let Some(param_list) = create_function.param_list() { + for param in param_list.params() { + if let Some(mode) = param.mode() { + highlight_param_mode(&mut out, mode); + } + if let Some(name) = param.name() { + out.push_token( + name.syntax().clone().into(), + SemanticTokenType::Parameter, + ); + } + if let Some(ty) = param.ty() { + highlight_type(&mut out, ty); + } + } + } + + if let Some(ret_type) = create_function.ret_type() { + if let Some(ty) = ret_type.ty() { + highlight_type(&mut out, ty); + } + if let Some(table_arg_list) = ret_type.table_arg_list() { + for arg in table_arg_list.args() { + if let ast::TableArg::Column(column) = arg + && let Some(ty) = column.ty() + { + highlight_type(&mut out, ty); + } + } + } + } + } } 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, - }) + out.push_token(token.into(), SemanticTokenType::PositionalParam); } } Leave(_) => {} } } - out + + out.build() } #[cfg(test)] @@ -143,12 +299,61 @@ mod test { } #[test] - fn create_function() { - assert_snapshot!(semantic_tokens(" -create function add(a int, b int) returns int - as 'select $1 + $2' - language sql; -"), @""); + fn create_function_misc_params() { + assert_snapshot!(semantic_tokens( + " +create function add( + in a int = 1, + inout b text default 'x', + in out c varchar(10)[], + variadic d int[] +) returns int +as 'select $1 + $2' +language sql; +", + ), @r#" + "in" @ 24..26: Keyword + "a" @ 27..28: Parameter + "int" @ 29..32: Type + "inout" @ 40..45: Keyword + "b" @ 46..47: Parameter + "text" @ 48..52: Type + "in" @ 68..70: Keyword + "out" @ 71..74: Keyword + "c" @ 75..76: Parameter + "varchar" @ 77..84: Type + "variadic" @ 94..102: Keyword + "d" @ 103..104: Parameter + "int" @ 105..108: Type + "int" @ 121..124: Type + "#); + } + + #[test] + fn create_function_param_mode_type() { + assert_snapshot!(semantic_tokens( + " +create function f(int8 in int8) +returns void +as '' language sql; +", + ), @r#" + "int8" @ 19..23: Parameter + "in" @ 24..26: Keyword + "int8" @ 27..31: Type + "void" @ 41..45: Type + "#); + } + + #[test] + fn create_function_percent_type() { + assert_snapshot!(semantic_tokens( + " +create function f(a t.c%type) +returns t.b%type +as '' language plpgsql; +", + ), @r#""a" @ 19..20: Parameter"#); } #[test] @@ -170,4 +375,49 @@ select $1, $2; "$2" @ 12..14: PositionalParam "#) } + + #[test] + fn from_alias_column_types() { + assert_snapshot!(semantic_tokens( + " +select * +from f as t(a int, b jsonb, c text, x int, ca char(5)[], ia int[][], r jbpop); +", + ), @r#" + "int" @ 24..27: Type + "jsonb" @ 31..36: Type + "text" @ 40..44: Type + "int" @ 48..51: Type + "char" @ 56..60: Type + "int" @ 70..73: Type + "jbpop" @ 81..86: Type + "#); + } + + #[test] + fn cast_types() { + assert_snapshot!(semantic_tokens( + " +select '1'::jsonb, '2'::json, cast(1 as integer), cast(1 as int4[][]), cast(1 as varchar(10)); +", + ), @r#" + "jsonb" @ 13..18: Type + "json" @ 25..29: Type + "integer" @ 41..48: Type + "int4" @ 61..65: Type + "varchar" @ 82..89: Type + "#); + } + + #[test] + fn positional_param_and_cast_type() { + assert_snapshot!(semantic_tokens( + " +select $2::jsonb; +", + ), @r#" + "$2" @ 8..10: PositionalParam + "jsonb" @ 12..17: Type + "#); + } } diff --git a/crates/squawk_parser/src/grammar.rs b/crates/squawk_parser/src/grammar.rs index e821c5b26..654eb56b6 100644 --- a/crates/squawk_parser/src/grammar.rs +++ b/crates/squawk_parser/src/grammar.rs @@ -3664,12 +3664,12 @@ fn opt_column_list_with(p: &mut Parser<'_>, kind: ColumnDefKind) -> bool { fn column(p: &mut Parser<'_>, kind: &ColumnDefKind) -> CompletedMarker { assert!(p.at_ts(COLUMN_FIRST)); let m = p.start(); - // https://www.depesz.com/2024/10/03/waiting-for-postgresql-18-add-temporal-foreign-key-contraints/ - // TODO: add validation to ensure this is in the right position - p.eat(PERIOD_KW); match kind { ColumnDefKind::Name => name(p), ColumnDefKind::NameRef => { + // https://www.depesz.com/2024/10/03/waiting-for-postgresql-18-add-temporal-foreign-key-contraints/ + // TODO: add validation to ensure this is in the right position + p.eat(PERIOD_KW); // supports parsing things like: // INSERT INTO tictactoe (game, board[1:3][1:3]) name_ref(p).map(|lhs| postfix_expr(p, lhs, true)); diff --git a/crates/squawk_syntax/src/lib.rs b/crates/squawk_syntax/src/lib.rs index cd1ab7e4d..22c6d83ac 100644 --- a/crates/squawk_syntax/src/lib.rs +++ b/crates/squawk_syntax/src/lib.rs @@ -42,7 +42,7 @@ pub use squawk_parser::SyntaxKind; use ast::AstNode; use rowan::GreenNode; use syntax_error::SyntaxError; -pub use syntax_node::{SyntaxNode, SyntaxNodePtr, SyntaxToken}; +pub use syntax_node::{SyntaxElement, SyntaxNode, SyntaxNodePtr, SyntaxToken}; pub use token_text::TokenText; /// `Parse` is the result of the parsing: a syntax tree and a collection of diff --git a/crates/squawk_syntax/src/postgresql.ungram b/crates/squawk_syntax/src/postgresql.ungram index 9da7529fa..a2d98e33e 100644 --- a/crates/squawk_syntax/src/postgresql.ungram +++ b/crates/squawk_syntax/src/postgresql.ungram @@ -664,7 +664,6 @@ CompressionMethod = 'compression' ('#ident' | 'default') Column = - 'period'? ( Name WithOptions? constraints:ColumnConstraint* DeferrableConstraintOption? NotDeferrableConstraintOption? @@ -674,7 +673,7 @@ Column = DeferrableConstraintOption? NotDeferrableConstraintOption? InitiallyDeferredConstraintOption? InitiallyImmediateConstraintOption? NotEnforced? Enforced? - | NameRef + | 'period'? NameRef | IndexExpr ) ColumnConstraint = diff --git a/crates/squawk_syntax/src/syntax_node.rs b/crates/squawk_syntax/src/syntax_node.rs index cf59b8ec5..0413eb7b3 100644 --- a/crates/squawk_syntax/src/syntax_node.rs +++ b/crates/squawk_syntax/src/syntax_node.rs @@ -45,7 +45,7 @@ impl Language for Sql { pub type SyntaxNode = rowan::SyntaxNode; pub type SyntaxToken = rowan::SyntaxToken; pub type SyntaxNodePtr = rowan::ast::SyntaxNodePtr; -// pub type SyntaxElement = rowan::SyntaxElement; +pub type SyntaxElement = rowan::SyntaxElement; pub type SyntaxNodeChildren = rowan::SyntaxNodeChildren; // pub type SyntaxElementChildren = rowan::SyntaxElementChildren; // pub type PreorderWithTokens = rowan::api::PreorderWithTokens; diff --git a/squawk-vscode/syntaxes/pgsql.tmLanguage.json b/squawk-vscode/syntaxes/pgsql.tmLanguage.json index 3661d799a..09453257b 100644 --- a/squawk-vscode/syntaxes/pgsql.tmLanguage.json +++ b/squawk-vscode/syntaxes/pgsql.tmLanguage.json @@ -214,7 +214,7 @@ "1": { "name": "punctuation.definition.string.begin.pgsql" }, - "3": { + "2": { "name": "punctuation.definition.string.end.pgsql" } },