ide: semantic syntax highlighting

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;

crates/squawk_ide/src/semantic_tokens.rs

@@ -0,0 +1,173 @@
+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<SemanticTokenModifier>,
+}
+
+#[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<TextRange>,
+) -> Vec<SemanticToken> {
+    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)
+                    && 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)) => {
+                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: once we get modfifiers, we'll need to update this
+            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
+        "#)
+    }
+}

crates/squawk_parser/src/grammar.rs

@@ -2511,7 +2511,9 @@ fn expr_bp(p: &mut Parser<'_>, bp: u8, r: &Restrictions) -> Option<CompletedMark
         // could be start of `is distinct from`
         && !(p.at(IS_KW) && p.nth_at(1, DISTINCT_KW))
     {
+        let m = p.start();
         col_label(p);
+        m.complete(p, AS_NAME);
         return Some(lhs);
     }
     if r.order_by_allowed && p.at(ORDER_KW) {

crates/squawk_parser/tests/snapshots/tests__select_ok.snap

@@ -5780,8 +5780,9 @@ SOURCE_FILE
           LITERAL
             INT_NUMBER "1"
           WHITESPACE " "
-          NAME
-            IS_KW "is"
+          AS_NAME
+            NAME
+              IS_KW "is"
   SEMICOLON ";"
   WHITESPACE "\n\n"
   SELECT
@@ -5793,8 +5794,9 @@ SOURCE_FILE
           LITERAL
             INT_NUMBER "1"
           WHITESPACE " "
-          NAME
-            AND_KW "and"
+          AS_NAME
+            NAME
+              AND_KW "and"
   SEMICOLON ";"
   WHITESPACE "\n\n"
   SELECT
@@ -5806,8 +5808,9 @@ SOURCE_FILE
           LITERAL
             INT_NUMBER "1"
           WHITESPACE " "
-          NAME
-            OR_KW "or"
+          AS_NAME
+            NAME
+              OR_KW "or"
   SEMICOLON ";"
   WHITESPACE "\n\n"
   SELECT
@@ -5819,8 +5822,9 @@ SOURCE_FILE
           LITERAL
             INT_NUMBER "1"
           WHITESPACE " "
-          NAME
-            COLLATE_KW "collate"
+          AS_NAME
+            NAME
+              COLLATE_KW "collate"
   SEMICOLON ";"
   WHITESPACE "\n\n"
   SELECT

crates/squawk_server/src/global_state.rs

@@ -16,16 +16,16 @@ use squawk_thread::TaskPool;
 use lsp_types::request::{
     CodeActionRequest, Completion, DocumentDiagnosticRequest, DocumentSymbolRequest,
     FoldingRangeRequest, GotoDefinition, HoverRequest, InlayHintRequest, References,
-    SelectionRangeRequest, Shutdown,
+    SelectionRangeRequest, SemanticTokensFullRequest, SemanticTokensRangeRequest, Shutdown,
 };
 
 use crate::dispatch::{NotificationDispatcher, RequestDispatcher};
 use crate::handlers::{
     SyntaxTreeRequest, TokensRequest, handle_cancel, handle_code_action, handle_completion,
     handle_did_change, handle_did_close, handle_did_open, handle_document_diagnostic,
     handle_document_symbol, handle_folding_range, handle_goto_definition, handle_hover,
-    handle_inlay_hints, handle_references, handle_selection_range, handle_shutdown,
-    handle_syntax_tree, handle_tokens,
+    handle_inlay_hints, handle_references, handle_selection_range, handle_semantic_tokens_full,
+    handle_semantic_tokens_range, handle_shutdown, handle_syntax_tree, handle_tokens,
 };
 
 type ReqQueue = lsp_server::ReqQueue<(String, Instant), ()>;
@@ -230,6 +230,8 @@ impl GlobalState {
             .on::<NO_RETRY, SyntaxTreeRequest>(handle_syntax_tree)
             .on::<NO_RETRY, TokensRequest>(handle_tokens)
             .on::<NO_RETRY, References>(handle_references)
+            .on::<NO_RETRY, SemanticTokensFullRequest>(handle_semantic_tokens_full)
+            .on::<NO_RETRY, SemanticTokensRangeRequest>(handle_semantic_tokens_range)
             .finish();
     }
 }

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};

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<Option<SemanticTokensResult>> {
+    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<Option<SemanticTokensRangeResult>> {
+    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),
+    })))
+}

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;