Implement quickfix for 'UnknownModule'.

crates/aiken-lang/src/tipo/error.rs

@@ -964,7 +964,6 @@ impl ExtraData for Error {
             | Error::UnexpectedLabeledArg { .. }
             | Error::UnexpectedLabeledArgInPattern { .. }
             | Error::UnknownLabels { .. }
-            | Error::UnknownModule { .. }
             | Error::UnknownModuleField { .. }
             | Error::UnknownModuleType { .. }
             | Error::UnknownModuleValue { .. }
@@ -975,7 +974,9 @@ impl ExtraData for Error {
             | Error::UpdateMultiConstructorType { .. }
             | Error::ValidatorImported { .. }
             | Error::ValidatorMustReturnBool { .. } => None,
-            Error::UnknownVariable { name, .. } => Some(name.clone()),
+            Error::UnknownVariable { name, .. } | Error::UnknownModule { name, .. } => {
+                Some(name.clone())
+            }
         }
     }
 }

crates/aiken-lsp/src/edits.rs

@@ -10,6 +10,8 @@ pub struct ParsedDocument {
     line_numbers: LineNumbers,
 }
 
+pub type AnnotatedEdit = (String, lsp_types::TextEdit);
+
 /// 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> {
@@ -46,7 +48,7 @@ impl ParsedDocument {
         &self,
         import: &CheckedModule,
         unqualified: Option<&str>,
-    ) -> Option<(String, lsp_types::TextEdit)> {
+    ) -> Option<AnnotatedEdit> {
         let import_path = import.name.split('/').collect_vec();
 
         let mut last_import = None;
@@ -115,7 +117,7 @@ impl ParsedDocument {
         import: &CheckedModule,
         unqualified: &str,
         location: Span,
-    ) -> (String, lsp_types::TextEdit) {
+    ) -> AnnotatedEdit {
         let title = format!(
             "Insert new unqualified import '{}' to {}",
             unqualified, import.name
@@ -135,7 +137,7 @@ impl ParsedDocument {
         import: &CheckedModule,
         unqualified: &str,
         location: Span,
-    ) -> (String, lsp_types::TextEdit) {
+    ) -> AnnotatedEdit {
         let title = format!(
             "Add new unqualified import '{}' to {}",
             unqualified, import.name
@@ -155,7 +157,7 @@ impl ParsedDocument {
         import: &CheckedModule,
         unqualified: Option<&str>,
         location: Option<Span>,
-    ) -> (String, lsp_types::TextEdit) {
+    ) -> AnnotatedEdit {
         let import_line = format!(
             "use {}{}",
             import.name,

crates/aiken-lsp/src/quickfix.rs

@@ -1,4 +1,7 @@
-use crate::{edits, server::lsp_project::LspProject};
+use crate::{
+    edits::{self, AnnotatedEdit, ParsedDocument},
+    server::lsp_project::LspProject,
+};
 use std::collections::HashMap;
 
 const UNKNOWN_VARIABLE: &str = "aiken::check::unknown::variable";
@@ -7,40 +10,47 @@ const UNKNOWN_MODULE: &str = "aiken::check::unknown::module";
 
 /// Errors for which we can provide quickfixes
 pub enum Quickfix {
-    UnknownVariable,
+    UnknownIdentifier,
+    UnknownModule,
 }
 
 fn match_code(diagnostic: &lsp_types::Diagnostic, expected: &str) -> bool {
     diagnostic.code == Some(lsp_types::NumberOrString::String(expected.to_string()))
 }
 
+/// Assert whether a diagnostic can be automatically fixed. Note that diagnostics often comes in
+/// two severities, an error and hint; so we must be careful only addressing errors.
 pub fn assert(diagnostic: &lsp_types::Diagnostic) -> Option<Quickfix> {
     let is_error = diagnostic.severity == Some(lsp_types::DiagnosticSeverity::ERROR);
 
     if is_error && match_code(diagnostic, UNKNOWN_VARIABLE) {
-        return Some(Quickfix::UnknownVariable);
+        return Some(Quickfix::UnknownIdentifier);
     }
 
     if is_error && match_code(diagnostic, UNKNOWN_MODULE) {
-        todo!()
+        return Some(Quickfix::UnknownModule);
     }
 
     None
 }
 
-pub fn unknown_variable(
+pub fn quickfix(
     compiler: &LspProject,
     text_document: &lsp_types::TextDocumentIdentifier,
     diagnostic: &lsp_types::Diagnostic,
+    quickfix: &Quickfix,
 ) -> Vec<lsp_types::CodeAction> {
     let mut actions = Vec::new();
 
-    if let Some(parsed_document) = edits::parse_document(text_document) {
+    if let Some(ref parsed_document) = edits::parse_document(text_document) {
-        if let Some(serde_json::Value::String(ref var_name)) = diagnostic.data {
+        if let Some(serde_json::Value::String(ref data)) = diagnostic.data {
-            for module in compiler.project.modules() {
+            let edits = match quickfix {
+                Quickfix::UnknownIdentifier => unknown_identifier(compiler, parsed_document, data),
+                Quickfix::UnknownModule => unknown_module(compiler, parsed_document, data),
+            };
+
+            for (title, edit) in edits.into_iter() {
                 let mut changes = HashMap::new();
-                if module.ast.has_definition(var_name) {
-                    if let Some((title, edit)) = parsed_document.import(&module, Some(var_name)) {
                 changes.insert(text_document.uri.clone(), vec![edit]);
                 actions.push(lsp_types::CodeAction {
                     title,
@@ -59,8 +69,42 @@ pub fn unknown_variable(
             }
         }
     }
-        }
-    }
 
     actions
 }
+
+fn unknown_identifier(
+    compiler: &LspProject,
+    parsed_document: &ParsedDocument,
+    var_name: &str,
+) -> Vec<AnnotatedEdit> {
+    let mut edits = Vec::new();
+
+    for module in compiler.project.modules() {
+        if module.ast.has_definition(var_name) {
+            if let Some(edit) = parsed_document.import(&module, Some(var_name)) {
+                edits.push(edit)
+            }
+        }
+    }
+
+    edits
+}
+
+fn unknown_module(
+    compiler: &LspProject,
+    parsed_document: &ParsedDocument,
+    module_name: &str,
+) -> Vec<AnnotatedEdit> {
+    let mut edits = Vec::new();
+
+    for module in compiler.project.modules() {
+        if module.name.ends_with(module_name) {
+            if let Some(edit) = parsed_document.import(&module, None) {
+                edits.push(edit);
+            }
+        }
+    }
+
+    edits
+}

crates/aiken-lsp/src/server.rs

@@ -35,7 +35,7 @@ use crate::{
     cast::{cast_notification, cast_request},
     error::Error as ServerError,
     line_numbers::LineNumbers,
-    quickfix::{self, Quickfix},
+    quickfix,
     utils::{
         path_to_uri, span_to_lsp_range, text_edit_replace, uri_to_module_name,
         COMPILING_PROGRESS_TOKEN, CREATE_COMPILING_PROGRESS_TOKEN,
@@ -336,17 +336,15 @@ impl Server {
                         .expect("cast code action request");
 
                     for diagnostic in params.context.diagnostics.iter() {
-                        match quickfix::assert(diagnostic) {
+                        if let Some(strategy) = quickfix::assert(diagnostic) {
-                            Some(Quickfix::UnknownVariable) => {
+                            let quickfixes = quickfix::quickfix(
-                                let quickfixes = quickfix::unknown_variable(
                                 compiler,
                                 &params.text_document,
                                 diagnostic,
+                                &strategy,
                             );
                             actions.extend(quickfixes);
                         }
-                            None => (),
-                        }
                     }
                 }