From c0513da032b2840db01b39732368094e20809dfe Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Sat, 21 Oct 2023 14:24:47 +0200
Subject: [PATCH] Add quickfix for unknown constructors.

---
 crates/aiken-lang/src/ast.rs     | 16 ++++++++++++++++
 crates/aiken-lsp/src/edits.rs    | 15 +++------------
 crates/aiken-lsp/src/quickfix.rs | 30 +++++++++++++++++++++++++++++-
 3 files changed, 48 insertions(+), 13 deletions(-)

diff --git a/crates/aiken-lang/src/ast.rs b/crates/aiken-lang/src/ast.rs
index b9447de8..e820de51 100644
--- a/crates/aiken-lang/src/ast.rs
+++ b/crates/aiken-lang/src/ast.rs
@@ -90,6 +90,22 @@ impl TypedModule {
         })
     }
 
+    pub fn has_constructor(&self, name: &str) -> bool {
+        self.definitions.iter().any(|def| match def {
+            Definition::DataType(t) if t.public && !t.opaque => t
+                .constructors
+                .iter()
+                .any(|constructor| constructor.name == name),
+            Definition::DataType(_) => false,
+            Definition::Fn(_) => false,
+            Definition::TypeAlias(_) => false,
+            Definition::ModuleConstant(_) => false,
+            Definition::Use(_) => false,
+            Definition::Test(_) => false,
+            Definition::Validator(_) => false,
+        })
+    }
+
     pub fn validate_module_name(&self) -> Result<(), Error> {
         if self.name == "aiken" || self.name == "aiken/builtin" {
             return Err(Error::ReservedModuleName {
diff --git a/crates/aiken-lsp/src/edits.rs b/crates/aiken-lsp/src/edits.rs
index 8aa9093b..1ccb4a0b 100644
--- a/crates/aiken-lsp/src/edits.rs
+++ b/crates/aiken-lsp/src/edits.rs
@@ -131,10 +131,7 @@ impl ParsedDocument {
         unqualified: &str,
         location: Span,
     ) -> AnnotatedEdit {
-        let title = format!(
-            "Insert new unqualified import '{}' to {}",
-            unqualified, import.name
-        );
+        let title = format!("Use '{}' from {}", unqualified, import.name);
         (
             title,
             insert_text(
@@ -151,10 +148,7 @@ impl ParsedDocument {
         unqualified: &str,
         location: Span,
     ) -> AnnotatedEdit {
-        let title = format!(
-            "Insert new unqualified import '{}' to {}",
-            unqualified, import.name
-        );
+        let title = format!("Use '{}' from {}", unqualified, import.name);
         (
             title,
             insert_text(
@@ -171,10 +165,7 @@ impl ParsedDocument {
         unqualified: &str,
         location: Span,
     ) -> AnnotatedEdit {
-        let title = format!(
-            "Add new unqualified import '{}' to {}",
-            unqualified, import.name
-        );
+        let title = format!("Use '{}' from {}", unqualified, import.name);
         (
             title,
             insert_text(
diff --git a/crates/aiken-lsp/src/quickfix.rs b/crates/aiken-lsp/src/quickfix.rs
index f21eb48b..a5c6184d 100644
--- a/crates/aiken-lsp/src/quickfix.rs
+++ b/crates/aiken-lsp/src/quickfix.rs
@@ -5,13 +5,16 @@ use crate::{
 use std::collections::HashMap;
 
 const UNKNOWN_VARIABLE: &str = "aiken::check::unknown::variable";
-const UNKNOWN_MODULE: &str = "aiken::check::unknown::module";
 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";
 
 /// Errors for which we can provide quickfixes
+#[allow(clippy::enum_variant_names)]
 pub enum Quickfix {
     UnknownIdentifier,
     UnknownModule,
+    UnknownConstructor,
 }
 
 fn match_code(diagnostic: &lsp_types::Diagnostic, expected: &str) -> bool {
@@ -35,6 +38,10 @@ pub fn assert(diagnostic: &lsp_types::Diagnostic) -> Option<Quickfix> {
         return Some(Quickfix::UnknownIdentifier);
     }
 
+    if match_code(diagnostic, UNKNOWN_CONSTRUCTOR) {
+        return Some(Quickfix::UnknownConstructor);
+    }
+
     if match_code(diagnostic, UNKNOWN_MODULE) {
         return Some(Quickfix::UnknownModule);
     }
@@ -55,6 +62,9 @@ pub fn quickfix(
             let edits = match quickfix {
                 Quickfix::UnknownIdentifier => unknown_identifier(compiler, parsed_document, data),
                 Quickfix::UnknownModule => unknown_module(compiler, parsed_document, data),
+                Quickfix::UnknownConstructor => {
+                    unknown_constructor(compiler, parsed_document, data)
+                }
             };
 
             for (title, edit) in edits.into_iter() {
@@ -99,6 +109,24 @@ fn unknown_identifier(
     edits
 }
 
+fn unknown_constructor(
+    compiler: &LspProject,
+    parsed_document: &ParsedDocument,
+    constructor_name: &str,
+) -> Vec<AnnotatedEdit> {
+    let mut edits = Vec::new();
+
+    for module in compiler.project.modules() {
+        if module.ast.has_constructor(constructor_name) {
+            if let Some(edit) = parsed_document.import(&module, Some(constructor_name)) {
+                edits.push(edit)
+            }
+        }
+    }
+
+    edits
+}
+
 fn unknown_module(
     compiler: &LspProject,
     parsed_document: &ParsedDocument,