nx_post_support/server/crates/lsp-main/src/server.rs

use anyhow::Result;
use log;
use lsp_server::{ Connection, Message, Response };
use lsp_server::Notification as ServerNotification;
use lsp_server::Request as ServerRequest;
use lsp_types as lsp;
use lsp_types::notification::{
    DidChangeTextDocument,
    DidCloseTextDocument,
    DidOpenTextDocument,
    LogMessage,
    Notification as LspNotification,
};
use lsp_types::request::{
    Completion,
    DocumentSymbolRequest,
    GotoDefinition,
    Request as LspRequest,
};
use serde_json::Value;

use analysis::{ collect_procs, ProcIndex };
use features::{ completion, definition, document_symbol };
use text::DocumentStore;

pub struct Server {
    conn: Connection,
    docs: DocumentStore,
    procs: ProcIndex,
}

impl Server {
    pub fn new(conn: Connection) -> Self {
        Self {
            conn,
            docs: DocumentStore::default(),
            procs: ProcIndex::default(),
        }
    }

    pub fn run(mut self) -> Result<()> {
        // Initialize handshake
        let (id, _params) = self.conn.initialize_start()?;
        let result = lsp::InitializeResult {
            capabilities: server_capabilities(),
            server_info: Some(lsp::ServerInfo {
                name: "tcl-lsp".into(),
                version: Some(env!("CARGO_PKG_VERSION").into()),
            }),
        };
        self.conn.initialize_finish(id, serde_json::to_value(result)?)?;

        let _ = self.log("tcl-lsp ready");

        let receiver = self.conn.receiver.clone();

        // main loop
        for msg in receiver.iter() {
            match msg {
                lsp_server::Message::Request(req) => {
                    if self.conn.handle_shutdown(&req)? {
                        break;
                    }
                    if let Err(e) = self.on_request(req) {
                        log::error!("request error: {e:#}");
                    }
                }
                lsp_server::Message::Notification(n) => {
                    if let Err(e) = self.on_notification(n) {
                        log::error!("notification error: {e:#}");
                    }
                }
                lsp_server::Message::Response(_) => {}
            }
        }

        let _ = self.log("tcl-lsp stopped");
        Ok(())
    }

    fn on_notification(&mut self, n: ServerNotification) -> Result<()> {
        match n.method.as_str() {
            DidOpenTextDocument::METHOD => {
                let params: lsp::DidOpenTextDocumentParams = serde_json::from_value(n.params)?;
                let uri = params.text_document.uri.clone();
                self.docs.open(params.text_document)?;
                self.reindex(&uri)?;
                self.log(&format!("Opened {:?}", uri))?;
            }
            DidChangeTextDocument::METHOD => {
                let params: lsp::DidChangeTextDocumentParams = serde_json::from_value(n.params)?;
                let uri = params.text_document.uri.clone();
                // FULL sync: we expect one change with whole text
                self.docs.apply_full_change(&uri, &params.content_changes)?;
                self.reindex(&uri)?;
            }
            DidCloseTextDocument::METHOD => {
                let params: lsp::DidCloseTextDocumentParams = serde_json::from_value(n.params)?;
                self.procs.remove(&params.text_document.uri);
                self.docs.close(&params.text_document.uri);
            }
            _ => {}
        }
        Ok(())
    }

    fn on_request(&mut self, req: ServerRequest) -> Result<()> {
        match req.method.as_str() {
            DocumentSymbolRequest::METHOD => {
                let params: lsp::DocumentSymbolParams = serde_json::from_value(req.params)?;
                let uri = params.text_document.uri;
                let result = document_symbol(&self.docs, &self.procs, &uri)?;
                send_ok(&self.conn, req.id, result)?;
            }
            GotoDefinition::METHOD => {
                let params: lsp::GotoDefinitionParams = serde_json::from_value(req.params)?;
                let result = definition(&self.docs, &self.procs, &params)?;
                send_ok(&self.conn, req.id, result)?;
            }
            Completion::METHOD => {
                let params: lsp::CompletionParams = serde_json::from_value(req.params)?;
                let _ = self.log("Completion");
                let result = completion(&self.docs, &self.procs, &params)?;
                send_ok(&self.conn, req.id, result)?;
            }
            other => {
                let resp = Response::new_err(
                    req.id,
                    lsp::error_codes::REQUEST_FAILED as i32,
                    format!("Unknown method: {other}")
                );
                self.conn.sender.send(Message::Response(resp))?;
            }
        }
        Ok(())
    }

    fn reindex(&mut self, uri: &lsp::Uri) -> Result<()> {
        if let Some(doc) = self.docs.get(uri) {
            if let Some(tree) = doc.tree() {
                let text = doc.text();
                let procs = collect_procs(tree, &text);
                self.procs.insert(uri.clone(), procs);
            }
        }
        Ok(())
    }

    fn log(&self, message: &str) -> Result<()> {
        let notif = LogMessage::METHOD.to_string();
        let params = lsp::LogMessageParams {
            typ: lsp::MessageType::INFO,
            message: message.to_string(),
        };
        self.conn.sender.send(Message::Notification(ServerNotification::new(notif, params)))?;
        Ok(())
    }
}

fn send_ok<T: serde::Serialize>(
    conn: &Connection,
    id: lsp_server::RequestId,
    result: T
) -> Result<()> {
    let v: Value = serde_json::to_value(result)?;
    conn.sender.send(Message::Response(Response::new_ok(id, v)))?;
    Ok(())
}

fn server_capabilities() -> lsp::ServerCapabilities {
    lsp::ServerCapabilities {
        completion_provider: Some(lsp::CompletionOptions {
            ..Default::default()
        }),
        ..Default::default()
    }
}