Skip to main content

squawk_server/
lib.rs

1use anyhow::{Context, Result};
2use line_index::LineIndex;
3use log::info;
4use lsp_server::{Connection, Message, Notification, Response};
5use lsp_types::{
6    CodeAction, CodeActionKind, CodeActionOptions, CodeActionOrCommand, CodeActionParams,
7    CodeActionProviderCapability, CodeActionResponse, Command, CompletionOptions, CompletionParams,
8    CompletionResponse, Diagnostic, DidChangeTextDocumentParams, DidCloseTextDocumentParams,
9    DidOpenTextDocumentParams, DocumentSymbol, DocumentSymbolParams, GotoDefinitionParams,
10    GotoDefinitionResponse, Hover, HoverContents, HoverParams, HoverProviderCapability,
11    InitializeParams, InlayHint, InlayHintKind, InlayHintLabel, InlayHintLabelPart,
12    InlayHintParams, LanguageString, Location, MarkedString, OneOf, PublishDiagnosticsParams,
13    ReferenceParams, SelectionRangeParams, SelectionRangeProviderCapability, ServerCapabilities,
14    SymbolKind, TextDocumentSyncCapability, TextDocumentSyncKind, Url, WorkDoneProgressOptions,
15    WorkspaceEdit,
16    notification::{
17        DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument, Notification as _,
18        PublishDiagnostics,
19    },
20    request::{
21        CodeActionRequest, Completion, DocumentSymbolRequest, GotoDefinition, HoverRequest,
22        InlayHintRequest, References, Request, SelectionRangeRequest,
23    },
24};
25use rowan::TextRange;
26use squawk_ide::code_actions::code_actions;
27use squawk_ide::completion::completion;
28use squawk_ide::document_symbols::{DocumentSymbolKind, document_symbols};
29use squawk_ide::find_references::find_references;
30use squawk_ide::goto_definition::goto_definition;
31use squawk_ide::hover::hover;
32use squawk_ide::inlay_hints::inlay_hints;
33use squawk_syntax::SourceFile;
34use std::collections::HashMap;
35
36use diagnostic::DIAGNOSTIC_NAME;
37
38use crate::diagnostic::AssociatedDiagnosticData;
39mod diagnostic;
40mod ignore;
41mod lint;
42mod lsp_utils;
43
44struct DocumentState {
45    content: String,
46    version: i32,
47}
48
49pub fn run() -> Result<()> {
50    info!("Starting Squawk LSP server");
51
52    let (connection, io_threads) = Connection::stdio();
53
54    let server_capabilities = serde_json::to_value(&ServerCapabilities {
55        text_document_sync: Some(TextDocumentSyncCapability::Kind(
56            TextDocumentSyncKind::INCREMENTAL,
57        )),
58        code_action_provider: Some(CodeActionProviderCapability::Options(CodeActionOptions {
59            code_action_kinds: Some(vec![
60                CodeActionKind::QUICKFIX,
61                CodeActionKind::REFACTOR_REWRITE,
62            ]),
63            work_done_progress_options: WorkDoneProgressOptions {
64                work_done_progress: None,
65            },
66            resolve_provider: None,
67        })),
68        selection_range_provider: Some(SelectionRangeProviderCapability::Simple(true)),
69        references_provider: Some(OneOf::Left(true)),
70        definition_provider: Some(OneOf::Left(true)),
71        hover_provider: Some(HoverProviderCapability::Simple(true)),
72        inlay_hint_provider: Some(OneOf::Left(true)),
73        document_symbol_provider: Some(OneOf::Left(true)),
74        completion_provider: Some(CompletionOptions {
75            resolve_provider: Some(false),
76            trigger_characters: Some(vec![".".to_owned()]),
77            all_commit_characters: None,
78            work_done_progress_options: WorkDoneProgressOptions {
79                work_done_progress: None,
80            },
81            completion_item: None,
82        }),
83        ..Default::default()
84    })
85    .unwrap();
86
87    info!("LSP server initializing connection...");
88    let initialization_params = connection.initialize(server_capabilities)?;
89    info!("LSP server initialized, entering main loop");
90
91    main_loop(connection, initialization_params)?;
92
93    info!("LSP server shutting down");
94
95    io_threads.join()?;
96    Ok(())
97}
98
99fn main_loop(connection: Connection, params: serde_json::Value) -> Result<()> {
100    info!("Server main loop");
101
102    let init_params: InitializeParams = serde_json::from_value(params).unwrap_or_default();
103    info!("Client process ID: {:?}", init_params.process_id);
104    let client_name = init_params.client_info.map(|x| x.name);
105    info!("Client name: {client_name:?}");
106
107    let mut documents: HashMap<Url, DocumentState> = HashMap::new();
108
109    for msg in &connection.receiver {
110        match msg {
111            Message::Request(req) => {
112                info!("Received request: method={}, id={:?}", req.method, req.id);
113
114                if connection.handle_shutdown(&req)? {
115                    info!("Received shutdown request, exiting");
116                    return Ok(());
117                }
118
119                match req.method.as_ref() {
120                    GotoDefinition::METHOD => {
121                        handle_goto_definition(&connection, req, &documents)?;
122                    }
123                    HoverRequest::METHOD => {
124                        handle_hover(&connection, req, &documents)?;
125                    }
126                    CodeActionRequest::METHOD => {
127                        handle_code_action(&connection, req, &documents)?;
128                    }
129                    SelectionRangeRequest::METHOD => {
130                        handle_selection_range(&connection, req, &documents)?;
131                    }
132                    InlayHintRequest::METHOD => {
133                        handle_inlay_hints(&connection, req, &documents)?;
134                    }
135                    DocumentSymbolRequest::METHOD => {
136                        handle_document_symbol(&connection, req, &documents)?;
137                    }
138                    Completion::METHOD => {
139                        handle_completion(&connection, req, &documents)?;
140                    }
141                    "squawk/syntaxTree" => {
142                        handle_syntax_tree(&connection, req, &documents)?;
143                    }
144                    "squawk/tokens" => {
145                        handle_tokens(&connection, req, &documents)?;
146                    }
147                    References::METHOD => {
148                        handle_references(&connection, req, &documents)?;
149                    }
150                    _ => {
151                        info!("Ignoring unhandled request: {}", req.method);
152                    }
153                }
154            }
155            Message::Response(resp) => {
156                info!("Received response: id={:?}", resp.id);
157            }
158            Message::Notification(notif) => {
159                info!("Received notification: method={}", notif.method);
160                match notif.method.as_ref() {
161                    DidOpenTextDocument::METHOD => {
162                        handle_did_open(&connection, notif, &mut documents)?;
163                    }
164                    DidChangeTextDocument::METHOD => {
165                        handle_did_change(&connection, notif, &mut documents)?;
166                    }
167                    DidCloseTextDocument::METHOD => {
168                        handle_did_close(&connection, notif, &mut documents)?;
169                    }
170                    _ => {
171                        info!("Ignoring unhandled notification: {}", notif.method);
172                    }
173                }
174            }
175        }
176    }
177    Ok(())
178}
179
180fn handle_goto_definition(
181    connection: &Connection,
182    req: lsp_server::Request,
183    documents: &HashMap<Url, DocumentState>,
184) -> Result<()> {
185    let params: GotoDefinitionParams = serde_json::from_value(req.params)?;
186    let uri = params.text_document_position_params.text_document.uri;
187    let position = params.text_document_position_params.position;
188
189    let content = documents.get(&uri).map_or("", |doc| &doc.content);
190    let parse = SourceFile::parse(content);
191    let file = parse.tree();
192    let line_index = LineIndex::new(content);
193    let offset = lsp_utils::offset(&line_index, position).unwrap();
194
195    let ranges = goto_definition(file, offset)
196        .into_iter()
197        .map(|target_range| {
198            debug_assert!(
199                !target_range.contains(offset),
200                "Our target destination range must not include the source range otherwise go to def won't work in vscode."
201            );
202            Location {
203                uri: uri.clone(),
204                range: lsp_utils::range(&line_index, target_range),
205            }
206        })
207        .collect();
208
209    let result = GotoDefinitionResponse::Array(ranges);
210    let resp = Response {
211        id: req.id,
212        result: Some(serde_json::to_value(&result).unwrap()),
213        error: None,
214    };
215
216    connection.sender.send(Message::Response(resp))?;
217    Ok(())
218}
219
220fn handle_hover(
221    connection: &Connection,
222    req: lsp_server::Request,
223    documents: &HashMap<Url, DocumentState>,
224) -> Result<()> {
225    let params: HoverParams = serde_json::from_value(req.params)?;
226    let uri = params.text_document_position_params.text_document.uri;
227    let position = params.text_document_position_params.position;
228
229    let content = documents.get(&uri).map_or("", |doc| &doc.content);
230    let parse = SourceFile::parse(content);
231    let file = parse.tree();
232    let line_index = LineIndex::new(content);
233    let offset = lsp_utils::offset(&line_index, position).unwrap();
234
235    let type_info = hover(&file, offset);
236
237    let result = type_info.map(|type_str| Hover {
238        contents: HoverContents::Scalar(MarkedString::LanguageString(LanguageString {
239            language: "sql".to_string(),
240            value: type_str,
241        })),
242        range: None,
243    });
244
245    let resp = Response {
246        id: req.id,
247        result: Some(serde_json::to_value(&result).unwrap()),
248        error: None,
249    };
250
251    connection.sender.send(Message::Response(resp))?;
252    Ok(())
253}
254
255fn handle_inlay_hints(
256    connection: &Connection,
257    req: lsp_server::Request,
258    documents: &HashMap<Url, DocumentState>,
259) -> Result<()> {
260    let params: InlayHintParams = serde_json::from_value(req.params)?;
261    let uri = params.text_document.uri;
262
263    let content = documents.get(&uri).map_or("", |doc| &doc.content);
264    let parse = SourceFile::parse(content);
265    let file = parse.tree();
266    let line_index = LineIndex::new(content);
267
268    let hints = inlay_hints(&file);
269
270    let lsp_hints: Vec<InlayHint> = hints
271        .into_iter()
272        .map(|hint| {
273            let line_col = line_index.line_col(hint.position);
274            let position = lsp_types::Position::new(line_col.line, line_col.col);
275            let kind = match hint.kind {
276                squawk_ide::inlay_hints::InlayHintKind::Type => InlayHintKind::TYPE,
277                squawk_ide::inlay_hints::InlayHintKind::Parameter => InlayHintKind::PARAMETER,
278            };
279
280            let label = if let Some(target_range) = hint.target {
281                InlayHintLabel::LabelParts(vec![InlayHintLabelPart {
282                    value: hint.label,
283                    location: Some(Location {
284                        uri: uri.clone(),
285                        range: lsp_utils::range(&line_index, target_range),
286                    }),
287                    tooltip: None,
288                    command: None,
289                }])
290            } else {
291                InlayHintLabel::String(hint.label)
292            };
293
294            InlayHint {
295                position,
296                label,
297                kind: Some(kind),
298                text_edits: None,
299                tooltip: None,
300                padding_left: None,
301                padding_right: None,
302                data: None,
303            }
304        })
305        .collect();
306
307    let resp = Response {
308        id: req.id,
309        result: Some(serde_json::to_value(&lsp_hints).unwrap()),
310        error: None,
311    };
312
313    connection.sender.send(Message::Response(resp))?;
314    Ok(())
315}
316
317fn handle_document_symbol(
318    connection: &Connection,
319    req: lsp_server::Request,
320    documents: &HashMap<Url, DocumentState>,
321) -> Result<()> {
322    let params: DocumentSymbolParams = serde_json::from_value(req.params)?;
323    let uri = params.text_document.uri;
324
325    let content = documents.get(&uri).map_or("", |doc| &doc.content);
326    let parse = SourceFile::parse(content);
327    let file = parse.tree();
328    let line_index = LineIndex::new(content);
329
330    let symbols = document_symbols(&file);
331
332    fn convert_symbol(
333        sym: squawk_ide::document_symbols::DocumentSymbol,
334        line_index: &LineIndex,
335    ) -> DocumentSymbol {
336        let range = lsp_utils::range(line_index, sym.full_range);
337        let selection_range = lsp_utils::range(line_index, sym.focus_range);
338
339        let children = sym
340            .children
341            .into_iter()
342            .map(|child| convert_symbol(child, line_index))
343            .collect::<Vec<_>>();
344
345        let children = (!children.is_empty()).then_some(children);
346
347        DocumentSymbol {
348            name: sym.name,
349            detail: sym.detail,
350            kind: match sym.kind {
351                DocumentSymbolKind::Schema => SymbolKind::NAMESPACE,
352                DocumentSymbolKind::Table => SymbolKind::STRUCT,
353                DocumentSymbolKind::View => SymbolKind::STRUCT,
354                DocumentSymbolKind::MaterializedView => SymbolKind::STRUCT,
355                DocumentSymbolKind::Function => SymbolKind::FUNCTION,
356                DocumentSymbolKind::Aggregate => SymbolKind::FUNCTION,
357                DocumentSymbolKind::Procedure => SymbolKind::FUNCTION,
358                DocumentSymbolKind::Type => SymbolKind::CLASS,
359                DocumentSymbolKind::Enum => SymbolKind::ENUM,
360                DocumentSymbolKind::Index => SymbolKind::KEY,
361                DocumentSymbolKind::Domain => SymbolKind::CLASS,
362                DocumentSymbolKind::Sequence => SymbolKind::CONSTANT,
363                DocumentSymbolKind::Trigger => SymbolKind::EVENT,
364                DocumentSymbolKind::Tablespace => SymbolKind::NAMESPACE,
365                DocumentSymbolKind::Database => SymbolKind::MODULE,
366                DocumentSymbolKind::Server => SymbolKind::OBJECT,
367                DocumentSymbolKind::Extension => SymbolKind::PACKAGE,
368                DocumentSymbolKind::Column => SymbolKind::FIELD,
369                DocumentSymbolKind::Variant => SymbolKind::ENUM_MEMBER,
370                DocumentSymbolKind::Cursor => SymbolKind::VARIABLE,
371                DocumentSymbolKind::PreparedStatement => SymbolKind::VARIABLE,
372                DocumentSymbolKind::Channel => SymbolKind::EVENT,
373                DocumentSymbolKind::EventTrigger => SymbolKind::EVENT,
374                DocumentSymbolKind::Role => SymbolKind::CLASS,
375                DocumentSymbolKind::Policy => SymbolKind::VARIABLE,
376            },
377            tags: None,
378            range,
379            selection_range,
380            children,
381            #[allow(deprecated)]
382            deprecated: None,
383        }
384    }
385
386    let lsp_symbols: Vec<DocumentSymbol> = symbols
387        .into_iter()
388        .map(|sym| convert_symbol(sym, &line_index))
389        .collect();
390
391    let resp = Response {
392        id: req.id,
393        result: Some(serde_json::to_value(&lsp_symbols).unwrap()),
394        error: None,
395    };
396
397    connection.sender.send(Message::Response(resp))?;
398    Ok(())
399}
400
401fn handle_selection_range(
402    connection: &Connection,
403    req: lsp_server::Request,
404    documents: &HashMap<Url, DocumentState>,
405) -> Result<()> {
406    let params: SelectionRangeParams = serde_json::from_value(req.params)?;
407    let uri = params.text_document.uri;
408
409    let content = documents.get(&uri).map_or("", |doc| &doc.content);
410    let parse = SourceFile::parse(content);
411    let root = parse.syntax_node();
412    let line_index = LineIndex::new(content);
413
414    let mut selection_ranges = vec![];
415
416    for position in params.positions {
417        let Some(offset) = lsp_utils::offset(&line_index, position) else {
418            continue;
419        };
420
421        let mut ranges = Vec::new();
422        {
423            let mut range = TextRange::new(offset, offset);
424            loop {
425                ranges.push(range);
426                let next = squawk_ide::expand_selection::extend_selection(&root, range);
427                if next == range {
428                    break;
429                } else {
430                    range = next
431                }
432            }
433        }
434
435        let mut range = lsp_types::SelectionRange {
436            range: lsp_utils::range(&line_index, *ranges.last().unwrap()),
437            parent: None,
438        };
439        for &r in ranges.iter().rev().skip(1) {
440            range = lsp_types::SelectionRange {
441                range: lsp_utils::range(&line_index, r),
442                parent: Some(Box::new(range)),
443            }
444        }
445        selection_ranges.push(range);
446    }
447
448    let resp = Response {
449        id: req.id,
450        result: Some(serde_json::to_value(&selection_ranges).unwrap()),
451        error: None,
452    };
453
454    connection.sender.send(Message::Response(resp))?;
455    Ok(())
456}
457
458fn handle_references(
459    connection: &Connection,
460    req: lsp_server::Request,
461    documents: &HashMap<Url, DocumentState>,
462) -> Result<()> {
463    let params: ReferenceParams = serde_json::from_value(req.params)?;
464    let uri = params.text_document_position.text_document.uri;
465    let position = params.text_document_position.position;
466
467    let content = documents.get(&uri).map_or("", |doc| &doc.content);
468    let parse = SourceFile::parse(content);
469    let file = parse.tree();
470    let line_index = LineIndex::new(content);
471    let offset = lsp_utils::offset(&line_index, position).unwrap();
472
473    let ranges = find_references(&file, offset);
474    let include_declaration = params.context.include_declaration;
475
476    let locations: Vec<Location> = ranges
477        .into_iter()
478        .filter(|range| include_declaration || !range.contains(offset))
479        .map(|range| Location {
480            uri: uri.clone(),
481            range: lsp_utils::range(&line_index, range),
482        })
483        .collect();
484
485    let resp = Response {
486        id: req.id,
487        result: Some(serde_json::to_value(&locations).unwrap()),
488        error: None,
489    };
490
491    connection.sender.send(Message::Response(resp))?;
492    Ok(())
493}
494
495fn handle_completion(
496    connection: &Connection,
497    req: lsp_server::Request,
498    documents: &HashMap<Url, DocumentState>,
499) -> Result<()> {
500    let params: CompletionParams = serde_json::from_value(req.params)?;
501    let uri = params.text_document_position.text_document.uri;
502    let position = params.text_document_position.position;
503
504    let content = documents.get(&uri).map_or("", |doc| &doc.content);
505    let parse = SourceFile::parse(content);
506    let file = parse.tree();
507    let line_index = LineIndex::new(content);
508
509    let Some(offset) = lsp_utils::offset(&line_index, position) else {
510        let resp = Response {
511            id: req.id,
512            result: Some(serde_json::to_value(CompletionResponse::Array(vec![])).unwrap()),
513            error: None,
514        };
515        connection.sender.send(Message::Response(resp))?;
516        return Ok(());
517    };
518
519    let completion_items = completion(&file, offset)
520        .into_iter()
521        .map(lsp_utils::completion_item)
522        .collect();
523
524    let result = CompletionResponse::Array(completion_items);
525
526    let resp = Response {
527        id: req.id,
528        result: Some(serde_json::to_value(&result).unwrap()),
529        error: None,
530    };
531
532    connection.sender.send(Message::Response(resp))?;
533    Ok(())
534}
535
536fn handle_code_action(
537    connection: &Connection,
538    req: lsp_server::Request,
539    documents: &HashMap<Url, DocumentState>,
540) -> Result<()> {
541    let params: CodeActionParams = serde_json::from_value(req.params)?;
542    let uri = params.text_document.uri;
543
544    let mut actions: CodeActionResponse = Vec::new();
545
546    let content = documents.get(&uri).map_or("", |doc| &doc.content);
547    let parse = SourceFile::parse(content);
548    let file = parse.tree();
549    let line_index = LineIndex::new(content);
550    let offset = lsp_utils::offset(&line_index, params.range.start).unwrap();
551
552    let ide_actions = code_actions(file, offset).unwrap_or_default();
553
554    for action in ide_actions {
555        let lsp_action = lsp_utils::code_action(&line_index, uri.clone(), action);
556        actions.push(CodeActionOrCommand::CodeAction(lsp_action));
557    }
558
559    for mut diagnostic in params
560        .context
561        .diagnostics
562        .into_iter()
563        .filter(|diagnostic| diagnostic.source.as_deref() == Some(DIAGNOSTIC_NAME))
564    {
565        let Some(rule_name) = diagnostic.code.as_ref().map(|x| match x {
566            lsp_types::NumberOrString::String(s) => s.clone(),
567            lsp_types::NumberOrString::Number(n) => n.to_string(),
568        }) else {
569            continue;
570        };
571        let Some(data) = diagnostic.data.take() else {
572            continue;
573        };
574
575        let associated_data: AssociatedDiagnosticData =
576            serde_json::from_value(data).context("deserializing diagnostic data")?;
577
578        if let Some(ignore_line_edit) = associated_data.ignore_line_edit {
579            let disable_line_action = CodeAction {
580                title: format!("Disable {rule_name} for this line"),
581                kind: Some(CodeActionKind::QUICKFIX),
582                diagnostics: Some(vec![diagnostic.clone()]),
583                edit: Some(WorkspaceEdit {
584                    changes: Some({
585                        let mut changes = HashMap::new();
586                        changes.insert(uri.clone(), vec![ignore_line_edit]);
587                        changes
588                    }),
589                    ..Default::default()
590                }),
591                command: None,
592                is_preferred: Some(false),
593                disabled: None,
594                data: None,
595            };
596            actions.push(CodeActionOrCommand::CodeAction(disable_line_action));
597        }
598        if let Some(ignore_file_edit) = associated_data.ignore_file_edit {
599            let disable_file_action = CodeAction {
600                title: format!("Disable {rule_name} for the entire file"),
601                kind: Some(CodeActionKind::QUICKFIX),
602                diagnostics: Some(vec![diagnostic.clone()]),
603                edit: Some(WorkspaceEdit {
604                    changes: Some({
605                        let mut changes = HashMap::new();
606                        changes.insert(uri.clone(), vec![ignore_file_edit]);
607                        changes
608                    }),
609                    ..Default::default()
610                }),
611                command: None,
612                is_preferred: Some(false),
613                disabled: None,
614                data: None,
615            };
616            actions.push(CodeActionOrCommand::CodeAction(disable_file_action));
617        }
618
619        let title = format!("Show documentation for {rule_name}");
620        let documentation_action = CodeAction {
621            title: title.clone(),
622            kind: Some(CodeActionKind::QUICKFIX),
623            diagnostics: Some(vec![diagnostic.clone()]),
624            edit: None,
625            command: Some(Command {
626                title,
627                command: "vscode.open".to_string(),
628                arguments: Some(vec![serde_json::to_value(format!(
629                    "https://squawkhq.com/docs/{rule_name}"
630                ))?]),
631            }),
632            is_preferred: Some(false),
633            disabled: None,
634            data: None,
635        };
636        actions.push(CodeActionOrCommand::CodeAction(documentation_action));
637
638        if !associated_data.title.is_empty() && !associated_data.edits.is_empty() {
639            let fix_action = CodeAction {
640                title: associated_data.title,
641                kind: Some(CodeActionKind::QUICKFIX),
642                diagnostics: Some(vec![diagnostic.clone()]),
643                edit: Some(WorkspaceEdit {
644                    changes: Some({
645                        let mut changes = HashMap::new();
646                        changes.insert(uri.clone(), associated_data.edits);
647                        changes
648                    }),
649                    ..Default::default()
650                }),
651                command: None,
652                is_preferred: Some(true),
653                disabled: None,
654                data: None,
655            };
656            actions.push(CodeActionOrCommand::CodeAction(fix_action));
657        }
658    }
659
660    let result: CodeActionResponse = actions;
661    let resp = Response {
662        id: req.id,
663        result: Some(serde_json::to_value(&result).unwrap()),
664        error: None,
665    };
666
667    connection.sender.send(Message::Response(resp))?;
668    Ok(())
669}
670
671fn publish_diagnostics(
672    connection: &Connection,
673    uri: Url,
674    version: i32,
675    diagnostics: Vec<Diagnostic>,
676) -> Result<()> {
677    let publish_params = PublishDiagnosticsParams {
678        uri,
679        diagnostics,
680        version: Some(version),
681    };
682
683    let notification = Notification {
684        method: PublishDiagnostics::METHOD.to_owned(),
685        params: serde_json::to_value(publish_params)?,
686    };
687
688    connection
689        .sender
690        .send(Message::Notification(notification))?;
691    Ok(())
692}
693
694fn handle_did_open(
695    connection: &Connection,
696    notif: lsp_server::Notification,
697    documents: &mut HashMap<Url, DocumentState>,
698) -> Result<()> {
699    let params: DidOpenTextDocumentParams = serde_json::from_value(notif.params)?;
700    let uri = params.text_document.uri;
701    let content = params.text_document.text;
702    let version = params.text_document.version;
703
704    documents.insert(uri.clone(), DocumentState { content, version });
705
706    let content = documents.get(&uri).map_or("", |doc| &doc.content);
707
708    // TODO: we need a better setup for "run func when input changed"
709    let diagnostics = lint::lint(content);
710    publish_diagnostics(connection, uri, version, diagnostics)?;
711
712    Ok(())
713}
714
715fn handle_did_change(
716    connection: &Connection,
717    notif: lsp_server::Notification,
718    documents: &mut HashMap<Url, DocumentState>,
719) -> Result<()> {
720    let params: DidChangeTextDocumentParams = serde_json::from_value(notif.params)?;
721    let uri = params.text_document.uri;
722    let version = params.text_document.version;
723
724    let Some(doc_state) = documents.get_mut(&uri) else {
725        return Ok(());
726    };
727
728    doc_state.content =
729        lsp_utils::apply_incremental_changes(&doc_state.content, params.content_changes);
730    doc_state.version = version;
731
732    let diagnostics = lint::lint(&doc_state.content);
733    publish_diagnostics(connection, uri, version, diagnostics)?;
734
735    Ok(())
736}
737
738fn handle_did_close(
739    connection: &Connection,
740    notif: lsp_server::Notification,
741    documents: &mut HashMap<Url, DocumentState>,
742) -> Result<()> {
743    let params: DidCloseTextDocumentParams = serde_json::from_value(notif.params)?;
744    let uri = params.text_document.uri;
745
746    documents.remove(&uri);
747
748    let publish_params = PublishDiagnosticsParams {
749        uri,
750        diagnostics: vec![],
751        version: None,
752    };
753
754    let notification = Notification {
755        method: PublishDiagnostics::METHOD.to_owned(),
756        params: serde_json::to_value(publish_params)?,
757    };
758
759    connection
760        .sender
761        .send(Message::Notification(notification))?;
762
763    Ok(())
764}
765
766#[derive(serde::Deserialize)]
767struct SyntaxTreeParams {
768    #[serde(rename = "textDocument")]
769    text_document: lsp_types::TextDocumentIdentifier,
770}
771
772fn handle_syntax_tree(
773    connection: &Connection,
774    req: lsp_server::Request,
775    documents: &HashMap<Url, DocumentState>,
776) -> Result<()> {
777    let params: SyntaxTreeParams = serde_json::from_value(req.params)?;
778    let uri = params.text_document.uri;
779
780    info!("Generating syntax tree for: {uri}");
781
782    let content = documents.get(&uri).map_or("", |doc| &doc.content);
783
784    let parse = SourceFile::parse(content);
785    let syntax_tree = format!("{:#?}", parse.syntax_node());
786
787    let resp = Response {
788        id: req.id,
789        result: Some(serde_json::to_value(&syntax_tree).unwrap()),
790        error: None,
791    };
792
793    connection.sender.send(Message::Response(resp))?;
794    Ok(())
795}
796
797#[derive(serde::Deserialize)]
798struct TokensParams {
799    #[serde(rename = "textDocument")]
800    text_document: lsp_types::TextDocumentIdentifier,
801}
802
803fn handle_tokens(
804    connection: &Connection,
805    req: lsp_server::Request,
806    documents: &HashMap<Url, DocumentState>,
807) -> Result<()> {
808    let params: TokensParams = serde_json::from_value(req.params)?;
809    let uri = params.text_document.uri;
810
811    info!("Generating tokens for: {uri}");
812
813    let content = documents.get(&uri).map_or("", |doc| &doc.content);
814
815    let tokens = squawk_lexer::tokenize(content);
816
817    let mut output = Vec::new();
818    let mut char_pos = 0;
819    for token in tokens {
820        let token_start = char_pos;
821        let token_end = token_start + token.len as usize;
822        let token_text = &content[token_start..token_end];
823        output.push(format!(
824            "{:?}@{}..{} {:?}",
825            token.kind, token_start, token_end, token_text
826        ));
827        char_pos = token_end;
828    }
829
830    let tokens_output = output.join("\n");
831
832    let resp = Response {
833        id: req.id,
834        result: Some(serde_json::to_value(&tokens_output).unwrap()),
835        error: None,
836    };
837
838    connection.sender.send(Message::Response(resp))?;
839    Ok(())
840}