math-lsp/src/main.rs

use std::collections::HashMap;

use anyhow::Error;
use lsp_server::{Connection, Message, Response};
use lsp_types::{
    DidChangeTextDocumentParams, DidCloseTextDocumentParams, DidOpenTextDocumentParams, InlayHint,
    InlayHintParams, Position, TextDocumentSyncCapability, TextDocumentSyncKind,
};
use lsp_types::{InitializeParams, ServerCapabilities};
use numbat::module_importer::BuiltinModuleImporter;
use numbat::resolver::CodeSource;
use numbat::InterpreterResult;

fn main() -> Result<(), Error> {
    let (connection, io_threads) = Connection::stdio();

    let server_capabilities = serde_json::to_value(ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Kind(TextDocumentSyncKind::FULL)),
        inlay_hint_provider: Some(lsp_types::OneOf::Left(true)),
        ..Default::default()
    })?;

    let initialization_params = connection.initialize(server_capabilities)?;
    main_loop(connection, initialization_params)?;
    io_threads.join()?;

    Ok(())
}

fn main_loop(connection: Connection, params: serde_json::Value) -> Result<(), Error> {
    let _params: InitializeParams = serde_json::from_value(params).unwrap();

    let mut docs: HashMap<lsp_types::Url, Vec<String>> = HashMap::new();
    let mut context = {
        let importer = BuiltinModuleImporter::default();
        numbat::Context::new(importer)
    };

    let mut settings = numbat::InterpreterSettings {
        print_fn: Box::new(|_| ()),
    };

    let _ = context
        .interpret_with_settings(&mut settings, "use prelude", CodeSource::Text)
        .expect("failed to import prelude");

    for msg in &connection.receiver {
        match msg {
            Message::Request(req) => {
                if connection.handle_shutdown(&req)? {
                    return Ok(());
                }

                if req.method == "textDocument/inlayHint" {
                    let params: InlayHintParams = serde_json::from_value(req.params)?;
                    let doc = docs.get(&params.text_document.uri).unwrap();
                    let mut hints = vec![];

                    for line_number in params.range.start.line..params.range.end.line {
                        let line = &doc[line_number as usize];

                        if !line.chars().next().is_some_and(char::is_whitespace) {
                            continue;
                        }

                        let result =
                            context.interpret_with_settings(&mut settings, line, CodeSource::Text);

                        let label = match result {
                            Ok((_, result)) => match result {
                                InterpreterResult::Value(value) => Some(format!("= {}", value)),
                                InterpreterResult::Continue => None,
                            },
                            Err(numbat::NumbatError::NameResolutionError(err)) => {
                                Some(format!("-> {}", err))
                            }
                            Err(numbat::NumbatError::ResolverError(err)) => {
                                Some(format!("-> {}", err))
                            }
                            Err(numbat::NumbatError::TypeCheckError(err)) => {
                                Some(format!("-> {}", err))
                            }
                            Err(numbat::NumbatError::RuntimeError(err)) => {
                                Some(format!("-> {}", err))
                            }
                        };

                        if let Some(label) = label {
                            let hint = InlayHint {
                                position: Position {
                                    line: line_number,
                                    character: line.len() as u32,
                                },
                                label: lsp_types::InlayHintLabel::String(label),
                                kind: None,
                                text_edits: None,
                                tooltip: None,
                                data: None,
                                padding_left: Some(true),
                                padding_right: None,
                            };

                            hints.push(hint);
                        }
                    }

                    let hints = Response {
                        id: req.id,
                        result: Some(serde_json::to_value(hints)?),
                        error: None,
                    };

                    connection.sender.send(Message::Response(hints))?;
                }
            }
            Message::Response(_) => {}
            Message::Notification(not) => match not.method.as_str() {
                "textDocument/didChange" => {
                    let params: DidChangeTextDocumentParams = serde_json::from_value(not.params)?;
                    docs.insert(
                        params.text_document.uri,
                        params.content_changes[0]
                            .text
                            .split('\n')
                            .map(String::from)
                            .collect(),
                    );
                }
                "textDocument/didOpen" => {
                    let params: DidOpenTextDocumentParams = serde_json::from_value(not.params)?;
                    docs.insert(
                        params.text_document.uri,
                        params
                            .text_document
                            .text
                            .split('\n')
                            .map(String::from)
                            .collect(),
                    );
                }
                "textDocument/didClose" => {
                    let params: DidCloseTextDocumentParams = serde_json::from_value(not.params)?;
                    docs.remove(&params.text_document.uri);
                }
                _ => {}
            },
        }
    }
    Ok(())
}