From dca633da48a2db466bf50bea7ec91f75cbf1b2e4 Mon Sep 17 00:00:00 2001
From: KtorZ <matthias.benkort@gmail.com>
Date: Fri, 23 Dec 2022 00:05:48 +0100
Subject: [PATCH] Refactor 'UnknownVariable' and 'UnknownTypeConstructor' as
 smart-constructor.

---
 crates/aiken-lang/src/tipo/environment.rs | 14 +---
 crates/aiken-lang/src/tipo/error.rs       | 92 ++++++++++++++++-------
 crates/aiken-lang/src/tipo/expr.rs        | 10 ++-
 crates/aiken-lang/src/tipo/pattern.rs     | 10 ++-
 4 files changed, 76 insertions(+), 50 deletions(-)

diff --git a/crates/aiken-lang/src/tipo/environment.rs b/crates/aiken-lang/src/tipo/environment.rs
index a6055a72..95766f8b 100644
--- a/crates/aiken-lang/src/tipo/environment.rs
+++ b/crates/aiken-lang/src/tipo/environment.rs
@@ -317,19 +317,7 @@ impl<'a> Environment<'a> {
     ) -> Result<&ValueConstructor, Error> {
         match module {
             None => self.scope.get(name).ok_or_else(|| {
-                if name.chars().into_iter().next().unwrap().is_uppercase() {
-                    Error::UnknownTypeConstructor {
-                        name: name.to_string(),
-                        variables: self.local_value_names(),
-                        location,
-                    }
-                } else {
-                    Error::UnknownVariable {
-                        name: name.to_string(),
-                        variables: self.local_value_names(),
-                        location,
-                    }
-                }
+                Error::unknown_variable_or_type(location, name, self.local_value_names())
             }),
 
             Some(m) => {
diff --git a/crates/aiken-lang/src/tipo/error.rs b/crates/aiken-lang/src/tipo/error.rs
index ae6f0fa3..2ed3e122 100644
--- a/crates/aiken-lang/src/tipo/error.rs
+++ b/crates/aiken-lang/src/tipo/error.rs
@@ -225,46 +225,24 @@ pub enum Error {
         types: Vec<String>,
     },
 
-    #[error("Unknown variable\n\n    {name}\n")]
+    #[error("Unknown variable '{name}'\n")]
+    #[diagnostic()]
     UnknownVariable {
         #[label]
         location: Span,
         name: String,
-        variables: Vec<String>,
+        #[help]
+        hint: String,
     },
 
     #[error("Unknown data-type constructor '{name}'\n")]
-    #[diagnostic(help(
-        r#"Did you forget to import it?
-
-Data-type constructors are not automatically imported, even if their type
-is imported. So, if a module `aiken/pet` defines the following type:
-
- ┍━ aiken/pet.ak ━━━━━━━━
- │ pub type Pet {{
- │   Cat
- │   Dog
- │ }}
-
-You must import its constructors explicitly to use them, or prefix them
-with the module's name.
-
- ┍━ foo.ak ━━━━━━━━
- │ use aiken/pet.{{Pet, Dog}}
- │
- │ fn foo(pet : Pet) {{
- │   when pet is {{
- │     pet.Cat -> // ...
- │     Dog -> // ...
- │   }}
- │ }}
-"#
-    ))]
+    #[diagnostic()]
     UnknownTypeConstructor {
         #[label]
         location: Span,
         name: String,
-        variables: Vec<String>,
+        #[help]
+        hint: String,
     },
 
     #[error("Unnecessary spread operator\n")]
@@ -455,6 +433,62 @@ impl Error {
             hint,
         }
     }
+
+    pub fn unknown_variable_or_type(location: Span, name: &str, variables: Vec<String>) -> Self {
+        let hint = variables
+            .iter()
+            .map(|s| (s, levenshtein::distance(name, s)))
+            .min_by(|(_, a), (_, b)| a.cmp(b))
+            .and_then(|(suggestion, distance)| {
+                if distance <= 3 {
+                    Some(format!("Did you mean '{suggestion}'?"))
+                } else {
+                    None
+                }
+            });
+
+        if name.chars().into_iter().next().unwrap().is_uppercase() {
+            let hint = hint.unwrap_or_else(|| {
+                r#"Did you forget to import it?
+
+Data-type constructors are not automatically imported, even if their type
+is imported. So, if a module `aiken/pet` defines the following type:
+
+ ┍━ aiken/pet.ak ━━━━━━━━
+ │ pub type Pet {{
+ │   Cat
+ │   Dog
+ │ }}
+
+You must import its constructors explicitly to use them, or prefix them
+with the module's name.
+
+ ┍━ foo.ak ━━━━━━━━
+ │ use aiken/pet.{{Pet, Dog}}
+ │
+ │ fn foo(pet : Pet) {{
+ │   when pet is {{
+ │     pet.Cat -> // ...
+ │     Dog -> // ...
+ │   }}
+ │ }}"#
+                    .to_string()
+            });
+
+            Self::UnknownTypeConstructor {
+                name: name.to_string(),
+                hint,
+                location,
+            }
+        } else {
+            let hint = hint.unwrap_or_else(|| "Did you forget to import it?".to_string());
+            Self::UnknownVariable {
+                name: name.to_string(),
+                location,
+                hint,
+            }
+        }
+    }
 }
 
 #[derive(Debug, PartialEq, Clone, thiserror::Error, Diagnostic)]
diff --git a/crates/aiken-lang/src/tipo/expr.rs b/crates/aiken-lang/src/tipo/expr.rs
index 02d99011..24b18eb8 100644
--- a/crates/aiken-lang/src/tipo/expr.rs
+++ b/crates/aiken-lang/src/tipo/expr.rs
@@ -1781,10 +1781,12 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                     self.environment
                         .get_variable(name)
                         .cloned()
-                        .ok_or_else(|| Error::UnknownVariable {
-                            location: *location,
-                            name: name.to_string(),
-                            variables: self.environment.local_value_names(),
+                        .ok_or_else(|| {
+                            Error::unknown_variable_or_type(
+                                *location,
+                                name,
+                                self.environment.local_value_names(),
+                            )
                         })?;
 
                 // Note whether we are using an ungeneralised function so that we can
diff --git a/crates/aiken-lang/src/tipo/pattern.rs b/crates/aiken-lang/src/tipo/pattern.rs
index a15c3064..f778be63 100644
--- a/crates/aiken-lang/src/tipo/pattern.rs
+++ b/crates/aiken-lang/src/tipo/pattern.rs
@@ -242,10 +242,12 @@ impl<'a, 'b> PatternTyper<'a, 'b> {
                     .environment
                     .get_variable(&name)
                     .cloned()
-                    .ok_or_else(|| Error::UnknownVariable {
-                    location,
-                    name: name.to_string(),
-                    variables: self.environment.local_value_names(),
+                    .ok_or_else(|| {
+                    Error::unknown_variable_or_type(
+                        location,
+                        &name,
+                        self.environment.local_value_names(),
+                    )
                 })?;
 
                 self.environment.increment_usage(&name);