Refactor and relocate document edits function for imports.

It's a bit 'off-topic' to keep these in aiken-lang as those functions are really just about lsp. Plus, it removes a bit some of the boilerplate and make the entire edition more readable and re-usable. Now we can tackle other similar errors with the same quickfix.

crates/aiken-lang/src/ast.rs

@@ -69,95 +69,6 @@ impl UntypedModule {
             })
             .collect()
     }
-
-    /// Find a suitable location (Span) in the import list. The boolean in the answer indicates
-    /// whether the import is a newline or not. It is set to 'false' when adding a qualified import
-    /// to an existing list.
-    pub fn edit_import(&self, module: &[&str], unqualified: Option<&str>) -> Option<EditImport> {
-        let mut last_import = None;
-
-        for def in self.definitions() {
-            match def {
-                Definition::Use(Use {
-                    location,
-                    module: existing_module,
-                    unqualified: unqualified_list,
-                    ..
-                }) => {
-                    last_import = Some(*location);
-
-                    if module != existing_module.as_slice() {
-                        continue;
-                    }
-
-                    match unqualified {
-                        // There's already a matching qualified import, so we have nothing to do.
-                        None => return None,
-                        Some(unqualified) => {
-                            // Insert lexicographically, assuming unqualified imports are already
-                            // ordered. If they are not, it doesn't really matter where we insert
-                            // anyway.
-                            for existing_unqualified in unqualified_list {
-                                let existing_name = existing_unqualified
-                                    .as_name
-                                    .as_ref()
-                                    .unwrap_or(&existing_unqualified.name);
-
-                                // The unqualified import already exist, nothing to do.
-                                if unqualified == existing_name {
-                                    return None;
-                                // Current import is lexicographically smaller, we can insert after.
-                                } else if existing_name.as_str() < unqualified {
-                                    return Some(EditImport {
-                                        location: Span {
-                                            start: existing_unqualified.location.end,
-                                            end: existing_unqualified.location.end,
-                                        },
-                                        is_new_line: false,
-                                        is_first_unqualified: false,
-                                    });
-                                } else {
-                                    continue;
-                                }
-                            }
-
-                            // Only happens if 'unqualified_list' is empty, in which case, we
-                            // simply create a new unqualified list of import.
-                            return Some(EditImport {
-                                location: Span {
-                                    start: location.end,
-                                    end: location.end,
-                                },
-                                is_new_line: false,
-                                is_first_unqualified: true,
-                            });
-                        }
-                    }
-                }
-                _ => continue,
-            }
-        }
-
-        // If the search above didn't lead to anything, we simply insert the import either:
-        //
-        // (a) After the last import statement if any;
-        // (b) As the first statement in the module.
-        Some(EditImport {
-            location: match last_import {
-                None => Span { start: 0, end: 0 },
-                Some(Span { end, .. }) => Span { start: end, end },
-            },
-            is_new_line: true,
-            is_first_unqualified: false,
-        })
-    }
-}
-
-#[derive(Debug, Clone, PartialEq)]
-pub struct EditImport {
-    pub location: Span,
-    pub is_new_line: bool,
-    pub is_first_unqualified: bool,
 }
 
 impl TypedModule {

crates/aiken-lsp/src/edits.rs

@@ -0,0 +1,180 @@
+use crate::{line_numbers::LineNumbers, utils::span_to_lsp_range};
+use aiken_lang::ast::{Definition, ModuleKind, Span, UntypedDefinition, Use};
+use aiken_project::module::CheckedModule;
+use itertools::Itertools;
+use std::fs;
+
+/// A freshly parsed module alongside its line numbers.
+pub struct ParsedDocument {
+    definitions: Vec<UntypedDefinition>,
+    line_numbers: LineNumbers,
+}
+
+/// Parse the target document as an 'UntypedModule' alongside its line numbers. This is useful in
+/// case we need to manipulate the AST for a quickfix.
+pub fn parse_document(document: &lsp_types::TextDocumentIdentifier) -> Option<ParsedDocument> {
+    let file_path = document
+        .uri
+        .to_file_path()
+        .expect("invalid text document uri?");
+
+    let source_code = fs::read_to_string(file_path).ok()?;
+
+    let line_numbers = LineNumbers::new(&source_code);
+
+    // NOTE: The 'ModuleKind' second argument doesn't matter. This is just added to the final
+    // object but has no influence on the parsing.
+    let (untyped_module, _) = aiken_lang::parser::module(&source_code, ModuleKind::Lib).ok()?;
+
+    Some(ParsedDocument {
+        definitions: untyped_module.definitions,
+        line_numbers,
+    })
+}
+
+/// Insert some text at the given location.
+fn insert_text(at: usize, line_numbers: &LineNumbers, new_text: String) -> lsp_types::TextEdit {
+    let range = span_to_lsp_range(Span { start: at, end: at }, line_numbers);
+    lsp_types::TextEdit { range, new_text }
+}
+
+/// Find a suitable location (Span) in the import list. The boolean in the answer indicates
+/// whether the import is a newline or not. It is set to 'false' when adding a qualified import
+/// to an existing list.
+impl ParsedDocument {
+    pub fn import(
+        &self,
+        import: &CheckedModule,
+        unqualified: Option<&str>,
+    ) -> Option<(String, lsp_types::TextEdit)> {
+        let import_path = import.name.split('/').collect_vec();
+
+        let mut last_import = None;
+
+        for def in self.definitions.iter() {
+            match def {
+                Definition::Use(Use {
+                    location,
+                    module: existing_module,
+                    unqualified: unqualified_list,
+                    ..
+                }) => {
+                    last_import = Some(*location);
+
+                    if import_path != existing_module.as_slice() {
+                        continue;
+                    }
+
+                    match unqualified {
+                        // There's already a matching qualified import, so we have nothing to do.
+                        None => return None,
+                        Some(unqualified) => {
+                            // Insert lexicographically, assuming unqualified imports are already
+                            // ordered. If they are not, it doesn't really matter where we insert
+                            // anyway.
+                            for existing_unqualified in unqualified_list {
+                                let existing_name = existing_unqualified
+                                    .as_name
+                                    .as_ref()
+                                    .unwrap_or(&existing_unqualified.name);
+
+                                // The unqualified import already exist, nothing to do.
+                                if unqualified == existing_name {
+                                    return None;
+                                // Current import is lexicographically smaller, we can insert after.
+                                } else if existing_name.as_str() < unqualified {
+                                    return Some(self.insert_into_qualified(
+                                        import,
+                                        unqualified,
+                                        existing_unqualified.location,
+                                    ));
+                                } else {
+                                    continue;
+                                }
+                            }
+
+                            // Only happens if 'unqualified_list' is empty, in which case, we
+                            // simply create a new unqualified list of import.
+                            return Some(self.add_new_qualified(import, unqualified, *location));
+                        }
+                    }
+                }
+                _ => continue,
+            }
+        }
+
+        // If the search above didn't lead to anything, we simply insert the import either:
+        //
+        // (a) After the last import statement if any;
+        // (b) As the first statement in the module.
+        Some(self.add_new_import_line(import, unqualified, last_import))
+    }
+
+    fn insert_into_qualified(
+        &self,
+        import: &CheckedModule,
+        unqualified: &str,
+        location: Span,
+    ) -> (String, lsp_types::TextEdit) {
+        let title = format!(
+            "Insert new unqualified import '{}' to {}",
+            unqualified, import.name
+        );
+        (
+            title,
+            insert_text(
+                location.end,
+                &self.line_numbers,
+                format!(", {}", unqualified),
+            ),
+        )
+    }
+
+    fn add_new_qualified(
+        &self,
+        import: &CheckedModule,
+        unqualified: &str,
+        location: Span,
+    ) -> (String, lsp_types::TextEdit) {
+        let title = format!(
+            "Add new unqualified import '{}' to {}",
+            unqualified, import.name
+        );
+        (
+            title,
+            insert_text(
+                location.end,
+                &self.line_numbers,
+                format!(".{{{}}}", unqualified),
+            ),
+        )
+    }
+
+    fn add_new_import_line(
+        &self,
+        import: &CheckedModule,
+        unqualified: Option<&str>,
+        location: Option<Span>,
+    ) -> (String, lsp_types::TextEdit) {
+        let import_line = format!(
+            "use {}{}",
+            import.name,
+            match unqualified {
+                Some(unqualified) => format!(".{{{}}}", unqualified),
+                None => String::new(),
+            }
+        );
+
+        let title = format!("Add new import line: {import_line}");
+
+        (
+            title,
+            match location {
+                None => insert_text(0, &self.line_numbers, format!("{import_line}\n")),
+                Some(Span { end, .. }) => {
+                    insert_text(end, &self.line_numbers, format!("\n{import_line}"))
+                }
+            },
+        )
+    }
+}

crates/aiken-lsp/src/lib.rs

@@ -5,6 +5,7 @@ use lsp_server::Connection;
 use std::env;
 
 mod cast;
+mod edits;
 pub mod error;
 mod line_numbers;
 pub mod server;

crates/aiken-lsp/src/server.rs

@@ -33,6 +33,7 @@ use miette::Diagnostic;
 
 use crate::{
     cast::{cast_notification, cast_request},
+    edits,
     error::Error as ServerError,
     line_numbers::LineNumbers,
     utils::{
@@ -409,63 +410,17 @@ impl Server {
     ) -> Option<Vec<lsp_types::CodeAction>> {
         let compiler = self.compiler.as_ref()?;
 
-        let file_path = text_document
+        let parsed_document = edits::parse_document(&text_document)?;
-            .uri
-            .to_file_path()
-            .expect("invalid text document uri?");
-
-        let source_code = fs::read_to_string(&file_path).ok()?;
-
-        let line_numbers = LineNumbers::new(&source_code);
-
-        // NOTE: The 'ModuleKind' second argument doesn't matter. This is just added to the final
-        // object but has no influence on the parsing.
-        let (untyped_module, _) = aiken_lang::parser::module(&source_code, ModuleKind::Lib).ok()?;
 
         let mut actions = Vec::new();
 
         if let Some(serde_json::Value::String(ref var_name)) = diagnostic.data {
             for module in compiler.project.modules() {
                 let mut changes = HashMap::new();
-                let module_path = module.name.split('/').collect_vec();
 
                 if module.ast.has_definition(var_name) {
-                    if let Some(edit) =
+                    if let Some((title, edit)) = parsed_document.import(&module, Some(var_name)) {
-                        untyped_module.edit_import(module_path.as_slice(), Some(var_name))
+                        changes.insert(text_document.uri.clone(), vec![edit]);
-                    {
-                        let (title, new_text) = if edit.is_new_line {
-                            (
-                                format!(
-                                    "Add new import line: use {}.{{{}}}",
-                                    module.name, var_name
-                                ),
-                                format!("\nuse {}.{{{}}}", module.name, var_name),
-                            )
-                        } else if edit.is_first_unqualified {
-                            (
-                                format!(
-                                    "Add new unqualified import '{}' to {}",
-                                    var_name, module.name
-                                ),
-                                format!(".{{{}}}", var_name),
-                            )
-                        } else {
-                            (
-                                format!(
-                                    "Add new unqualified import '{}' to {}",
-                                    var_name, module.name
-                                ),
-                                format!(", {}", var_name),
-                            )
-                        };
-
-                        let range = span_to_lsp_range(edit.location, &line_numbers);
-
-                        changes.insert(
-                            text_document.uri.clone(),
-                            vec![lsp_types::TextEdit { range, new_text }],
-                        );
-
                         actions.push(lsp_types::CodeAction {
                             title,
                             kind: Some(lsp_types::CodeActionKind::QUICKFIX),