fix: when formatting we need to use the in memory edited version

crates/lsp/src/error.rs

@@ -0,0 +1,29 @@
+use std::io;
+
+use crossbeam_channel::SendError;
+use lsp_server::{ExtractError, Message};
+
+#[derive(Debug, thiserror::Error, miette::Diagnostic)]
+pub enum Error {
+    #[error(transparent)]
+    #[diagnostic(code(aiken::lsp::server_capabilities))]
+    ServerCapabilities(#[from] serde_json::Error),
+    #[error(transparent)]
+    #[diagnostic(code(aiken::lsp::server_init))]
+    ServerInit(#[from] lsp_server::ProtocolError),
+    #[error(transparent)]
+    #[diagnostic(code(aiken::lsp::io))]
+    Io(#[from] io::Error),
+    #[error("Unsupported LSP request: {request}")]
+    #[diagnostic(code(aiken::lsp::unsupported_lsp_request))]
+    UnsupportedLspRequest { request: String },
+    #[error(transparent)]
+    #[diagnostic(code(aiken::lsp::cast_request))]
+    CastRequest(#[from] ExtractError<lsp_server::Request>),
+    #[error(transparent)]
+    #[diagnostic(code(aiken::lsp::cast_notification))]
+    CastNotification(#[from] ExtractError<lsp_server::Notification>),
+    #[error(transparent)]
+    #[diagnostic(code(aiken::lsp::send))]
+    Send(#[from] SendError<Message>),
+}

crates/lsp/src/lib.rs

@@ -1,34 +1,15 @@
-use std::{fs, io, path::Path};
+use lsp_server::Connection;
-
-use aiken_lang::{ast::ModuleKind, parser};
-use crossbeam_channel::SendError;
-use lsp_server::{Connection, ExtractError, Message};
 use lsp_types::{
-    request::{Formatting, Request},
+    OneOf, SaveOptions, ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind,
-    DocumentFormattingParams, InitializeParams, OneOf, ServerCapabilities, TextEdit,
+    TextDocumentSyncOptions, TextDocumentSyncSaveOptions,
 };
 
-#[derive(Debug, thiserror::Error, miette::Diagnostic)]
+pub mod error;
-pub enum Error {
+pub mod server;
-    #[error(transparent)]
+
-    #[diagnostic(code(aiken::lsp::server_capabilities))]
+use error::Error;
-    ServerCapabilities(#[from] serde_json::Error),
+
-    #[error(transparent)]
+use crate::server::Server;
-    #[diagnostic(code(aiken::lsp::server_init))]
-    ServerInit(#[from] lsp_server::ProtocolError),
-    #[error(transparent)]
-    #[diagnostic(code(aiken::lsp::io))]
-    Io(#[from] io::Error),
-    #[error("Unsupported LSP request: {request}")]
-    #[diagnostic(code(aiken::lsp::unsupported_lsp_request))]
-    UnsupportedLspRequest { request: String },
-    #[error(transparent)]
-    #[diagnostic(code(aiken::lsp::cast_request))]
-    CastRequest(#[from] ExtractError<lsp_server::Request>),
-    #[error(transparent)]
-    #[diagnostic(code(aiken::lsp::send))]
-    Send(#[from] SendError<Message>),
-}
 
 pub fn start() -> Result<(), Error> {
     tracing::info!("Aiken language server starting");
@@ -38,16 +19,14 @@ pub fn start() -> Result<(), Error> {
     let (connection, io_threads) = Connection::stdio();
 
     // Run the server and wait for the two threads to end (typically by trigger LSP Exit event).
-    let server_capabilities = serde_json::to_value(&ServerCapabilities {
+    let server_capabilities = serde_json::to_value(&capabilities())?;
-        definition_provider: Some(OneOf::Left(true)),
-        document_formatting_provider: Some(OneOf::Left(true)),
-        ..Default::default()
-    })?;
 
     let initialization_params = connection.initialize(server_capabilities)?;
-    let initialization_params = serde_json::from_value(initialization_params)?;
+    let initialize_params = serde_json::from_value(initialization_params)?;
 
-    main_loop(connection, initialization_params)?;
+    let mut server = Server::new(initialize_params, None);
+
+    server.listen(connection)?;
 
     io_threads.join()?;
 
@@ -56,97 +35,21 @@ pub fn start() -> Result<(), Error> {
     Ok(())
 }
 
-fn main_loop(connection: Connection, _params: InitializeParams) -> Result<(), Error> {
+fn capabilities() -> ServerCapabilities {
-    for msg in &connection.receiver {
+    ServerCapabilities {
-        tracing::debug!("Got message: {:#?}", msg);
+        text_document_sync: Some(TextDocumentSyncCapability::Options(
-
+            TextDocumentSyncOptions {
-        match msg {
+                open_close: None,
-            Message::Request(req) => {
+                change: Some(TextDocumentSyncKind::FULL),
-                if connection.handle_shutdown(&req)? {
+                will_save: None,
-                    return Ok(());
+                will_save_wait_until: None,
-                }
+                save: Some(TextDocumentSyncSaveOptions::SaveOptions(SaveOptions {
-
+                    include_text: Some(false),
-                tracing::debug!("Get request: {:#?}", req);
+                })),
-
-                let response = handle_request(req)?;
-
-                connection.sender.send(Message::Response(response))?;
-            }
-            Message::Response(_) => todo!(),
-            Message::Notification(notif) => tracing::debug!("Get notification: {:#?}", notif),
-        }
-    }
-
-    Ok(())
-}
-
-fn handle_request(request: lsp_server::Request) -> Result<lsp_server::Response, Error> {
-    let id = request.id.clone();
-
-    match request.method.as_str() {
-        Formatting::METHOD => {
-            let params = cast_request::<Formatting>(request)?;
-
-            let result = format(params);
-
-            match result {
-                Ok(text_edit) => {
-                    let result = serde_json::to_value(text_edit)?;
-
-                    Ok(lsp_server::Response {
-                        id,
-                        error: None,
-                        result: Some(result),
-                    })
-                }
-                Err(_) => {
-                    todo!("transform project errors in lsp diagnostic")
-                }
-            }
-        }
-        unsupported => Err(Error::UnsupportedLspRequest {
-            request: unsupported.to_string(),
-        }),
-    }
-}
-
-fn cast_request<R>(request: lsp_server::Request) -> Result<R::Params, Error>
-where
-    R: lsp_types::request::Request,
-    R::Params: serde::de::DeserializeOwned,
-{
-    let (_, params) = request.extract(R::METHOD)?;
-
-    Ok(params)
-}
-
-fn format(params: DocumentFormattingParams) -> Result<Vec<TextEdit>, aiken_project::error::Error> {
-    let path = params.text_document.uri.path();
-    let mut new_text = String::new();
-
-    let src = fs::read_to_string(path)?;
-
-    let (module, extra) = parser::module(&src, ModuleKind::Lib).map_err(|errs| {
-        aiken_project::error::Error::from_parse_errors(errs, Path::new(path), &src)
-    })?;
-
-    aiken_lang::format::pretty(&mut new_text, module, extra, &src);
-
-    Ok(vec![text_edit_replace(new_text)])
-}
-
-fn text_edit_replace(new_text: String) -> TextEdit {
-    TextEdit {
-        range: lsp_types::Range {
-            start: lsp_types::Position {
-                line: 0,
-                character: 0,
             },
-            end: lsp_types::Position {
+        )),
-                line: u32::MAX,
+        definition_provider: Some(OneOf::Left(true)),
-                character: 0,
+        document_formatting_provider: Some(OneOf::Left(true)),
-            },
+        ..Default::default()
-        },
-        new_text,
     }
 }

crates/lsp/src/server.rs

@@ -0,0 +1,178 @@
+use std::{collections::HashMap, fs, path::Path};
+
+use aiken_lang::{ast::ModuleKind, parser};
+use aiken_project::config;
+use lsp_server::{Connection, Message};
+use lsp_types::{
+    notification::{DidChangeTextDocument, Notification},
+    request::{Formatting, Request},
+    DocumentFormattingParams, InitializeParams, TextEdit,
+};
+
+use crate::error::Error;
+
+#[allow(dead_code)]
+pub struct Server {
+    config: Option<config::Config>,
+
+    /// Files that have been edited in memory
+    edited: HashMap<String, String>,
+
+    initialize_params: InitializeParams,
+}
+
+impl Server {
+    pub fn new(initialize_params: InitializeParams, config: Option<config::Config>) -> Self {
+        Self {
+            config,
+            edited: HashMap::new(),
+            initialize_params,
+        }
+    }
+
+    pub fn listen(&mut self, connection: Connection) -> Result<(), Error> {
+        for msg in &connection.receiver {
+            tracing::debug!("Got message: {:#?}", msg);
+
+            match msg {
+                Message::Request(req) => {
+                    if connection.handle_shutdown(&req)? {
+                        return Ok(());
+                    }
+
+                    tracing::debug!("Get request: {:#?}", req);
+
+                    let response = self.handle_request(req)?;
+
+                    connection.sender.send(Message::Response(response))?;
+                }
+                Message::Response(_) => todo!(),
+                Message::Notification(notification) => {
+                    self.handle_notification(&connection, notification)?
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    fn handle_request(
+        &mut self,
+        request: lsp_server::Request,
+    ) -> Result<lsp_server::Response, Error> {
+        let id = request.id.clone();
+
+        match request.method.as_str() {
+            Formatting::METHOD => {
+                let params = cast_request::<Formatting>(request)?;
+
+                let result = self.format(params);
+
+                match result {
+                    Ok(text_edit) => {
+                        let result = serde_json::to_value(text_edit)?;
+
+                        Ok(lsp_server::Response {
+                            id,
+                            error: None,
+                            result: Some(result),
+                        })
+                    }
+                    Err(_) => {
+                        todo!("transform project errors in lsp diagnostic")
+                    }
+                }
+            }
+            unsupported => Err(Error::UnsupportedLspRequest {
+                request: unsupported.to_string(),
+            }),
+        }
+    }
+
+    fn handle_notification(
+        &mut self,
+        _connection: &lsp_server::Connection,
+        notification: lsp_server::Notification,
+    ) -> Result<(), Error> {
+        match notification.method.as_str() {
+            DidChangeTextDocument::METHOD => {
+                let params = cast_notification::<DidChangeTextDocument>(notification)?;
+
+                // A file has changed in the editor so store a copy of the new content in memory
+                let path = params.text_document.uri.path().to_string();
+
+                if let Some(changes) = params.content_changes.into_iter().next() {
+                    let _ = self.edited.insert(path, changes.text);
+                }
+
+                Ok(())
+            }
+            _ => Ok(()),
+        }
+    }
+
+    fn format(
+        &mut self,
+        params: DocumentFormattingParams,
+    ) -> Result<Vec<TextEdit>, aiken_project::error::Error> {
+        let path = params.text_document.uri.path();
+        let mut new_text = String::new();
+
+        match self.edited.get(path) {
+            Some(src) => {
+                let (module, extra) = parser::module(src, ModuleKind::Lib).map_err(|errs| {
+                    aiken_project::error::Error::from_parse_errors(errs, Path::new(path), src)
+                })?;
+
+                aiken_lang::format::pretty(&mut new_text, module, extra, src);
+            }
+            None => {
+                let src = fs::read_to_string(path)?;
+
+                let (module, extra) = parser::module(&src, ModuleKind::Lib).map_err(|errs| {
+                    aiken_project::error::Error::from_parse_errors(errs, Path::new(path), &src)
+                })?;
+
+                aiken_lang::format::pretty(&mut new_text, module, extra, &src);
+            }
+        }
+
+        Ok(vec![text_edit_replace(new_text)])
+    }
+}
+
+fn cast_request<R>(request: lsp_server::Request) -> Result<R::Params, Error>
+where
+    R: lsp_types::request::Request,
+    R::Params: serde::de::DeserializeOwned,
+{
+    let (_, params) = request.extract(R::METHOD)?;
+
+    Ok(params)
+}
+
+fn cast_notification<N>(notification: lsp_server::Notification) -> Result<N::Params, Error>
+where
+    N: lsp_types::notification::Notification,
+    N::Params: serde::de::DeserializeOwned,
+{
+    let params = notification.extract::<N::Params>(N::METHOD)?;
+
+    Ok(params)
+}
+
+fn text_edit_replace(new_text: String) -> TextEdit {
+    TextEdit {
+        range: lsp_types::Range {
+            start: lsp_types::Position {
+                line: 0,
+                character: 0,
+            },
+            end: lsp_types::Position {
+                line: u32::MAX,
+                character: 0,
+            },
+        },
+        new_text,
+    }
+}