roc/crates/lang_srv/src/server.rs

use analysis::HIGHLIGHT_TOKENS_LEGEND;

use log::{debug, trace};
use registry::Registry;
use std::future::Future;
use std::time::Duration;

use tower_lsp::jsonrpc::Result;
use tower_lsp::lsp_types::*;
use tower_lsp::{Client, LanguageServer, LspService, Server};

use crate::analysis::{global_analysis, DocInfo};

mod analysis;
mod convert;
mod registry;

#[derive(Debug)]
struct RocServer {
    pub state: RocServerState,
    client: Client,
}

///This exists so we can test most of RocLs without anything LSP related
#[derive(Debug)]
struct RocServerState {
    registry: Registry,
}

impl std::panic::RefUnwindSafe for RocServer {}

impl RocServer {
    pub fn new(client: Client) -> Self {
        Self {
            state: RocServerState::new(),
            client,
        }
    }
    pub fn capabilities() -> ServerCapabilities {
        let text_document_sync = TextDocumentSyncCapability::Options(
            // TODO: later on make this incremental
            TextDocumentSyncOptions {
                open_close: Some(true),
                change: Some(TextDocumentSyncKind::FULL),
                ..TextDocumentSyncOptions::default()
            },
        );
        let hover_provider = HoverProviderCapability::Simple(true);
        let definition_provider = DefinitionOptions {
            work_done_progress_options: WorkDoneProgressOptions {
                work_done_progress: None,
            },
        };
        let document_formatting_provider = DocumentFormattingOptions {
            work_done_progress_options: WorkDoneProgressOptions {
                work_done_progress: None,
            },
        };
        let semantic_tokens_provider =
            SemanticTokensServerCapabilities::SemanticTokensOptions(SemanticTokensOptions {
                work_done_progress_options: WorkDoneProgressOptions {
                    work_done_progress: None,
                },
                legend: SemanticTokensLegend {
                    token_types: HIGHLIGHT_TOKENS_LEGEND.into(),
                    token_modifiers: vec![],
                },
                range: None,
                full: Some(SemanticTokensFullOptions::Bool(true)),
            });
        let completion_provider = CompletionOptions {
            resolve_provider: Some(false),
            trigger_characters: Some(vec![".".to_string()]),
            //TODO: what is this?
            all_commit_characters: None,
            work_done_progress_options: WorkDoneProgressOptions {
                work_done_progress: None,
            },
        };
        ServerCapabilities {
            text_document_sync: Some(text_document_sync),
            hover_provider: Some(hover_provider),
            definition_provider: Some(OneOf::Right(definition_provider)),
            document_formatting_provider: Some(OneOf::Right(document_formatting_provider)),
            semantic_tokens_provider: Some(semantic_tokens_provider),
            completion_provider: Some(completion_provider),
            ..ServerCapabilities::default()
        }
    }

    /// Records a document content change.
    async fn change(&self, fi: Url, text: String, version: i32) {
        let updating_result = self.state.change(&fi, text, version).await;

        //The analysis task can be cancelled by another change coming in which will update the watched variable
        if let Err(e) = updating_result {
            debug!("cancelled change. Reason:{:?}", e);
            return;
        }
        debug!("applied_change getting and returning diagnostics");

        let diagnostics = self.state.registry.diagnostics(&fi).await;

        self.client
            .publish_diagnostics(fi, diagnostics, Some(version))
            .await;
    }
}

impl RocServerState {
    pub fn new() -> RocServerState {
        Self {
            registry: Registry::default(),
        }
    }

    async fn registry(&self) -> &Registry {
        &self.registry
    }

    async fn close(&self, _fi: Url) {}

    pub async fn change(
        &self,
        fi: &Url,
        text: String,
        version: i32,
    ) -> std::result::Result<(), String> {
        debug!("V{:?}:starting change", version);
        let doc_info = DocInfo::new(fi.clone(), text, version);

        self.registry
            .apply_doc_info_changes(fi.clone(), doc_info.clone())
            .await;

        debug!(
            "V{:?}:finished updating docinfo, starting analysis ",
            version
        );

        let inner_ref = self;
        let updating_result = async {
            //This reduces wasted computation by waiting to allow a new change to come in and update the version before we check, but does delay the final analysis. Ideally this would be replaced with cancelling the analysis when a new one comes in.
            let debounce_ms = std::env::var("ROCLS_DEBOUNCE")
                .map_or(100, |a| str::parse::<u64>(&a).unwrap_or(100));
            tokio::time::sleep(Duration::from_millis(debounce_ms)).await;
            let is_latest = inner_ref
                .registry
                .get_latest_version(fi)
                .await
                .map(|latest| latest == version)
                .unwrap_or(true);
            if !is_latest {
                return Err("Not latest version skipping analysis".to_string());
            }

            let results = match tokio::task::spawn_blocking(|| global_analysis(doc_info)).await {
                Err(e) => return Err(format!("Document analysis failed. reason:{:?}", e)),
                Ok(a) => a,
            };
            let latest_version = inner_ref.registry.get_latest_version(fi).await;

            //if this version is not the latest another change must have come in and this analysis is useless
            //if there is no older version we can just proceed with the update
            if let Some(latest_version) = latest_version {
                if latest_version != version {
                    return Err(format!(
                        "version {0} doesn't match latest: {1} discarding analysis  ",
                        version, latest_version
                    ));
                }
            }
            debug!(
                "V{:?}:finished document analysis applying changes ",
                version
            );

            inner_ref.registry.apply_changes(results, fi.clone()).await;
            Ok(())
        }
        .await;
        debug!("V{:?}:finished document change process", version);
        updating_result
    }
}

#[tower_lsp::async_trait]
impl LanguageServer for RocServer {
    async fn initialize(&self, _: InitializeParams) -> Result<InitializeResult> {
        Ok(InitializeResult {
            capabilities: Self::capabilities(),
            ..InitializeResult::default()
        })
    }

    async fn initialized(&self, _: InitializedParams) {
        self.client
            .log_message(MessageType::INFO, "Roc language server initialized.")
            .await;
    }

    async fn did_open(&self, params: DidOpenTextDocumentParams) {
        let TextDocumentItem {
            uri, text, version, ..
        } = params.text_document;
        self.change(uri, text, version).await;
    }

    async fn did_change(&self, params: DidChangeTextDocumentParams) {
        let VersionedTextDocumentIdentifier { uri, version, .. } = params.text_document;

        // NOTE: We specify that we expect full-content syncs in the server capabilities,
        // so here we assume the only change passed is a change of the entire document's content.
        let TextDocumentContentChangeEvent { text, .. } =
            params.content_changes.into_iter().next().unwrap();

        trace!("got did_change");
        self.change(uri, text, version).await;
    }

    async fn did_close(&self, params: DidCloseTextDocumentParams) {
        let TextDocumentIdentifier { uri } = params.text_document;
        self.state.close(uri).await;
    }

    async fn shutdown(&self) -> Result<()> {
        Ok(())
    }

    async fn hover(&self, params: HoverParams) -> Result<Option<Hover>> {
        let HoverParams {
            text_document_position_params:
                TextDocumentPositionParams {
                    text_document,
                    position,
                },
            work_done_progress_params: _,
        } = params;

        panic_wrapper_async(|| async {
            self.state
                .registry
                .hover(&text_document.uri, position)
                .await
        })
        .await
    }

    async fn goto_definition(
        &self,
        params: GotoDefinitionParams,
    ) -> Result<Option<GotoDefinitionResponse>> {
        let GotoDefinitionParams {
            text_document_position_params:
                TextDocumentPositionParams {
                    text_document,
                    position,
                },
            work_done_progress_params: _,
            partial_result_params: _,
        } = params;

        panic_wrapper_async(|| async {
            self.state
                .registry()
                .await
                .goto_definition(&text_document.uri, position)
                .await
        })
        .await
    }

    async fn formatting(&self, params: DocumentFormattingParams) -> Result<Option<Vec<TextEdit>>> {
        let DocumentFormattingParams {
            text_document,
            options: _,
            work_done_progress_params: _,
        } = params;

        panic_wrapper_async(|| async {
            self.state
                .registry()
                .await
                .formatting(&text_document.uri)
                .await
        })
        .await
    }

    async fn semantic_tokens_full(
        &self,
        params: SemanticTokensParams,
    ) -> Result<Option<SemanticTokensResult>> {
        let SemanticTokensParams {
            text_document,
            work_done_progress_params: _,
            partial_result_params: _,
        } = params;

        panic_wrapper_async(|| async {
            self.state
                .registry()
                .await
                .semantic_tokens(&text_document.uri)
                .await
        })
        .await
    }

    async fn completion(&self, params: CompletionParams) -> Result<Option<CompletionResponse>> {
        let doc = params.text_document_position;
        trace!("got completion request");

        panic_wrapper_async(|| async {
            self.state
                .registry
                .completion_items(&doc.text_document.uri, doc.position)
                .await
        })
        .await
    }
}

async fn panic_wrapper_async<Fut, T>(
    f: impl FnOnce() -> Fut + std::panic::UnwindSafe,
) -> Result<Option<T>>
where
    Fut: Future<Output = Option<T>>,
{
    match std::panic::catch_unwind(f) {
        Ok(r) => Ok(r.await),
        Err(_) => Err(tower_lsp::jsonrpc::Error::internal_error()),
    }
}

#[tokio::main]
async fn main() {
    env_logger::Builder::from_env("ROCLS_LOG").init();

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

    let (service, socket) = LspService::new(RocServer::new);
    Server::new(stdin, stdout, socket).serve(service).await;
}

#[cfg(test)]
mod tests {
    use std::sync::Once;

    use expect_test::expect;
    use indoc::indoc;
    use log::info;

    use super::*;

    fn completion_resp_to_labels(resp: CompletionResponse) -> Vec<String> {
        match resp {
            CompletionResponse::Array(list) => list.into_iter(),
            CompletionResponse::List(list) => list.items.into_iter(),
        }
        .map(|item| item.label)
        .collect::<Vec<_>>()
    }
    ///Gets completion and returns only the label for each completion
    async fn get_completion_labels(
        reg: &Registry,
        url: &Url,
        position: Position,
    ) -> Option<Vec<String>> {
        reg.completion_items(url, position)
            .await
            .map(completion_resp_to_labels)
    }

    const DOC_LIT: &str = indoc! {r#"
        interface Test
          exposes []
          imports []
        "#};
    static INIT: Once = Once::new();
    async fn test_setup(doc: String) -> (RocServerState, Url) {
        INIT.call_once(|| {
            env_logger::builder()
                .is_test(true)
                .filter_level(log::LevelFilter::Trace)
                .init();
        });
        info!("doc is:\n{0}", doc);
        let url = Url::parse("file:/Test.roc").unwrap();

        let inner = RocServerState::new();
        //setup the file
        inner.change(&url, doc, 0).await.unwrap();
        (inner, url)
    }
    ///Test that completion works properly when we apply an "as" pattern to an identifier
    #[tokio::test]
    async fn test_completion_as_identifier() {
        let suffix = DOC_LIT.to_string()
            + indoc! {r#"
            main =
              when a is
                inn as outer -> 
                  "#};
        let (inner, url) = test_setup(suffix.clone()).await;
        let position = Position::new(6, 7);
        let reg = &inner.registry;

        let change = suffix.clone() + "o";
        inner.change(&url, change, 1).await.unwrap();
        let comp1 = get_completion_labels(reg, &url, position).await;

        let c = suffix.clone() + "i";
        inner.change(&url, c, 2).await.unwrap();
        let comp2 = get_completion_labels(reg, &url, position).await;

        let actual = [comp1, comp2];
        expect![[r#"
            [
                Some(
                    [
                        "outer",
                    ],
                ),
                Some(
                    [
                        "inn",
                        "outer",
                    ],
                ),
            ]
        "#]]
        .assert_debug_eq(&actual)
    }

    ///Test that completion works properly when we apply an "as" pattern to a record
    #[tokio::test]
    async fn test_completion_as_record() {
        let doc = DOC_LIT.to_string()
            + indoc! {r#"
            main =
              when a is
                {one,two} as outer -> 
                  "#};

        let (inner, url) = test_setup(doc.clone()).await;
        let position = Position::new(6, 7);
        let reg = &inner.registry;

        let change = doc.clone() + "o";
        inner.change(&url, change, 1).await.unwrap();
        let comp1 = get_completion_labels(reg, &url, position).await;

        let c = doc.clone() + "t";
        inner.change(&url, c, 2).await.unwrap();
        let comp2 = get_completion_labels(reg, &url, position).await;
        let actual = [comp1, comp2];

        expect![[r#"
            [
                Some(
                    [
                        "one",
                        "two",
                        "outer",
                    ],
                ),
                Some(
                    [
                        "one",
                        "two",
                        "outer",
                    ],
                ),
            ]
        "#]]
        .assert_debug_eq(&actual);
    }
}