server: refactor lib.rs into separate files

crates/squawk_server/src/handlers.rs

@@ -0,0 +1,25 @@
+mod code_action;
+mod completion;
+mod document_symbol;
+mod folding_range;
+mod goto_definition;
+mod hover;
+mod inlay_hints;
+mod notifications;
+mod references;
+mod selection_range;
+mod syntax_tree;
+mod tokens;
+
+pub(crate) use code_action::handle_code_action;
+pub(crate) use completion::handle_completion;
+pub(crate) use document_symbol::handle_document_symbol;
+pub(crate) use folding_range::handle_folding_range;
+pub(crate) use goto_definition::handle_goto_definition;
+pub(crate) use hover::handle_hover;
+pub(crate) use inlay_hints::handle_inlay_hints;
+pub(crate) use notifications::{handle_did_change, handle_did_close, handle_did_open};
+pub(crate) use references::handle_references;
+pub(crate) use selection_range::handle_selection_range;
+pub(crate) use syntax_tree::handle_syntax_tree;
+pub(crate) use tokens::handle_tokens;

crates/squawk_server/src/handlers/code_action.rs

@@ -0,0 +1,147 @@
+use anyhow::{Context, Result};
+use lsp_server::{Connection, Message, Response};
+use lsp_types::{
+    CodeAction, CodeActionKind, CodeActionOrCommand, CodeActionParams, CodeActionResponse, Command,
+    WorkspaceEdit,
+};
+use squawk_ide::code_actions::code_actions;
+use squawk_ide::db::line_index;
+use std::collections::HashMap;
+
+use crate::diagnostic::{AssociatedDiagnosticData, DIAGNOSTIC_NAME};
+use crate::lsp_utils;
+use crate::system::System;
+
+pub(crate) fn handle_code_action(
+    connection: &Connection,
+    req: lsp_server::Request,
+    system: &impl System,
+) -> Result<()> {
+    let params: CodeActionParams = serde_json::from_value(req.params)?;
+    let uri = params.text_document.uri;
+
+    let mut actions: CodeActionResponse = vec![];
+
+    let db = system.db();
+    let file = system.file(&uri).unwrap();
+    let line_index = line_index(db, file);
+    let offset = lsp_utils::offset(&line_index, params.range.start).unwrap();
+
+    let ide_actions = code_actions(db, file, offset).unwrap_or_default();
+
+    for action in ide_actions {
+        let lsp_action = lsp_utils::code_action(&line_index, uri.clone(), action);
+        actions.push(CodeActionOrCommand::CodeAction(lsp_action));
+    }
+
+    for mut diagnostic in params
+        .context
+        .diagnostics
+        .into_iter()
+        .filter(|diagnostic| diagnostic.source.as_deref() == Some(DIAGNOSTIC_NAME))
+    {
+        let Some(rule_name) = diagnostic.code.as_ref().map(|x| match x {
+            lsp_types::NumberOrString::String(s) => s.clone(),
+            lsp_types::NumberOrString::Number(n) => n.to_string(),
+        }) else {
+            continue;
+        };
+        let Some(data) = diagnostic.data.take() else {
+            continue;
+        };
+
+        let associated_data: AssociatedDiagnosticData =
+            serde_json::from_value(data).context("deserializing diagnostic data")?;
+
+        if let Some(ignore_line_edit) = associated_data.ignore_line_edit {
+            let disable_line_action = CodeAction {
+                title: format!("Disable {rule_name} for this line"),
+                kind: Some(CodeActionKind::QUICKFIX),
+                diagnostics: Some(vec![diagnostic.clone()]),
+                edit: Some(WorkspaceEdit {
+                    changes: Some({
+                        let mut changes = HashMap::new();
+                        changes.insert(uri.clone(), vec![ignore_line_edit]);
+                        changes
+                    }),
+                    ..Default::default()
+                }),
+                command: None,
+                is_preferred: Some(false),
+                disabled: None,
+                data: None,
+            };
+            actions.push(CodeActionOrCommand::CodeAction(disable_line_action));
+        }
+        if let Some(ignore_file_edit) = associated_data.ignore_file_edit {
+            let disable_file_action = CodeAction {
+                title: format!("Disable {rule_name} for the entire file"),
+                kind: Some(CodeActionKind::QUICKFIX),
+                diagnostics: Some(vec![diagnostic.clone()]),
+                edit: Some(WorkspaceEdit {
+                    changes: Some({
+                        let mut changes = HashMap::new();
+                        changes.insert(uri.clone(), vec![ignore_file_edit]);
+                        changes
+                    }),
+                    ..Default::default()
+                }),
+                command: None,
+                is_preferred: Some(false),
+                disabled: None,
+                data: None,
+            };
+            actions.push(CodeActionOrCommand::CodeAction(disable_file_action));
+        }
+
+        let title = format!("Show documentation for {rule_name}");
+        let documentation_action = CodeAction {
+            title: title.clone(),
+            kind: Some(CodeActionKind::QUICKFIX),
+            diagnostics: Some(vec![diagnostic.clone()]),
+            edit: None,
+            command: Some(Command {
+                title,
+                command: "vscode.open".to_string(),
+                arguments: Some(vec![serde_json::to_value(format!(
+                    "https://squawkhq.com/docs/{rule_name}"
+                ))?]),
+            }),
+            is_preferred: Some(false),
+            disabled: None,
+            data: None,
+        };
+        actions.push(CodeActionOrCommand::CodeAction(documentation_action));
+
+        if !associated_data.title.is_empty() && !associated_data.edits.is_empty() {
+            let fix_action = CodeAction {
+                title: associated_data.title,
+                kind: Some(CodeActionKind::QUICKFIX),
+                diagnostics: Some(vec![diagnostic.clone()]),
+                edit: Some(WorkspaceEdit {
+                    changes: Some({
+                        let mut changes = HashMap::new();
+                        changes.insert(uri.clone(), associated_data.edits);
+                        changes
+                    }),
+                    ..Default::default()
+                }),
+                command: None,
+                is_preferred: Some(true),
+                disabled: None,
+                data: None,
+            };
+            actions.push(CodeActionOrCommand::CodeAction(fix_action));
+        }
+    }
+
+    let result: CodeActionResponse = actions;
+    let resp = Response {
+        id: req.id,
+        result: Some(serde_json::to_value(&result).unwrap()),
+        error: None,
+    };
+
+    connection.sender.send(Message::Response(resp))?;
+    Ok(())
+}

crates/squawk_server/src/handlers/completion.rs

@@ -0,0 +1,48 @@
+use anyhow::Result;
+use lsp_server::{Connection, Message, Response};
+use lsp_types::{CompletionParams, CompletionResponse};
+use squawk_ide::completion::completion;
+use squawk_ide::db::line_index;
+
+use crate::lsp_utils;
+use crate::system::System;
+
+pub(crate) fn handle_completion(
+    connection: &Connection,
+    req: lsp_server::Request,
+    system: &impl System,
+) -> Result<()> {
+    let params: CompletionParams = serde_json::from_value(req.params)?;
+    let uri = params.text_document_position.text_document.uri;
+    let position = params.text_document_position.position;
+
+    let db = system.db();
+    let file = system.file(&uri).unwrap();
+    let line_index = line_index(db, file);
+
+    let Some(offset) = lsp_utils::offset(&line_index, position) else {
+        let resp = Response {
+            id: req.id,
+            result: Some(serde_json::to_value(CompletionResponse::Array(vec![])).unwrap()),
+            error: None,
+        };
+        connection.sender.send(Message::Response(resp))?;
+        return Ok(());
+    };
+
+    let completion_items = completion(db, file, offset)
+        .into_iter()
+        .map(lsp_utils::completion_item)
+        .collect();
+
+    let result = CompletionResponse::Array(completion_items);
+
+    let resp = Response {
+        id: req.id,
+        result: Some(serde_json::to_value(&result).unwrap()),
+        error: None,
+    };
+
+    connection.sender.send(Message::Response(resp))?;
+    Ok(())
+}

crates/squawk_server/src/handlers/document_symbol.rs

@@ -0,0 +1,92 @@
+use ::line_index::LineIndex;
+use anyhow::Result;
+use lsp_server::{Connection, Message, Response};
+use lsp_types::{DocumentSymbol, DocumentSymbolParams, SymbolKind};
+use squawk_ide::db::line_index;
+use squawk_ide::document_symbols::{DocumentSymbolKind, document_symbols};
+
+use crate::lsp_utils;
+use crate::system::System;
+
+pub(crate) fn handle_document_symbol(
+    connection: &Connection,
+    req: lsp_server::Request,
+    system: &impl System,
+) -> Result<()> {
+    let params: DocumentSymbolParams = serde_json::from_value(req.params)?;
+    let uri = params.text_document.uri;
+
+    let db = system.db();
+    let file = system.file(&uri).unwrap();
+    let line_index = line_index(db, file);
+
+    let symbols = document_symbols(db, file);
+
+    fn convert_symbol(
+        sym: squawk_ide::document_symbols::DocumentSymbol,
+        line_index: &LineIndex,
+    ) -> DocumentSymbol {
+        let range = lsp_utils::range(line_index, sym.full_range);
+        let selection_range = lsp_utils::range(line_index, sym.focus_range);
+
+        let children = sym
+            .children
+            .into_iter()
+            .map(|child| convert_symbol(child, line_index))
+            .collect::<Vec<_>>();
+
+        let children = (!children.is_empty()).then_some(children);
+
+        DocumentSymbol {
+            name: sym.name,
+            detail: sym.detail,
+            kind: match sym.kind {
+                DocumentSymbolKind::Schema => SymbolKind::NAMESPACE,
+                DocumentSymbolKind::Table => SymbolKind::STRUCT,
+                DocumentSymbolKind::View => SymbolKind::STRUCT,
+                DocumentSymbolKind::MaterializedView => SymbolKind::STRUCT,
+                DocumentSymbolKind::Function => SymbolKind::FUNCTION,
+                DocumentSymbolKind::Aggregate => SymbolKind::FUNCTION,
+                DocumentSymbolKind::Procedure => SymbolKind::FUNCTION,
+                DocumentSymbolKind::Type => SymbolKind::CLASS,
+                DocumentSymbolKind::Enum => SymbolKind::ENUM,
+                DocumentSymbolKind::Index => SymbolKind::KEY,
+                DocumentSymbolKind::Domain => SymbolKind::CLASS,
+                DocumentSymbolKind::Sequence => SymbolKind::CONSTANT,
+                DocumentSymbolKind::Trigger => SymbolKind::EVENT,
+                DocumentSymbolKind::Tablespace => SymbolKind::NAMESPACE,
+                DocumentSymbolKind::Database => SymbolKind::MODULE,
+                DocumentSymbolKind::Server => SymbolKind::OBJECT,
+                DocumentSymbolKind::Extension => SymbolKind::PACKAGE,
+                DocumentSymbolKind::Column => SymbolKind::FIELD,
+                DocumentSymbolKind::Variant => SymbolKind::ENUM_MEMBER,
+                DocumentSymbolKind::Cursor => SymbolKind::VARIABLE,
+                DocumentSymbolKind::PreparedStatement => SymbolKind::VARIABLE,
+                DocumentSymbolKind::Channel => SymbolKind::EVENT,
+                DocumentSymbolKind::EventTrigger => SymbolKind::EVENT,
+                DocumentSymbolKind::Role => SymbolKind::CLASS,
+                DocumentSymbolKind::Policy => SymbolKind::VARIABLE,
+            },
+            tags: None,
+            range,
+            selection_range,
+            children,
+            #[allow(deprecated)]
+            deprecated: None,
+        }
+    }
+
+    let lsp_symbols: Vec<DocumentSymbol> = symbols
+        .into_iter()
+        .map(|sym| convert_symbol(sym, &line_index))
+        .collect();
+
+    let resp = Response {
+        id: req.id,
+        result: Some(serde_json::to_value(&lsp_symbols).unwrap()),
+        error: None,
+    };
+
+    connection.sender.send(Message::Response(resp))?;
+    Ok(())
+}

crates/squawk_server/src/handlers/folding_range.rs

@@ -0,0 +1,35 @@
+use anyhow::Result;
+use lsp_server::{Connection, Message, Response};
+use lsp_types::FoldingRange;
+use squawk_ide::db::line_index;
+use squawk_ide::folding_ranges::folding_ranges;
+
+use crate::lsp_utils;
+use crate::system::System;
+
+pub(crate) fn handle_folding_range(
+    connection: &Connection,
+    req: lsp_server::Request,
+    system: &impl System,
+) -> Result<()> {
+    let params: lsp_types::FoldingRangeParams = serde_json::from_value(req.params)?;
+    let uri = params.text_document.uri;
+
+    let db = system.db();
+    let file = system.file(&uri).unwrap();
+    let line_idx = line_index(db, file);
+
+    let lsp_folds: Vec<FoldingRange> = folding_ranges(db, file)
+        .into_iter()
+        .map(|fold| lsp_utils::folding_range(&line_idx, fold))
+        .collect();
+
+    let resp = Response {
+        id: req.id,
+        result: Some(serde_json::to_value(&lsp_folds).unwrap()),
+        error: None,
+    };
+
+    connection.sender.send(Message::Response(resp))?;
+    Ok(())
+}

crates/squawk_server/src/handlers/goto_definition.rs

@@ -0,0 +1,44 @@
+use anyhow::Result;
+use lsp_server::{Connection, Message, Response};
+use lsp_types::{GotoDefinitionParams, GotoDefinitionResponse};
+use squawk_ide::db::line_index;
+use squawk_ide::goto_definition::goto_definition;
+
+use crate::lsp_utils::{self, to_location};
+use crate::system::System;
+
+pub(crate) fn handle_goto_definition(
+    connection: &Connection,
+    req: lsp_server::Request,
+    system: &impl System,
+) -> Result<()> {
+    let params: GotoDefinitionParams = serde_json::from_value(req.params)?;
+    let uri = params.text_document_position_params.text_document.uri;
+    let position = params.text_document_position_params.position;
+
+    let db = system.db();
+    let file = system.file(&uri).unwrap();
+    let line_index = line_index(db, file);
+    let offset = lsp_utils::offset(&line_index, position).unwrap();
+
+    let ranges = goto_definition(db, file, offset)
+        .into_iter()
+        .filter_map(|location| {
+            debug_assert!(
+                !location.range.contains(offset),
+                "Our target destination range must not include the source range otherwise go to def won't work in vscode."
+            );
+            to_location(db, system, &uri, location)
+        })
+        .collect();
+
+    let result = GotoDefinitionResponse::Array(ranges);
+    let resp = Response {
+        id: req.id,
+        result: Some(serde_json::to_value(&result).unwrap()),
+        error: None,
+    };
+
+    connection.sender.send(Message::Response(resp))?;
+    Ok(())
+}