Implement quickfixes for redundant imports.

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

@@ -1591,14 +1591,18 @@ impl ExtraData for Warning {
             | Warning::Todo { .. }
             | Warning::UnexpectedTypeHole { .. }
             | Warning::UnusedConstructor { .. }
-            | Warning::UnusedImportedModule { .. }
-            | Warning::UnusedImportedValue { .. }
             | Warning::UnusedPrivateFunction { .. }
             | Warning::UnusedPrivateModuleConstant { .. }
             | Warning::UnusedType { .. }
             | Warning::UnusedVariable { .. }
             | Warning::Utf8ByteArrayIsValidHexString { .. }
             | Warning::ValidatorInLibraryModule { .. } => None,
+            Warning::UnusedImportedModule { location, .. } => {
+                Some(format!("{},{}", false, location.start))
+            }
+            Warning::UnusedImportedValueOrType { location, .. } => {
+                Some(format!("{},{}", true, location.start))
+            }
         }
     }
 }

crates/aiken-lsp/src/edits.rs

@@ -8,6 +8,7 @@ use std::fs;
 pub struct ParsedDocument {
     definitions: Vec<UntypedDefinition>,
     line_numbers: LineNumbers,
+    source_code: String,
 }
 
 pub type AnnotatedEdit = (String, lsp_types::TextEdit);
@@ -31,6 +32,7 @@ pub fn parse_document(document: &lsp_types::TextDocumentIdentifier) -> Option<Pa
     Some(ParsedDocument {
         definitions: untyped_module.definitions,
         line_numbers,
+        source_code,
     })
 }
 
@@ -125,6 +127,48 @@ impl ParsedDocument {
         Some(self.add_new_import_line(import, unqualified, last_import))
     }
 
+    pub fn remove_import(&self, start: usize, is_qualified: bool) -> AnnotatedEdit {
+        let offset = if is_qualified {
+            let import_len = self
+                .source_code
+                .chars()
+                .skip(start)
+                .take_while(|c| c != &',' && c != &'}')
+                .count();
+
+            let has_trailing_comma = self
+                .source_code
+                .chars()
+                .skip(start + import_len)
+                .collect::<String>()
+                .starts_with(',');
+
+            import_len + if has_trailing_comma { 1 } else { 0 }
+        } else {
+            1 + self
+                .source_code
+                .chars()
+                .skip(start)
+                .take_while(|c| c != &'\n')
+                .count()
+        };
+
+        let range = span_to_lsp_range(
+            Span {
+                start,
+                end: start + offset,
+            },
+            &self.line_numbers,
+        );
+
+        let new_text = String::new();
+
+        (
+            "Remove redundant import".to_string(),
+            lsp_types::TextEdit { range, new_text },
+        )
+    }
+
     fn insert_qualified_before(
         &self,
         import: &CheckedModule,

crates/aiken-lsp/src/quickfix.rs

@@ -2,48 +2,56 @@ use crate::{
     edits::{self, AnnotatedEdit, ParsedDocument},
     server::lsp_project::LspProject,
 };
-use std::collections::HashMap;
+use std::{collections::HashMap, str::FromStr};
 
 const UNKNOWN_VARIABLE: &str = "aiken::check::unknown::variable";
 const UNKNOWN_TYPE: &str = "aiken::check::unknown::type";
 const UNKNOWN_CONSTRUCTOR: &str = "aiken::check::unknown::type_constructor";
 const UNKNOWN_MODULE: &str = "aiken::check::unknown::module";
+const UNUSED_IMPORT_VALUE: &str = "aiken::check::unused:import::value";
+const UNUSED_IMPORT_MODULE: &str = "aiken::check::unused::import::module";
 
 /// Errors for which we can provide quickfixes
 #[allow(clippy::enum_variant_names)]
 pub enum Quickfix {
-    UnknownIdentifier,
+    UnknownIdentifier(lsp_types::Diagnostic),
-    UnknownModule,
+    UnknownModule(lsp_types::Diagnostic),
-    UnknownConstructor,
+    UnknownConstructor(lsp_types::Diagnostic),
+    UnusedImports(Vec<lsp_types::Diagnostic>),
 }
 
-fn match_code(diagnostic: &lsp_types::Diagnostic, expected: &str) -> bool {
+fn match_code(
+    diagnostic: &lsp_types::Diagnostic,
+    severity: lsp_types::DiagnosticSeverity,
+    expected: &str,
+) -> bool {
     diagnostic.code == Some(lsp_types::NumberOrString::String(expected.to_string()))
+        && diagnostic.severity == Some(severity)
 }
 
 /// 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> {
+pub fn assert(diagnostic: lsp_types::Diagnostic) -> Option<Quickfix> {
-    let is_error = diagnostic.severity == Some(lsp_types::DiagnosticSeverity::ERROR);
+    use lsp_types::DiagnosticSeverity as Severity;
 
-    if !is_error {
+    if match_code(&diagnostic, Severity::ERROR, UNKNOWN_VARIABLE)
-        return None;
+        || match_code(&diagnostic, Severity::ERROR, UNKNOWN_TYPE)
+    {
+        return Some(Quickfix::UnknownIdentifier(diagnostic));
     }
 
-    if match_code(diagnostic, UNKNOWN_VARIABLE) {
+    if match_code(&diagnostic, Severity::ERROR, UNKNOWN_CONSTRUCTOR) {
-        return Some(Quickfix::UnknownIdentifier);
+        return Some(Quickfix::UnknownConstructor(diagnostic));
     }
 
-    if match_code(diagnostic, UNKNOWN_TYPE) {
+    if match_code(&diagnostic, Severity::ERROR, UNKNOWN_MODULE) {
-        return Some(Quickfix::UnknownIdentifier);
+        return Some(Quickfix::UnknownModule(diagnostic));
     }
 
-    if match_code(diagnostic, UNKNOWN_CONSTRUCTOR) {
+    if match_code(&diagnostic, Severity::WARNING, UNUSED_IMPORT_VALUE)
-        return Some(Quickfix::UnknownConstructor);
+        || match_code(&diagnostic, Severity::WARNING, UNUSED_IMPORT_MODULE)
-    }
+    {
-
+        return Some(Quickfix::UnusedImports(vec![diagnostic]));
-    if match_code(diagnostic, UNKNOWN_MODULE) {
-        return Some(Quickfix::UnknownModule);
     }
 
     None
@@ -52,24 +60,62 @@ pub fn assert(diagnostic: &lsp_types::Diagnostic) -> Option<Quickfix> {
 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(ref parsed_document) = edits::parse_document(text_document) {
-        if let Some(serde_json::Value::String(ref data)) = diagnostic.data {
+        match quickfix {
-            let edits = match quickfix {
+            Quickfix::UnknownIdentifier(diagnostic) => {
-                Quickfix::UnknownIdentifier => unknown_identifier(compiler, parsed_document, data),
+                each_as_distinct_action(
-                Quickfix::UnknownModule => unknown_module(compiler, parsed_document, data),
+                    &mut actions,
-                Quickfix::UnknownConstructor => {
+                    text_document,
-                    unknown_constructor(compiler, parsed_document, data)
+                    diagnostic,
+                    unknown_identifier(compiler, parsed_document, diagnostic.data.as_ref()),
+                );
             }
+            Quickfix::UnknownModule(diagnostic) => each_as_distinct_action(
+                &mut actions,
+                text_document,
+                diagnostic,
+                unknown_module(compiler, parsed_document, diagnostic.data.as_ref()),
+            ),
+            Quickfix::UnknownConstructor(diagnostic) => each_as_distinct_action(
+                &mut actions,
+                text_document,
+                diagnostic,
+                unknown_constructor(compiler, parsed_document, diagnostic.data.as_ref()),
+            ),
+            Quickfix::UnusedImports(diagnostics) => as_single_action(
+                &mut actions,
+                text_document,
+                diagnostics.to_owned(),
+                "Remove redundant imports",
+                unused_imports(
+                    parsed_document,
+                    diagnostics
+                        .iter()
+                        .map(|diagnostic| diagnostic.data.as_ref())
+                        .collect(),
+                ),
+            ),
         };
+    }
 
+    actions
+}
+
+fn each_as_distinct_action(
+    actions: &mut Vec<lsp_types::CodeAction>,
+    text_document: &lsp_types::TextDocumentIdentifier,
+    diagnostic: &lsp_types::Diagnostic,
+    edits: Vec<AnnotatedEdit>,
+) {
     for (title, edit) in edits.into_iter() {
         let mut changes = HashMap::new();
+
         changes.insert(text_document.uri.clone(), vec![edit]);
+
         actions.push(lsp_types::CodeAction {
             title,
             kind: Some(lsp_types::CodeActionKind::QUICKFIX),
@@ -85,19 +131,46 @@ pub fn quickfix(
             }),
         });
     }
-        }
+}
-    }
 
-    actions
+fn as_single_action(
+    actions: &mut Vec<lsp_types::CodeAction>,
+    text_document: &lsp_types::TextDocumentIdentifier,
+    diagnostics: Vec<lsp_types::Diagnostic>,
+    title: &str,
+    edits: Vec<AnnotatedEdit>,
+) {
+    let mut changes = HashMap::new();
+
+    changes.insert(
+        text_document.uri.clone(),
+        edits.into_iter().map(|(_, b)| b).collect(),
+    );
+
+    actions.push(lsp_types::CodeAction {
+        title: title.to_string(),
+        kind: Some(lsp_types::CodeActionKind::QUICKFIX),
+        diagnostics: Some(diagnostics),
+        is_preferred: Some(true),
+        disabled: None,
+        data: None,
+        command: None,
+        edit: Some(lsp_types::WorkspaceEdit {
+            changes: Some(changes),
+            document_changes: None,
+            change_annotations: None,
+        }),
+    });
 }
 
 fn unknown_identifier(
     compiler: &LspProject,
     parsed_document: &ParsedDocument,
-    var_name: &str,
+    data: Option<&serde_json::Value>,
 ) -> Vec<AnnotatedEdit> {
     let mut edits = Vec::new();
 
+    if let Some(serde_json::Value::String(ref var_name)) = data {
         for module in compiler.project.modules() {
             if module.ast.has_definition(var_name) {
                 if let Some(edit) = parsed_document.import(&module, Some(var_name)) {
@@ -105,6 +178,7 @@ fn unknown_identifier(
                 }
             }
         }
+    }
 
     edits
 }
@@ -112,10 +186,11 @@ fn unknown_identifier(
 fn unknown_constructor(
     compiler: &LspProject,
     parsed_document: &ParsedDocument,
-    constructor_name: &str,
+    data: Option<&serde_json::Value>,
 ) -> Vec<AnnotatedEdit> {
     let mut edits = Vec::new();
 
+    if let Some(serde_json::Value::String(ref constructor_name)) = data {
         for module in compiler.project.modules() {
             if module.ast.has_constructor(constructor_name) {
                 if let Some(edit) = parsed_document.import(&module, Some(constructor_name)) {
@@ -123,6 +198,7 @@ fn unknown_constructor(
                 }
             }
         }
+    }
 
     edits
 }
@@ -130,10 +206,11 @@ fn unknown_constructor(
 fn unknown_module(
     compiler: &LspProject,
     parsed_document: &ParsedDocument,
-    module_name: &str,
+    data: Option<&serde_json::Value>,
 ) -> Vec<AnnotatedEdit> {
     let mut edits = Vec::new();
 
+    if let Some(serde_json::Value::String(ref module_name)) = data {
         for module in compiler.project.modules() {
             if module.name.ends_with(module_name) {
                 if let Some(edit) = parsed_document.import(&module, None) {
@@ -141,6 +218,37 @@ fn unknown_module(
                 }
             }
         }
+    }
+
+    edits
+}
+
+fn unused_imports(
+    parsed_document: &ParsedDocument,
+    datas: Vec<Option<&serde_json::Value>>,
+) -> Vec<AnnotatedEdit> {
+    let mut edits = Vec::new();
+
+    for data in datas.iter().rev().flatten() {
+        if let serde_json::Value::String(ref args) = data {
+            let args = args.split(',').collect::<Vec<&str>>();
+            match args.as_slice() {
+                &[is_qualified, start] => {
+                    let start = start
+                        .parse::<usize>()
+                        .expect("malformed unused_imports argument: not a usize");
+
+                    let is_qualified = FromStr::from_str(is_qualified)
+                        .expect("malformed unused_imports argument: not a bool");
+
+                    edits.push(parsed_document.remove_import(start, is_qualified));
+                }
+                _ => {
+                    panic!("malformed unused_imports arguments: not a 2-tuple");
+                }
+            }
+        }
+    }
 
     edits
 }

crates/aiken-lsp/src/server.rs

@@ -1,3 +1,4 @@
+use crate::quickfix::Quickfix;
 use std::{
     collections::{HashMap, HashSet},
     fs,
@@ -335,18 +336,36 @@ impl Server {
                     let params = cast_request::<CodeActionRequest>(request)
                         .expect("cast code action request");
 
-                    for diagnostic in params.context.diagnostics.iter() {
+                    let mut unused_imports = Vec::new();
-                        if let Some(strategy) = quickfix::assert(diagnostic) {
+
+                    for diagnostic in params.context.diagnostics.into_iter() {
+                        match quickfix::assert(diagnostic) {
+                            None => (),
+                            Some(Quickfix::UnusedImports(diagnostics)) => {
+                                unused_imports.extend(diagnostics);
+                            }
+                            Some(strategy) => {
+                                let quickfixes =
+                                    quickfix::quickfix(compiler, &params.text_document, &strategy);
+                                actions.extend(quickfixes);
+                            }
+                        }
+                    }
+
+                    // NOTE: We handle unused imports separately because we want them to be done in
+                    // a single action. Otherwise they might mess up with spans and alignments as
+                    // they remove content from the document.
+                    // Plus, it's also just much more convenient that way instead of removing each
+                    // import one-by-one.
+                    if !unused_imports.is_empty() {
                         let quickfixes = quickfix::quickfix(
                             compiler,
                             &params.text_document,
-                                diagnostic,
+                            &Quickfix::UnusedImports(unused_imports),
-                                &strategy,
                         );
                         actions.extend(quickfixes);
                     }
                 }
-                }
 
                 Ok(lsp_server::Response {
                     id,