aster/aster-lsp/src/server.rs

use std::collections::HashMap;

use aster_core::lexer::Lexer;
use aster_core::parser::Parser;
use aster_core::error::RuntimeError;

use lsp_server::{Connection, Message, Notification};
use lsp_types::{
    Diagnostic, DiagnosticSeverity, DidChangeTextDocumentParams,
    DidCloseTextDocumentParams, DidOpenTextDocumentParams, InitializeParams,
    Position, PublishDiagnosticsParams, Range, ServerCapabilities, TextDocumentSyncCapability,
    TextDocumentSyncKind, Uri,
};

use crate::completion;
use crate::goto_def;
use crate::hover;
use crate::references;
use crate::rename;
use crate::signature;

pub fn run() {
    eprintln!("Aster LSP server starting...");

    let (connection, io_threads) = Connection::stdio();

    let capabilities = ServerCapabilities {
        text_document_sync: Some(TextDocumentSyncCapability::Kind(
            TextDocumentSyncKind::FULL,
        )),
        completion_provider: Some(lsp_types::CompletionOptions {
            trigger_characters: Some(
                (b'a'..=b'z').chain(b'A'..=b'Z')
                    .map(|c| (c as char).to_string())
                    .chain(std::iter::once(".".to_string()))
                    .chain(std::iter::once("_".to_string()))
                    .collect()
            ),
            resolve_provider: Some(false),
            ..Default::default()
        }),
        hover_provider: Some(lsp_types::HoverProviderCapability::Simple(true)),
        signature_help_provider: Some(lsp_types::SignatureHelpOptions::default()),
        definition_provider: Some(lsp_types::OneOf::Left(true)),
        references_provider: Some(lsp_types::OneOf::Left(true)),
        rename_provider: Some(lsp_types::OneOf::Left(true)),
        ..ServerCapabilities::default()
    };

    let init_result = connection
        .initialize(serde_json::to_value(&capabilities).unwrap())
        .unwrap();
    let _init_params: InitializeParams = serde_json::from_value(init_result).unwrap();

    eprintln!("Aster LSP initialized successfully");

    let mut documents: HashMap<Uri, String> = HashMap::new();

    for msg in &connection.receiver {
        match msg {
            Message::Request(req) => {
                if connection.handle_shutdown(&req).unwrap() {
                    break;
                }
                handle_request(&documents, &connection, req);
            }
            Message::Notification(notif) => {
                handle_notification(&mut documents, &connection, notif);
            }
            Message::Response(_resp) => {}
        }
    }

    io_threads.join().unwrap();
    eprintln!("Aster LSP server stopped.");
}

fn handle_request(
    documents: &HashMap<Uri, String>,
    connection: &Connection,
    req: lsp_server::Request,
) {
    match req.method.as_str() {
        "textDocument/completion" => {
            completion::handle_completion(documents, connection, req);
        }
        "textDocument/hover" => {
            hover::handle_hover(documents, connection, req);
        }
        "textDocument/signatureHelp" => {
            signature::handle_signature_help(documents, connection, req);
        }
        "textDocument/definition" => {
            goto_def::handle_goto_definition(documents, connection, req);
        }
        "textDocument/references" => {
            references::handle_references(documents, connection, req);
        }
        "textDocument/rename" => {
            rename::handle_rename(documents, connection, req);
        }
        _ => {
            eprintln!("Unhandled request: {}", req.method);
        }
    }
}

fn handle_notification(
    documents: &mut HashMap<Uri, String>,
    connection: &Connection,
    notif: Notification,
) {
    match notif.method.as_str() {
        "textDocument/didOpen" => {
            let params: DidOpenTextDocumentParams =
                serde_json::from_value(notif.params).unwrap();
            let uri = params.text_document.uri.clone();
            let text = params.text_document.text.clone();
            documents.insert(uri.clone(), text.clone());
            publish_diagnostics(connection, &uri, &text);
        }
        "textDocument/didChange" => {
            let params: DidChangeTextDocumentParams =
                serde_json::from_value(notif.params).unwrap();
            let uri = params.text_document.uri.clone();
            if let Some(change) = params.content_changes.into_iter().last() {
                documents.insert(uri.clone(), change.text.clone());
                publish_diagnostics(connection, &uri, &change.text);
            }
        }
        "textDocument/didClose" => {
            let params: DidCloseTextDocumentParams =
                serde_json::from_value(notif.params).unwrap();
            documents.remove(&params.text_document.uri);
            let clear = PublishDiagnosticsParams {
                uri: params.text_document.uri.clone(),
                diagnostics: vec![],
                version: None,
            };
            send_notification(connection, "textDocument/publishDiagnostics", clear);
        }
        _ => {}
    }
}

fn publish_diagnostics(connection: &Connection, uri: &Uri, text: &str) {
    let (tokens, lex_errors) = Lexer::new(text).tokenize();
    let mut errors = lex_errors;

    let mut parser = Parser::new(tokens);
    let (_stmts, parse_errors) = parser.parse();
    errors.extend(parse_errors);

    let diagnostics: Vec<Diagnostic> = errors
        .iter()
        .filter_map(error_to_diagnostic)
        .collect();

    let params = PublishDiagnosticsParams {
        uri: uri.clone(),
        diagnostics,
        version: None,
    };

    send_notification(connection, "textDocument/publishDiagnostics", params);
}

fn error_to_diagnostic(error: &RuntimeError) -> Option<Diagnostic> {
    match error {
        RuntimeError::LexError { message, line, column } => {
            let l = (*line as u32).saturating_sub(1);
            let c = (*column as u32).saturating_sub(1);
            Some(Diagnostic {
                range: Range {
                    start: Position { line: l, character: c },
                    end: Position { line: l, character: c.saturating_add(1) },
                },
                severity: Some(DiagnosticSeverity::ERROR),
                message: format!("[Lex] {}", message),
                ..Diagnostic::default()
            })
        }
        RuntimeError::ParseError { message, token } => {
            let l = (token.line as u32).saturating_sub(1);
            let c = (token.column as u32).saturating_sub(1);
            Some(Diagnostic {
                range: Range {
                    start: Position { line: l, character: c },
                    end: Position { line: l, character: c.saturating_add(1) },
                },
                severity: Some(DiagnosticSeverity::ERROR),
                message: format!("[Parse] {}", message),
                ..Diagnostic::default()
            })
        }
        RuntimeError::RuntimeError { .. } => None,
    }
}

fn send_notification(connection: &Connection, method: &str, params: impl serde::Serialize) {
    let notification = Notification::new(method.to_string(), serde_json::to_value(params).unwrap());
    let _ = connection.sender.send(Message::Notification(notification));
}