ide: semantic token highlighting for func params & types (#1057)

sbdchd · web-flow · commit 902cab31a045 · 2026-04-13T22:40:16.000Z
diff --git 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<SemanticToken>,
+}
+
+impl SemanticTokenBuilder {
+    fn build(mut self) -> Vec<SemanticToken> {
+        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
@@ -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
@@ -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
@@ -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
diff --git a/squawk-vscode/syntaxes/pgsql.tmLanguage.json b/squawk-vscode/syntaxes/pgsql.tmLanguage.json
