plum

// LSP server for plum. Two diagnostic sources:

//

//   1. Syntax errors — read directly off the tree-sitter parse tree (ERROR /

//      MISSING nodes), which carries real byte/line/column spans.

//

//   2. Type errors — from `plum_checker::checkSource`. `CheckError` carries a

//      `line`/`column` pointing at the enclosing `fn`/`enum`/`test`

//      declaration (not the exact sub-expression/token within it — the

//      checker doesn't track spans that fine yet), so the diagnostic range

//      covers that declaration's line rather than the whole document.

//

// Both are republished on every `did_open`/`did_change`/`did_save`. Type

// checking goes through `plum_core::loadAndMergeWithOverrides`, passing the

// live in-memory buffer for the edited document while every other

// (transitively imported) file still reads from disk — so it reflects

// unsaved edits, not just the last save.

//

// Completion (see `completion.rs`) is a separate, coarser-grained feature:

// keywords + top-level names always; `expr.` member completion and

// not-yet-imported std lib symbols only when the document currently loads

// cleanly enough to build the checker's global tables.

#![allow(non_snake_case)]

use std::collections::HashMap;

use std::path::PathBuf;

use std::sync::OnceLock;

use tokio::sync::Mutex;

use tower_lsp::jsonrpc::Result as RpcResult;

use tower_lsp::lsp_types::*;

use tower_lsp::{Client, LanguageServer, LspService, Server};

mod completion;

struct Backend {

    client: Client,

    docs: Mutex<HashMap<Url, String>>,

    libPath: PathBuf,

    /// Lazily built on first completion request and kept for the server's

    /// lifetime — the std lib doesn't change while `plum-lsp` is running.

    stdIndex: OnceLock<HashMap<String, String>>,

}

impl Backend {

    fn parseTree(source: &str) -> tree_sitter::Tree {

        let mut parser = tree_sitter::Parser::new();

        parser

            .set_language(&tree_sitter_plum::LANGUAGE.into())

            .expect("failed to load plum grammar");

        parser

            .parse(source, None)

            .expect("tree-sitter parse returned None")

    }

    /// Real-span diagnostics straight from the parse tree: any ERROR or

    /// MISSING node becomes one diagnostic at its own range.

    fn syntaxDiagnostics(tree: &tree_sitter::Tree) -> Vec<Diagnostic> {

        let mut out = Vec::new();

        let mut cursor = tree.walk();

        let mut visit = |node: tree_sitter::Node| {

            if node.is_error() || node.is_missing() {

                let start = node.start_position();

                let end = node.end_position();

                out.push(Diagnostic {

                    range: Range {

                        start: Position::new(start.row as u32, start.column as u32),

                        end: Position::new(end.row as u32, end.column as u32),

                    },

                    severity: Some(DiagnosticSeverity::ERROR),

                    source: Some("plum-syntax".into()),

                    message: if node.is_missing() {

                        format!("missing {}", node.kind())

                    } else {

                        "syntax error".into()

                    },

                    ..Default::default()

                });

            }

        };

        // Manual depth-first walk since `TreeCursor` has no built-in visitor.

        loop {

            visit(cursor.node());

            if cursor.goto_first_child() {

                continue;

            }

            loop {

                if cursor.goto_next_sibling() {

                    break;

                }

                if !cursor.goto_parent() {

                    return out;

                }

            }

        }

    }

    /// Diagnostics from `plum_checker::checkSource`, run against the live

    /// in-memory buffer (via `loadAndMergeWithOverrides`) rather than

    /// whatever's on disk, so an unsaved edit is reflected immediately.

    /// Every other (transitively imported) file still resolves from disk.

    /// Returns `None` (rather than an error) if the document has no on-disk

    /// path or fails to load — those aren't type errors, so we don't want to

    /// flash a bogus diagnostic.

    fn typeDiagnostics(&self, uri: &Url, text: &str) -> Option<Vec<Diagnostic>> {

        let path = uri.to_file_path().ok()?;

        let canon = std::fs::canonicalize(&path).unwrap_or(path.clone());

        let mut overrides = HashMap::new();

        overrides.insert(canon, text.to_string());

        let ast = plum_core::loadAndMergeWithOverrides(&path, &self.libPath, &overrides).ok()?;

        let lastLine = text.lines().count().max(1) as u32 - 1;

        match plum_checker::checkSource(&ast) {

            Ok(()) => Some(Vec::new()),

            Err(errors) => Some(

                errors

                    .into_iter()

                    .map(|e| {

                        // `CheckError`'s position is 1-based and points at the

                        // enclosing declaration, not a sub-expression — clamp

                        // to the document and cover the whole line so an

                        // editor still shows something sensible even if the

                        // line itself is short/blank.

                        let line = e.line.saturating_sub(1).min(lastLine as usize) as u32;

                        Diagnostic {

                            range: Range {

                                start: Position::new(line, 0),

                                end: Position::new(line, u32::MAX),

                            },

                            severity: Some(DiagnosticSeverity::ERROR),

                            source: Some("plum-check".into()),

                            message: e.message,

                            ..Default::default()

                        }

                    })

                    .collect(),

            ),

        }

    }

    /// Builds `plum_checker`'s global tables for the live buffer, the same

    /// way `typeDiagnostics` does — used by member/std-lib completion, which

    /// both need method/field/import-target lookups, not just the document's

    /// own AST.

    fn mergedSource(&self, uri: &Url, text: &str) -> Option<plum_core::ast::Source> {

        let path = uri.to_file_path().ok()?;

        let canon = std::fs::canonicalize(&path).unwrap_or(path.clone());

        let mut overrides = HashMap::new();

        overrides.insert(canon, text.to_string());

        plum_core::loadAndMergeWithOverrides(&path, &self.libPath, &overrides).ok()

    }

    async fn completionItems(&self, uri: &Url, pos: Position) -> Vec<CompletionItem> {

        let text = match self.docs.lock().await.get(uri) {

            Some(t) => t.clone(),

            None => return Vec::new(),

        };

        let line_text = text.lines().nth(pos.line as usize).unwrap_or("");

        let sanitized = completion::sanitizeForParsing(&text, pos);

        let tree = Self::parseTree(&sanitized);

        let own_source = plum_core::AstParser::new(&sanitized).parseSource(tree.root_node());

        if let Some(receiver) = completion::receiverBeforeDot(line_text, pos.character as usize) {

            let Some(merged) = self.mergedSource(uri, &sanitized) else {

                return Vec::new();

            };

            let (global_env, methods, enum_variants, enum_params) =

                plum_checker::buildGlobalTables(&merged);

            let min_required = plum_checker::buildMinRequiredArgs(&merged);

            let ctx = plum_checker::CheckCtx {

                methods: &methods,

                enum_variants: &enum_variants,

                enum_params: &enum_params,

                min_required: &min_required,

            };

            let cursor_fn = completion::enclosingFn(&own_source, pos.line as usize + 1);

            return completion::memberItems(cursor_fn, &receiver, &global_env, &ctx);

        }

        let mut items = completion::keywordItems();

        let mut seen: std::collections::HashSet<String> = std::collections::HashSet::new();

        let pushUnique = |item: CompletionItem,

                          items: &mut Vec<CompletionItem>,

                          seen: &mut std::collections::HashSet<String>| {

            if seen.insert(item.label.clone()) {

                items.push(item);

            }

        };

        // The merged source (own file + imports) if it loads cleanly, else

        // just the document's own declarations — still useful mid-edit.

        let base_source = self.mergedSource(uri, &sanitized);

        for item in completion::topLevelItems(base_source.as_ref().unwrap_or(&own_source)) {

            pushUnique(item, &mut items, &mut seen);

        }

        let stdIndex = self

            .stdIndex

            .get_or_init(|| completion::stdLibIndex(&self.libPath));

        for (name, import_path) in stdIndex {

            if seen.contains(name) {

                continue;

            }

            let already_imported = own_source.imports.iter().any(|i| &i.path == import_path);

            if already_imported {

                continue;

            }

            items.push(CompletionItem {

                label: name.clone(),

                kind: Some(CompletionItemKind::CLASS),

                detail: Some(format!("{import_path} (not imported)")),

                additional_text_edits: Some(vec![completion::importInsertionEdit(

                    &text,

                    import_path,

                )]),

                ..Default::default()

            });

        }

        items

    }

    async fn publishAll(&self, uri: Url) {

        let text = match self.docs.lock().await.get(&uri) {

            Some(t) => t.clone(),

            None => return,

        };

        let tree = Self::parseTree(&text);

        let mut diagnostics = Self::syntaxDiagnostics(&tree);

        if diagnostics.is_empty() {

            if let Some(mut typeErrors) = self.typeDiagnostics(&uri, &text) {

                diagnostics.append(&mut typeErrors);

            }

        }

        self.client

            .publish_diagnostics(uri, diagnostics, None)

            .await;

    }

}

#[tower_lsp::async_trait]

impl LanguageServer for Backend {

    async fn initialize(&self, _: InitializeParams) -> RpcResult<InitializeResult> {

        Ok(InitializeResult {

            capabilities: ServerCapabilities {

                text_document_sync: Some(TextDocumentSyncCapability::Kind(

                    TextDocumentSyncKind::FULL,

                )),

                completion_provider: Some(CompletionOptions {

                    trigger_characters: Some(vec![".".to_string()]),

                    ..Default::default()

                }),

                ..Default::default()

            },

            server_info: Some(ServerInfo {

                name: "plum-lsp".into(),

                version: Some(env!("CARGO_PKG_VERSION").into()),

            }),

        })

    }

    async fn initialized(&self, _: InitializedParams) {

        self.client

            .log_message(MessageType::INFO, "plum-lsp initialized")

            .await;

    }

    async fn shutdown(&self) -> RpcResult<()> {

        Ok(())

    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {

        let uri = params.text_document.uri;

        self.docs

            .lock()

            .await

            .insert(uri.clone(), params.text_document.text);

        self.publishAll(uri).await;

    }

    async fn did_change(&self, mut params: DidChangeTextDocumentParams) {

        let uri = params.text_document.uri;

        // FULL sync: the last content change carries the whole new text.

        if let Some(change) = params.content_changes.pop() {

            self.docs.lock().await.insert(uri.clone(), change.text);

        }

        self.publishAll(uri).await;

    }

    async fn did_save(&self, params: DidSaveTextDocumentParams) {

        self.publishAll(params.text_document.uri).await;

    }

    async fn did_close(&self, params: DidCloseTextDocumentParams) {

        self.docs.lock().await.remove(&params.text_document.uri);

    }

    async fn completion(&self, params: CompletionParams) -> RpcResult<Option<CompletionResponse>> {

        let uri = params.text_document_position.text_document.uri;

        let pos = params.text_document_position.position;

        Ok(Some(CompletionResponse::Array(

            self.completionItems(&uri, pos).await,

        )))

    }

}

#[tokio::main]

async fn main() {

    // `--lib-path` mirrors `plum-cli`'s own required flag (see

    // plum-cli/src/main.rs's `Command::Run`/`Compile`) — the checker has no

    // built-in default location for the std lib. It's the directory

    // CONTAINING `plum-std/` (an `import std/Bool` resolves to

    // `<lib_path>/plum-std/Bool.plum` — see `plum_core::loader::loadImport`),

    // not `plum-std` itself; the repo root is `.` when run from a checkout.

    let libPath = std::env::args()

        .collect::<Vec<_>>()

        .windows(2)

        .find(|w| w[0] == "--lib-path")

        .map(|w| PathBuf::from(&w[1]))

        .unwrap_or_else(|| PathBuf::from("."));

    let stdin = tokio::io::stdin();

    let stdout = tokio::io::stdout();

    let (service, socket) = LspService::new(|client| Backend {

        client,

        docs: Mutex::new(HashMap::new()),

        libPath,

        stdIndex: OnceLock::new(),

    });

    Server::new(stdin, stdout, socket).serve(service).await;

}