feat(lsp): load all files at start and improve state updates

Scan the project during initialization and add all found files to the
state. Also, only remove files from the state when they're deleted, not
when closed.

crates/stef-lsp/Cargo.toml

@@ -13,6 +13,7 @@ license.workspace = true
 anyhow = "1.0.75"
 clap.workspace = true
 directories = "5.0.1"
+ignore = "0.4.21"
 line-index = "0.1.1"
 log = { version = "0.4.20", features = ["kv_unstable_std", "std"] }
 lsp-server = "0.7.5"

crates/stef-lsp/src/handlers/mod.rs

@@ -4,27 +4,58 @@ use anyhow::{Context, Result};
 use line_index::{LineIndex, TextRange};
 use log::{as_debug, as_display, debug, error, warn};
 use lsp_types::{
-    DidChangeConfigurationParams, DidChangeTextDocumentParams, DidCloseTextDocumentParams,
-    DidOpenTextDocumentParams, DocumentSymbolParams, DocumentSymbolResponse, Hover, HoverContents,
-    HoverParams, HoverProviderCapability, InitializeParams, InitializeResult, InitializedParams,
-    MarkupContent, MarkupKind, OneOf, PositionEncodingKind, Registration, SemanticTokens,
-    SemanticTokensFullOptions, SemanticTokensLegend, SemanticTokensOptions, SemanticTokensParams,
-    SemanticTokensResult, SemanticTokensServerCapabilities, ServerCapabilities, ServerInfo,
-    TextDocumentSyncCapability, TextDocumentSyncKind, WorkDoneProgressOptions,
+    DeleteFilesParams, DidChangeConfigurationParams, DidChangeTextDocumentParams,
+    DidCloseTextDocumentParams, DidOpenTextDocumentParams, DocumentSymbolParams,
+    DocumentSymbolResponse, FileOperationFilter, FileOperationPattern, FileOperationPatternKind,
+    FileOperationRegistrationOptions, Hover, HoverContents, HoverParams, HoverProviderCapability,
+    InitializeParams, InitializeResult, InitializedParams, MarkupContent, MarkupKind, OneOf,
+    PositionEncodingKind, Registration, SemanticTokens, SemanticTokensFullOptions,
+    SemanticTokensLegend, SemanticTokensOptions, SemanticTokensParams, SemanticTokensResult,
+    SemanticTokensServerCapabilities, ServerCapabilities, ServerInfo,
+    TextDocumentContentChangeEvent, TextDocumentSyncCapability, TextDocumentSyncKind, Url,
+    WorkDoneProgressOptions, WorkspaceFileOperationsServerCapabilities,
+    WorkspaceServerCapabilities,
 };
 use ropey::Rope;
 
-use crate::{state::FileBuilder, GlobalState};
+use crate::{
+    state::{self, FileBuilder},
+    GlobalState,
+};
 
 mod compile;
 mod document_symbols;
 mod hover;
 mod semantic_tokens;
 
 pub fn initialize(
-    _state: &mut GlobalState<'_>,
-    _params: InitializeParams,
+    state: &mut GlobalState<'_>,
+    params: InitializeParams,
 ) -> Result<InitializeResult> {
+    log::trace!("{params:#?}");
+
+    if let Some(root) = params.root_uri {
+        for path in ignore::Walk::new(root.path()) {
+            let Ok(path) = path else { continue };
+
+            if path.file_type().is_some_and(|ty| ty.is_file())
+                && path.path().extension().is_some_and(|ext| ext == "stef")
+            {
+                let Ok(text) = std::fs::read_to_string(path.path()) else {
+                    error!(path = as_debug!(path.path()); "failed reading file content");
+                    continue;
+                };
+
+                let Ok(uri) = Url::from_file_path(path.path()) else {
+                    error!(path = as_debug!(path.path()); "failed parsing file path as URI");
+                    continue;
+                };
+
+                state.files.insert(uri.clone(), create_file(uri, text));
+            }
+        }
+    }
+
     Ok(InitializeResult {
         server_info: Some(ServerInfo {
             name: env!("CARGO_CRATE_NAME").to_owned(),
@@ -37,6 +68,22 @@ pub fn initialize(
             )),
             hover_provider: Some(HoverProviderCapability::Simple(true)),
             document_symbol_provider: Some(OneOf::Left(true)),
+            workspace: Some(WorkspaceServerCapabilities {
+                workspace_folders: None,
+                file_operations: Some(WorkspaceFileOperationsServerCapabilities {
+                    did_delete: Some(FileOperationRegistrationOptions {
+                        filters: vec![FileOperationFilter {
+                            scheme: Some("file".to_owned()),
+                            pattern: FileOperationPattern {
+                                glob: "**/*.stef".to_owned(),
+                                matches: Some(FileOperationPatternKind::File),
+                                options: None,
+                            },
+                        }],
+                    }),
+                    ..WorkspaceFileOperationsServerCapabilities::default()
+                }),
+            }),
             semantic_tokens_provider: Some(
                 SemanticTokensServerCapabilities::SemanticTokensOptions(SemanticTokensOptions {
                     work_done_progress_options: WorkDoneProgressOptions {
@@ -76,18 +123,23 @@ pub fn did_open(state: &mut GlobalState<'_>, params: DidOpenTextDocumentParams)
     debug!(uri = as_display!(params.text_document.uri); "schema opened");
 
     let text = params.text_document.text;
-    let file = FileBuilder {
-        rope: Rope::from_str(&text),
-        index: LineIndex::new(&text),
-        content: text,
-        schema_builder: |index, schema| {
-            compile::compile(params.text_document.uri.clone(), schema, index)
-        },
-    }
-    .build();
+    let file = if let Some(file) = state
+        .files
+        .get(&params.text_document.uri)
+        .filter(|file| file.borrow_content() == &text)
+    {
+        file
+    } else {
+        debug!("file missing from state");
+
+        let file = create_file(params.text_document.uri.clone(), text);
+
+        state.files.insert(params.text_document.uri.clone(), file);
+        &state.files[&params.text_document.uri]
+    };
 
     if let Err(e) = state.client.publish_diagnostics(
-        params.text_document.uri.clone(),
+        params.text_document.uri,
         file.borrow_schema()
             .as_ref()
             .err()
@@ -97,63 +149,44 @@ pub fn did_open(state: &mut GlobalState<'_>, params: DidOpenTextDocumentParams)
     ) {
         error!(error = as_debug!(e); "failed publishing diagnostics");
     }
-
-    state.files.insert(params.text_document.uri, file);
 }
 
 pub fn did_change(state: &mut GlobalState<'_>, mut params: DidChangeTextDocumentParams) {
-    debug!(uri = as_display!(params.text_document.uri); "schema changed");
+    fn is_full(changes: &[TextDocumentContentChangeEvent]) -> bool {
+        changes.len() == 1 && changes.first().is_some_and(|change| change.range.is_none())
+    }
 
-    let file = if params.content_changes.len() == 1
-        && params
-            .content_changes
-            .first()
-            .is_some_and(|change| change.range.is_none())
-    {
+    debug!(
+        uri = as_display!(params.text_document.uri),
+        full = as_display!(is_full(&params.content_changes));
+        "schema changed",
+    );
+
+    let file = if is_full(&params.content_changes) {
         let text = params.content_changes.remove(0).text;
-        FileBuilder {
-            rope: Rope::from_str(&text),
-            index: LineIndex::new(&text),
-            content: text,
-            schema_builder: |index, schema| {
-                compile::compile(params.text_document.uri.clone(), schema, index)
-            },
-        }
-        .build()
+        create_file(params.text_document.uri.clone(), text)
     } else {
         let Some(file) = state.files.remove(&params.text_document.uri) else {
             warn!("missing state for changed file");
             return;
         };
 
-        let mut heads = file.into_heads();
-
-        for change in params.content_changes {
-            let range = match convert_range(&heads.index, change.range) {
-                Ok(range) => range,
-                Err(e) => {
-                    error!(error = as_debug!(e); "invalid change");
-                    continue;
-                }
-            };
-
-            let start = heads.rope.byte_to_char(range.start().into());
-            let end = heads.rope.byte_to_char(range.end().into());
-            heads.rope.remove(start..end);
-            heads.rope.insert(start, &change.text);
-        }
-
-        let text = String::from(&heads.rope);
-
-        FileBuilder {
-            rope: heads.rope,
-            index: LineIndex::new(&text),
-            content: text,
-            schema_builder: |index, schema| {
-                compile::compile(params.text_document.uri.clone(), schema, index)
-            },
-        }
-        .build()
+        update_file(params.text_document.uri.clone(), file, |rope, index| {
+            for change in params.content_changes {
+                let range = match convert_range(index, change.range) {
+                    Ok(range) => range,
+                    Err(e) => {
+                        error!(error = as_debug!(e); "invalid change");
+                        continue;
+                    }
+                };
+
+                let start = rope.byte_to_char(range.start().into());
+                let end = rope.byte_to_char(range.end().into());
+                rope.remove(start..end);
+                rope.insert(start, &change.text);
+            }
+        })
     };
 
     if let Err(e) = state.client.publish_diagnostics(
@@ -171,9 +204,15 @@ pub fn did_change(state: &mut GlobalState<'_>, mut params: DidChangeTextDocument
     state.files.insert(params.text_document.uri, file);
 }
 
-pub fn did_close(state: &mut GlobalState<'_>, params: DidCloseTextDocumentParams) {
+pub fn did_close(_state: &mut GlobalState<'_>, params: DidCloseTextDocumentParams) {
     debug!(uri = as_display!(params.text_document.uri); "schema closed");
-    state.files.remove(&params.text_document.uri);
+}
+
+pub fn did_delete(state: &mut GlobalState<'_>, params: DeleteFilesParams) {
+    debug!(files = as_debug!(params.files); "files deleted");
+    state
+        .files
+        .retain(|uri, _| !params.files.iter().any(|file| file.uri == uri.as_str()));
 }
 
 pub fn hover(state: &mut GlobalState<'_>, params: HoverParams) -> Result<Option<Hover>> {
@@ -259,7 +298,7 @@ pub fn did_change_configuration(
     }
 }
 
-pub fn convert_range(index: &LineIndex, range: Option<lsp_types::Range>) -> Result<TextRange> {
+fn convert_range(index: &LineIndex, range: Option<lsp_types::Range>) -> Result<TextRange> {
     let range = range.context("incremental change misses range")?;
 
     let start = index
@@ -292,3 +331,32 @@ pub fn convert_range(index: &LineIndex, range: Option<lsp_types::Range>) -> Resu
 
     Ok(TextRange::new(start, end))
 }
+
+fn create_file(uri: Url, text: String) -> state::File {
+    FileBuilder {
+        rope: Rope::from_str(&text),
+        index: LineIndex::new(&text),
+        content: text,
+        schema_builder: |index, schema| compile::compile(uri, schema, index),
+    }
+    .build()
+}
+
+fn update_file(
+    uri: Url,
+    file: state::File,
+    update: impl FnOnce(&mut Rope, &LineIndex),
+) -> state::File {
+    let mut heads = file.into_heads();
+
+    update(&mut heads.rope, &heads.index);
+    let text = String::from(&heads.rope);
+
+    FileBuilder {
+        rope: heads.rope,
+        index: LineIndex::new(&text),
+        content: text,
+        schema_builder: |index, schema| compile::compile(uri, schema, index),
+    }
+    .build()
+}

crates/stef-lsp/src/main.rs

@@ -8,8 +8,8 @@ use log::{as_debug, debug, error, info, warn};
 use lsp_server::{Connection, ErrorCode, ExtractError, Notification, Request, RequestId, Response};
 use lsp_types::{
     notification::{
-        DidChangeConfiguration, DidChangeTextDocument, DidCloseTextDocument, DidOpenTextDocument,
-        Initialized, Notification as LspNotification,
+        DidChangeConfiguration, DidChangeTextDocument, DidCloseTextDocument, DidDeleteFiles,
+        DidOpenTextDocument, Initialized, Notification as LspNotification,
     },
     request::{
         DocumentSymbolRequest, HoverRequest, Request as LspRequest, SemanticTokensFullRequest,
@@ -138,6 +138,9 @@ fn main_loop(conn: &Connection, mut state: GlobalState<'_>) -> Result<()> {
                 DidCloseTextDocument::METHOD => {
                     handle_notify::<DidCloseTextDocument>(&mut state, notif, handlers::did_close);
                 }
+                DidDeleteFiles::METHOD => {
+                    handle_notify::<DidDeleteFiles>(&mut state, notif, handlers::did_delete);
+                }
                 DidChangeConfiguration::METHOD => {
                     handle_notify::<DidChangeConfiguration>(
                         &mut state,

crates/stef-lsp/src/state.rs

@@ -22,7 +22,7 @@ pub struct GlobalState<'a> {
 pub struct File {
     rope: Rope,
     pub index: LineIndex,
-    content: String,
+    pub content: String,
     #[borrows(index, content)]
     #[covariant]
     pub schema: Result<Schema<'this>, Diagnostic>,