Custom on_enter LSP method (FuelLabs#4876)

## Description

Closes FuelLabs#4728

Moves the on_enter capabilities to a custom LSP method that simply
returns the edits to the client, rather than directly applying the edits
as part of did_change.

With this approach, there is still some lag if you hit enter many times,
but it seems to be working better than before.

## Checklist

- [x] I have linked to any relevant issues.
- [ ] I have commented my code, particularly in hard-to-understand
areas.
- [ ] I have updated the documentation where relevant (API docs, the
reference, and the Sway book).
- [ ] I have added tests that prove my fix is effective or that my
feature works.
- [ ] I have added (or requested a maintainer to add) the necessary
`Breaking*` or `New Feature` labels where relevant.
- [ ] I have done my best to ensure that my PR adheres to [the Fuel Labs
Code Review
Standards](https://github.com/FuelLabs/rfcs/blob/master/text/code-standards/external-contributors.md).
- [ ] I have requested a review from the relevant team or maintainers.

sway-lsp/src/capabilities/on_enter.rs

@@ -1,15 +1,12 @@
 use crate::{
     config::OnEnterConfig,
     core::{document::TextDocument, session::Session},
+    lsp_ext::OnEnterParams,
 };
 use std::sync::Arc;
-use tower_lsp::{
-    lsp_types::{
-        DidChangeTextDocumentParams, DocumentChanges, OneOf,
-        OptionalVersionedTextDocumentIdentifier, Position, Range, TextDocumentEdit, TextEdit, Url,
-        WorkspaceEdit,
-    },
-    Client,
+use tower_lsp::lsp_types::{
+    DocumentChanges, OneOf, OptionalVersionedTextDocumentIdentifier, Position, Range,
+    TextDocumentEdit, TextEdit, Url, WorkspaceEdit,
 };
 
 const NEWLINE: &str = "\n";
@@ -18,15 +15,14 @@ const DOC_COMMENT_START: &str = "///";
 
 /// If the change was an enter keypress or pasting multiple lines in a comment, it prefixes the line(s)
 /// with the appropriate comment start pattern (// or ///).
-pub(crate) async fn on_enter(
+pub(crate) fn on_enter(
     config: &OnEnterConfig,
-    client: &Client,
     session: &Arc<Session>,
     temp_uri: &Url,
-    params: &DidChangeTextDocumentParams,
-) {
+    params: &OnEnterParams,
+) -> Option<WorkspaceEdit> {
     if !(params.content_changes[0].text.contains(NEWLINE)) {
-        return;
+        return None;
     }
 
     let mut workspace_edit = None;
@@ -42,66 +38,63 @@ pub(crate) async fn on_enter(
         workspace_edit = get_comment_workspace_edit(COMMENT_START, params, &text_document);
     }
 
-    // Apply any edits.
-    if let Some(edit) = workspace_edit {
-        if let Err(err) = client.apply_edit(edit).await {
-            tracing::error!("on_enter failed to apply edit: {}", err);
-        }
-    }
+    workspace_edit
 }
 
 fn get_comment_workspace_edit(
     start_pattern: &str,
-    change_params: &DidChangeTextDocumentParams,
+    change_params: &OnEnterParams,
     text_document: &TextDocument,
 ) -> Option<WorkspaceEdit> {
     let range = change_params.content_changes[0]
         .range
         .expect("change is missing range");
     let line = text_document.get_line(range.start.line as usize);
-    if line.trim().starts_with(start_pattern) {
-        let uri = change_params.text_document.uri.clone();
-        let text = change_params.content_changes[0].text.clone();
-
-        let indentation = &line[..line.find(start_pattern).unwrap_or(0)];
-        let mut edits = vec![];
-
-        // To support pasting multiple lines in a comment, we need to add the comment start pattern after each newline,
-        // except the last one.
-        let lines: Vec<_> = text.split(NEWLINE).collect();
-        lines.iter().enumerate().for_each(|(i, _)| {
-            if i < lines.len() - 1 {
-                let position =
-                    Position::new(range.start.line + (i as u32) + 1, indentation.len() as u32);
-                edits.push(OneOf::Left(TextEdit {
-                    new_text: format!("{start_pattern} "),
-                    range: Range::new(position, position),
-                }));
-            }
-        });
-        let edit = TextDocumentEdit {
-            text_document: OptionalVersionedTextDocumentIdentifier {
-                // Use the original uri to make updates, not the temporary one from the session.
-                uri,
-                version: None,
-            },
-            edits,
-        };
-        Some(WorkspaceEdit {
-            document_changes: Some(DocumentChanges::Edits(vec![edit])),
-            ..Default::default()
-        })
-    } else {
-        None
+
+    // If the previous line doesn't start with a comment, return early.
+    if !line.trim().starts_with(start_pattern) {
+        return None;
     }
+
+    let uri = change_params.text_document.uri.clone();
+    let text = change_params.content_changes[0].text.clone();
+
+    let indentation = &line[..line.find(start_pattern).unwrap_or(0)];
+    let mut edits = vec![];
+
+    // To support pasting multiple lines in a comment, we need to add the comment start pattern after each newline,
+    // except the last one.
+    let lines: Vec<_> = text.split(NEWLINE).collect();
+    lines.iter().enumerate().for_each(|(i, _)| {
+        if i < lines.len() - 1 {
+            let position =
+                Position::new(range.start.line + (i as u32) + 1, indentation.len() as u32);
+            edits.push(OneOf::Left(TextEdit {
+                new_text: format!("{start_pattern} "),
+                range: Range::new(position, position),
+            }));
+        }
+    });
+
+    let edit = TextDocumentEdit {
+        text_document: OptionalVersionedTextDocumentIdentifier {
+            // Use the original uri to make updates, not the temporary one from the session.
+            uri,
+            version: None,
+        },
+        edits,
+    };
+
+    Some(WorkspaceEdit {
+        document_changes: Some(DocumentChanges::Edits(vec![edit])),
+        ..Default::default()
+    })
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use lsp_types::{
-        AnnotatedTextEdit, TextDocumentContentChangeEvent, VersionedTextDocumentIdentifier,
-    };
+    use lsp_types::{AnnotatedTextEdit, TextDocumentContentChangeEvent, TextDocumentIdentifier};
     use sway_lsp_test_utils::get_absolute_path;
 
     fn assert_text_edit(
@@ -132,8 +125,8 @@ mod tests {
         let uri = Url::from_file_path(path.clone()).unwrap();
         let text_document =
             TextDocument::build_from_path(path.as_str()).expect("failed to build document");
-        let params = DidChangeTextDocumentParams {
-            text_document: VersionedTextDocumentIdentifier { uri, version: 1 },
+        let params = OnEnterParams {
+            text_document: TextDocumentIdentifier { uri },
             content_changes: vec![TextDocumentContentChangeEvent {
                 range: Some(Range {
                     start: Position {
@@ -169,8 +162,8 @@ mod tests {
         let uri = Url::from_file_path(path.clone()).unwrap();
         let text_document =
             TextDocument::build_from_path(path.as_str()).expect("failed to build document");
-        let params = DidChangeTextDocumentParams {
-            text_document: VersionedTextDocumentIdentifier { uri, version: 1 },
+        let params = OnEnterParams {
+            text_document: TextDocumentIdentifier { uri },
             content_changes: vec![TextDocumentContentChangeEvent {
                 range: Some(Range {
                     start: Position {

sway-lsp/src/handlers/notification.rs

@@ -1,7 +1,7 @@
 //! This module is responsible for implementing handlers for Language Server
 //! Protocol. This module specifically handles notification messages sent by the Client.
 
-use crate::{capabilities, core::sync, server_state::ServerState};
+use crate::{core::sync, server_state::ServerState};
 use forc_pkg::PackageManifestFile;
 use lsp_types::{
     DidChangeTextDocumentParams, DidChangeWatchedFilesParams, DidOpenTextDocumentParams,
@@ -30,15 +30,11 @@ pub(crate) async fn handle_did_change_text_document(
     state: &ServerState,
     params: DidChangeTextDocumentParams,
 ) {
-    let config = state.config.read().on_enter.clone();
     match state
         .sessions
         .uri_and_session_from_workspace(&params.text_document.uri)
     {
         Ok((uri, session)) => {
-            // handle on_enter capabilities if they are enabled
-            capabilities::on_enter(&config, &state.client, &session, &uri.clone(), &params).await;
-
             // update this file with the new changes and write to disk
             match session.write_changes_to_file(&uri, params.content_changes) {
                 Ok(_) => {

sway-lsp/src/handlers/request.rs

@@ -413,3 +413,29 @@ pub(crate) fn handle_show_ast(
         }
     }
 }
+
+/// This method is triggered when the use hits enter or pastes a newline in the editor.
+pub(crate) fn on_enter(
+    state: &ServerState,
+    params: lsp_ext::OnEnterParams,
+) -> Result<Option<WorkspaceEdit>> {
+    let config = state.config.read().on_enter.clone();
+    match state
+        .sessions
+        .uri_and_session_from_workspace(&params.text_document.uri)
+    {
+        Ok((uri, session)) => {
+            // handle on_enter capabilities if they are enabled
+            Ok(capabilities::on_enter(
+                &config,
+                &session,
+                &uri.clone(),
+                &params,
+            ))
+        }
+        Err(err) => {
+            tracing::error!("{}", err.to_string());
+            Ok(None)
+        }
+    }
+}

sway-lsp/src/lib.rs

@@ -26,6 +26,7 @@ use tower_lsp::{LspService, Server};
 pub async fn start() {
     let (service, socket) = LspService::build(ServerState::new)
         .custom_method("sway/show_ast", ServerState::show_ast)
+        .custom_method("sway/on_enter", ServerState::on_enter)
         .finish();
     Server::new(tokio::io::stdin(), tokio::io::stdout(), socket)
         .serve(service)

sway-lsp/src/lsp_ext.rs

@@ -1,6 +1,6 @@
 //! sway-lsp extensions to the LSP.
 
-use lsp_types::{TextDocumentIdentifier, Url};
+use lsp_types::{TextDocumentContentChangeEvent, TextDocumentIdentifier, Url};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -10,3 +10,11 @@ pub struct ShowAstParams {
     pub ast_kind: String,
     pub save_path: Url,
 }
+
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(rename_all = "camelCase")]
+pub struct OnEnterParams {
+    pub text_document: TextDocumentIdentifier,
+    /// The actual content changes, including the newline.
+    pub content_changes: Vec<TextDocumentContentChangeEvent>,
+}

sway-lsp/src/server.rs

@@ -3,7 +3,7 @@
 
 use crate::{
     handlers::{notification, request},
-    lsp_ext::ShowAstParams,
+    lsp_ext::{OnEnterParams, ShowAstParams},
     server_state::ServerState,
 };
 use lsp_types::{
@@ -117,4 +117,8 @@ impl ServerState {
     pub async fn show_ast(&self, params: ShowAstParams) -> Result<Option<TextDocumentIdentifier>> {
         request::handle_show_ast(self, params)
     }
+
+    pub async fn on_enter(&self, params: OnEnterParams) -> Result<Option<WorkspaceEdit>> {
+        request::on_enter(self, params)
+    }
 }