Use (untyped) AST to find the right insert location for imports.

This removes the need to rely on the formatter to clear things up
  after insert a new import. While this is not so useful for imports, I
  wanted to experiment with the approach for future similar edits (for
  example, when suggesting an inline rewrite).

crates/aiken-lang/src/ast.rs

@@ -69,6 +69,95 @@ 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/server.rs

@@ -409,40 +409,78 @@ impl Server {
     ) -> Option<Vec<lsp_types::CodeAction>> {
         let compiler = self.compiler.as_ref()?;
 
+        let file_path = 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) {
-                    let import_line = format!("use {}.{{{}}}", module.name, var_name);
+                    if let Some(edit) =
+                        untyped_module.edit_import(module_path.as_slice(), Some(var_name))
+                    {
+                        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),
+                            )
+                        };
 
-                    changes.insert(
+                        let range = span_to_lsp_range(edit.location, &line_numbers);
-                        text_document.uri.clone(),
-                        vec![lsp_types::TextEdit {
-                            range: lsp_types::Range {
-                                start: lsp_types::Position::new(0, 0),
-                                end: lsp_types::Position::new(0, 0),
-                            },
-                            new_text: format!("{import_line}\n"),
-                        }],
-                    );
 
-                    actions.push(lsp_types::CodeAction {
+                        changes.insert(
-                        title: import_line,
+                            text_document.uri.clone(),
-                        kind: Some(lsp_types::CodeActionKind::QUICKFIX),
+                            vec![lsp_types::TextEdit { range, new_text }],
-                        diagnostics: Some(vec![diagnostic.clone()]),
+                        );
-                        is_preferred: Some(true),
+
-                        disabled: None,
+                        actions.push(lsp_types::CodeAction {
-                        data: None,
+                            title,
-                        command: None,
+                            kind: Some(lsp_types::CodeActionKind::QUICKFIX),
-                        edit: Some(lsp_types::WorkspaceEdit {
+                            diagnostics: Some(vec![diagnostic.clone()]),
-                            changes: Some(changes),
+                            is_preferred: Some(true),
-                            document_changes: None,
+                            disabled: None,
-                            change_annotations: None,
+                            data: None,
-                        }),
+                            command: None,
-                    });
+                            edit: Some(lsp_types::WorkspaceEdit {
+                                changes: Some(changes),
+                                document_changes: None,
+                                change_annotations: None,
+                            }),
+                        });
+                    }
                 }
             }
         }