server: refactor request handling to embrace types

crates/squawk_server/src/dispatch.rs

@@ -0,0 +1,78 @@
+// Based on https://github.com/rust-lang/rust-analyzer/blob/2efc80078029894eec0699f62ec8d5c1a56af763/crates/rust-analyzer/src/handlers/dispatch.rs#L277C5-L277C5
+
+use anyhow::Result;
+use log::{error, info};
+use lsp_server::{Connection, Message, Response};
+use lsp_types::request::Request as LspRequest;
+
+use crate::system::System;
+
+pub(crate) struct RequestDispatcher<'a> {
+    connection: &'a Connection,
+    req: Option<lsp_server::Request>,
+    system: &'a dyn System,
+}
+
+impl<'a> RequestDispatcher<'a> {
+    pub(crate) fn new(
+        connection: &'a Connection,
+        req: lsp_server::Request,
+        system: &'a dyn System,
+    ) -> Self {
+        Self {
+            connection,
+            req: Some(req),
+            system,
+        }
+    }
+
+    fn parse<R>(&mut self) -> Option<(lsp_server::RequestId, R::Params)>
+    where
+        R: LspRequest,
+    {
+        let req = self.req.take_if(|req| req.method.as_str() == R::METHOD)?;
+        let id = req.id.clone();
+
+        match req.extract(R::METHOD) {
+            Ok((id, params)) => Some((id, params)),
+            Err(err) => {
+                let response = lsp_server::Response::new_err(
+                    id,
+                    lsp_server::ErrorCode::InvalidParams as i32,
+                    err.to_string(),
+                );
+                if let Err(err) = self.connection.sender.send(Message::Response(response)) {
+                    error!("Failed to send parse error response: {err}");
+                }
+                None
+            }
+        }
+    }
+
+    pub(crate) fn on<R>(
+        mut self,
+        handler: fn(&dyn System, R::Params) -> Result<R::Result>,
+    ) -> Result<Self>
+    where
+        R: LspRequest,
+    {
+        if let Some((id, params)) = self.parse::<R>() {
+            // TODO: we probably need to send an error here
+            let result = handler(self.system, params)?;
+            let resp = Response {
+                id,
+                result: Some(serde_json::to_value(result)?),
+                error: None,
+            };
+            self.connection.sender.send(Message::Response(resp))?;
+        }
+
+        Ok(self)
+    }
+
+    pub(crate) fn finish(self) {
+        if let Some(req) = self.req {
+            info!("Ignoring unhandled request: {}", req.method);
+        }
+    }
+}

crates/squawk_server/src/handlers.rs

@@ -21,5 +21,5 @@ 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;
+pub(crate) use syntax_tree::{SyntaxTreeRequest, handle_syntax_tree};
+pub(crate) use tokens::{TokensRequest, handle_tokens};

crates/squawk_server/src/handlers/code_action.rs

@@ -1,5 +1,4 @@
 use anyhow::{Context, Result};
-use lsp_server::{Connection, Message, Response};
 use lsp_types::{
     CodeAction, CodeActionKind, CodeActionOrCommand, CodeActionParams, CodeActionResponse, Command,
     WorkspaceEdit,
@@ -13,11 +12,9 @@ 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)?;
+    system: &dyn System,
+    params: CodeActionParams,
+) -> Result<Option<CodeActionResponse>> {
     let uri = params.text_document.uri;
 
     let mut actions: CodeActionResponse = vec![];
@@ -135,13 +132,5 @@ pub(crate) fn handle_code_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(())
+    Ok(Some(actions))
 }

crates/squawk_server/src/handlers/completion.rs

@@ -1,5 +1,4 @@
 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;
@@ -8,11 +7,9 @@ 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)?;
+    system: &dyn System,
+    params: CompletionParams,
+) -> Result<Option<CompletionResponse>> {
     let uri = params.text_document_position.text_document.uri;
     let position = params.text_document_position.position;
 
@@ -21,28 +18,13 @@ pub(crate) fn handle_completion(
     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(());
+        return Ok(Some(CompletionResponse::Array(vec![])));
     };
 
     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(())
+    Ok(Some(CompletionResponse::Array(completion_items)))
 }

crates/squawk_server/src/handlers/document_symbol.rs

@@ -1,19 +1,16 @@
 use ::line_index::LineIndex;
 use anyhow::Result;
-use lsp_server::{Connection, Message, Response};
-use lsp_types::{DocumentSymbol, DocumentSymbolParams, SymbolKind};
+use lsp_types::{DocumentSymbol, DocumentSymbolParams, DocumentSymbolResponse, 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)?;
+    system: &dyn System,
+    params: DocumentSymbolParams,
+) -> Result<Option<DocumentSymbolResponse>> {
     let uri = params.text_document.uri;
 
     let db = system.db();
@@ -81,12 +78,5 @@ pub(crate) fn handle_document_symbol(
         .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(())
+    Ok(Some(DocumentSymbolResponse::Nested(lsp_symbols)))
 }

crates/squawk_server/src/handlers/folding_range.rs

@@ -1,18 +1,15 @@
 use anyhow::Result;
-use lsp_server::{Connection, Message, Response};
-use lsp_types::FoldingRange;
+use lsp_types::{FoldingRange, FoldingRangeParams};
 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)?;
+    system: &dyn System,
+    params: FoldingRangeParams,
+) -> Result<Option<Vec<FoldingRange>>> {
     let uri = params.text_document.uri;
 
     let db = system.db();
@@ -24,12 +21,5 @@ pub(crate) fn handle_folding_range(
         .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(())
+    Ok(Some(lsp_folds))
 }

crates/squawk_server/src/handlers/goto_definition.rs

@@ -1,5 +1,4 @@
 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;
@@ -8,11 +7,9 @@ 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)?;
+    system: &dyn System,
+    params: GotoDefinitionParams,
+) -> Result<Option<GotoDefinitionResponse>> {
     let uri = params.text_document_position_params.text_document.uri;
     let position = params.text_document_position_params.position;
 
@@ -32,13 +29,5 @@ pub(crate) fn handle_goto_definition(
         })
         .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(())
+    Ok(Some(GotoDefinitionResponse::Array(ranges)))
 }

crates/squawk_server/src/handlers/hover.rs

@@ -1,18 +1,12 @@
 use anyhow::Result;
-use lsp_server::{Connection, Message, Response};
 use lsp_types::{Hover, HoverContents, HoverParams, LanguageString, MarkedString};
 use squawk_ide::db::line_index;
 use squawk_ide::hover::hover;
 
 use crate::lsp_utils;
 use crate::system::System;
 
-pub(crate) fn handle_hover(
-    connection: &Connection,
-    req: lsp_server::Request,
-    system: &impl System,
-) -> Result<()> {
-    let params: HoverParams = serde_json::from_value(req.params)?;
+pub(crate) fn handle_hover(system: &dyn System, params: HoverParams) -> Result<Option<Hover>> {
     let uri = params.text_document_position_params.text_document.uri;
     let position = params.text_document_position_params.position;
 
@@ -23,20 +17,11 @@ pub(crate) fn handle_hover(
 
     let type_info = hover(db, file, offset);
 
-    let result = type_info.map(|type_str| Hover {
+    Ok(type_info.map(|type_str| Hover {
         contents: HoverContents::Scalar(MarkedString::LanguageString(LanguageString {
             language: "sql".to_string(),
             value: type_str,
         })),
         range: None,
-    });
-
-    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/inlay_hints.rs

@@ -1,5 +1,4 @@
 use anyhow::Result;
-use lsp_server::{Connection, Message, Response};
 use lsp_types::{
     InlayHint, InlayHintKind, InlayHintLabel, InlayHintLabelPart, InlayHintParams, Location,
 };
@@ -11,11 +10,9 @@ use crate::lsp_utils;
 use crate::system::System;
 
 pub(crate) fn handle_inlay_hints(
-    connection: &Connection,
-    req: lsp_server::Request,
-    system: &impl System,
-) -> Result<()> {
-    let params: InlayHintParams = serde_json::from_value(req.params)?;
+    system: &dyn System,
+    params: InlayHintParams,
+) -> Result<Option<Vec<InlayHint>>> {
     let uri = params.text_document.uri;
 
     let db = system.db();
@@ -41,8 +38,8 @@ pub(crate) fn handle_inlay_hints(
             };
 
             let kind: InlayHintKind = match hint.kind {
-                squawk_ide::inlay_hints::InlayHintKind::Type => InlayHintKind::TYPE,
                 squawk_ide::inlay_hints::InlayHintKind::Parameter => InlayHintKind::PARAMETER,
+                squawk_ide::inlay_hints::InlayHintKind::Type => InlayHintKind::TYPE,
             };
 
             let label = if let Some(target_range) = hint.target {
@@ -72,12 +69,5 @@ pub(crate) fn handle_inlay_hints(
         })
         .collect();
 
-    let resp = Response {
-        id: req.id,
-        result: Some(serde_json::to_value(&lsp_hints).unwrap()),
-        error: None,
-    };
-
-    connection.sender.send(Message::Response(resp))?;
-    Ok(())
+    Ok(Some(lsp_hints))
 }