Suggest wrapping value in Ok or Error if it makes sense

CHANGELOG.md

@@ -4,6 +4,43 @@
 
 ### Compiler
 
+- The compiler can now suggest to wrap a value in an `Ok` or `Error` if that can
+  solve a type mismatch error:
+
+  ```gleam
+  pub fn first(list: List(a)) -> Result(a, Nil) {
+    case number {
+      [] -> Error(Nil)
+      [first, ..rest] -> first
+    }
+  }
+  ```
+
+  Results in the following error:
+
+  ```txt
+  error: Type mismatch
+    ┌─ /src/one/two.gleam:5:5
+    │
+  5 │     [first, ..rest] -> first
+    │     ^^^^^^^^^^^^^^^^^^^^^^^^
+    │                        │
+    │                        Did you mean to wrap this in an `Ok`?
+
+  This case clause was found to return a different type than the previous
+  one, but all case clauses must return the same type.
+
+  Expected type:
+
+      Result(a, Nil)
+
+  Found type:
+
+      a
+  ```
+
+  ([Giacomo Cavalieri](https://github.com/giacomocavalieri))
+
 ### Build tool
 
 ### Language server

compiler-core/src/ast/typed.rs

@@ -336,8 +336,6 @@ impl TypedExpr {
         matches!(
             self,
             Self::Int{ value, .. } | Self::Float { value, .. } if NON_ZERO.get_or_init(||
-
-
                 Regex::new(r"[1-9]").expect("NON_ZERO regex")).is_match(value)
         )
     }

compiler-core/src/error.rs

@@ -1,6 +1,7 @@
 #![allow(clippy::unwrap_used, clippy::expect_used)]
 use crate::build::{Outcome, Runtime, Target};
 use crate::diagnostic::{Diagnostic, ExtraLabel, Label, Location};
+use crate::type_::collapse_links;
 use crate::type_::error::{
     MissingAnnotation, ModuleValueUsageContext, Named, UnknownField, UnknownTypeHint,
     UnsafeRecordUpdateReason,
@@ -21,6 +22,7 @@ use std::collections::HashSet;
 use std::fmt::{Debug, Display};
 use std::io::Write;
 use std::path::PathBuf;
+use std::sync::Arc;
 use termcolor::Buffer;
 use thiserror::Error;
 use vec1::Vec1;
@@ -1961,19 +1963,33 @@ But function expects:
                     text.push_str(&printer.print_type(expected));
                     text.push_str("\n\nFound type:\n\n    ");
                     text.push_str(&printer.print_type(given));
+
+                    let (main_message_location, main_message_text, extra_labels) = match situation {
+                        // When the mismatch error comes from a case clause we want to highlight the
+                        // entire branch (pattern included) when reporting the error; in addition,
+                        // if the error could be resolved just by wrapping the value in an `Ok`
+                        // or `Error` we want to add an additional label with this hint below the
+                        // offending value.
+                        Some(UnifyErrorSituation::CaseClauseMismatch{ clause_location }) => (clause_location, None, vec![]),
+                        // In all other cases we just highlight the offending expression, optionally
+                        // adding the wrapping hint if it makes sense.
+                        Some(_) | None =>
+                            (location, hint_wrap_value_in_result(expected, given), vec![])
+                    };
+
                     Diagnostic {
                         title: "Type mismatch".into(),
                         text,
                         hint: None,
                         level: Level::Error,
                         location: Some(Location {
                             label: Label {
-                                text: None,
-                                span: *location,
+                                text: main_message_text,
+                                span: *main_message_location,
                             },
                             path: path.clone(),
                             src: src.clone(),
-                            extra_labels: vec![],
+                            extra_labels,
                         }),
                     }
                 }
@@ -3948,6 +3964,19 @@ fn hint_alternative_operator(op: &BinOp, given: &Type) -> Option<String> {
     }
 }
 
+fn hint_wrap_value_in_result(expected: &Arc<Type>, given: &Arc<Type>) -> Option<String> {
+    let expected = collapse_links(expected.clone());
+    let (expected_ok_type, expected_error_type) = expected.result_types()?;
+
+    if given.same_as(expected_ok_type.as_ref()) {
+        Some("Did you mean to wrap this in an `Ok`?".into())
+    } else if given.same_as(expected_error_type.as_ref()) {
+        Some("Did you mean to wrap this in an `Error`?".into())
+    } else {
+        None
+    }
+}
+
 fn hint_numeric_message(alt: &str, type_: &str) -> String {
     format!("the {alt} operator can be used with {type_}s\n")
 }

compiler-core/src/type_.rs

@@ -138,6 +138,18 @@ impl Type {
         }
     }
 
+    pub fn result_types(&self) -> Option<(Arc<Type>, Arc<Type>)> {
+        match self {
+            Self::Named {
+                module, name, args, ..
+            } if "Result" == name && is_prelude_module(module) => {
+                Some((args.first().cloned()?, args.get(1).cloned()?))
+            }
+            Self::Var { type_ } => type_.borrow().result_types(),
+            Self::Named { .. } | Self::Tuple { .. } | Type::Fn { .. } => None,
+        }
+    }
+
     pub fn is_unbound(&self) -> bool {
         match self {
             Self::Var { type_ } => type_.borrow().is_unbound(),
@@ -433,6 +445,90 @@ impl Type {
             _ => None,
         }
     }
+
+    #[must_use]
+    /// Returns `true` is the two types are the same. This differs from the
+    /// standard `Eq` implementation as it also follows all links to check if
+    /// two types are really the same.
+    ///
+    pub fn same_as(&self, other: &Self) -> bool {
+        match (self, other) {
+            (Type::Named { .. }, Type::Fn { .. } | Type::Tuple { .. }) => false,
+            (one @ Type::Named { .. }, Type::Var { type_ }) => {
+                type_.as_ref().borrow().same_as_other_type(one)
+            }
+            (Type::Named { .. }, Type::Named { .. }) => self == other,
+
+            (Type::Fn { .. }, Type::Named { .. } | Type::Tuple { .. }) => false,
+            (one @ Type::Fn { .. }, Type::Var { type_ }) => {
+                type_.as_ref().borrow().same_as_other_type(one)
+            }
+            (
+                Type::Fn { args, retrn },
+                Type::Fn {
+                    args: other_args,
+                    retrn: other_retrn,
+                },
+            ) => {
+                args.len() == other_args.len()
+                    && args
+                        .iter()
+                        .zip(other_args)
+                        .all(|(one, other)| one.same_as(other))
+                    && retrn.same_as(other_retrn)
+            }
+
+            (Type::Var { type_ }, other) => type_.as_ref().borrow().same_as_other_type(other),
+
+            (Type::Tuple { .. }, Type::Fn { .. } | Type::Named { .. }) => false,
+            (one @ Type::Tuple { .. }, Type::Var { type_ }) => {
+                type_.as_ref().borrow().same_as_other_type(one)
+            }
+            (Type::Tuple { elems }, Type::Tuple { elems: other_elems }) => {
+                elems.len() == other_elems.len()
+                    && elems
+                        .iter()
+                        .zip(other_elems)
+                        .all(|(one, other)| one.same_as(other))
+            }
+        }
+    }
+}
+
+impl TypeVar {
+    #[must_use]
+    fn same_as_other_type(&self, other: &Type) -> bool {
+        match (self, other) {
+            (TypeVar::Unbound { .. }, _) => true,
+            (TypeVar::Link { type_ }, other) => type_.same_as(other),
+
+            (
+                TypeVar::Generic { .. },
+                Type::Named { .. } | Type::Fn { .. } | Type::Tuple { .. },
+            ) => false,
+
+            (one @ TypeVar::Generic { .. }, Type::Var { type_ }) => {
+                one.same_as(&type_.as_ref().borrow())
+            }
+        }
+    }
+
+    #[must_use]
+    fn same_as(&self, other: &Self) -> bool {
+        match (self, other) {
+            (TypeVar::Unbound { .. }, _) | (_, TypeVar::Unbound { .. }) => true,
+            (TypeVar::Link { type_ }, TypeVar::Link { type_: other_type }) => {
+                type_.same_as(other_type)
+            }
+            (TypeVar::Link { type_ }, other @ TypeVar::Generic { .. }) => {
+                other.same_as_other_type(type_)
+            }
+            (TypeVar::Generic { id }, TypeVar::Generic { id: other_id }) => id == other_id,
+            (one @ TypeVar::Generic { .. }, TypeVar::Link { type_ }) => {
+                one.same_as_other_type(type_)
+            }
+        }
+    }
 }
 
 pub fn collapse_links(t: Arc<Type>) -> Arc<Type> {
@@ -970,6 +1066,13 @@ impl TypeVar {
         }
     }
 
+    pub fn result_types(&self) -> Option<(Arc<Type>, Arc<Type>)> {
+        match self {
+            TypeVar::Link { type_ } => type_.result_types(),
+            TypeVar::Unbound { .. } | TypeVar::Generic { .. } => None,
+        }
+    }
+
     pub fn fn_types(&self) -> Option<(Vec<Arc<Type>>, Arc<Type>)> {
         match self {
             Self::Link { type_ } => type_.fn_types(),

compiler-core/src/type_/error.rs

@@ -1329,12 +1329,16 @@ fn flip_unify_error_test() {
         UnifyError::CouldNotUnify {
             expected: crate::type_::int(),
             given: crate::type_::float(),
-            situation: Some(UnifyErrorSituation::CaseClauseMismatch),
+            situation: Some(UnifyErrorSituation::CaseClauseMismatch {
+                clause_location: SrcSpan::default()
+            }),
         },
         flip_unify_error(UnifyError::CouldNotUnify {
             expected: crate::type_::float(),
             given: crate::type_::int(),
-            situation: Some(UnifyErrorSituation::CaseClauseMismatch),
+            situation: Some(UnifyErrorSituation::CaseClauseMismatch {
+                clause_location: SrcSpan::default()
+            }),
         })
     );
 }
@@ -1364,15 +1368,19 @@ fn unify_enclosed_type_test() {
         Err(UnifyError::CouldNotUnify {
             expected: crate::type_::int(),
             given: crate::type_::float(),
-            situation: Some(UnifyErrorSituation::CaseClauseMismatch)
+            situation: Some(UnifyErrorSituation::CaseClauseMismatch {
+                clause_location: SrcSpan::default()
+            })
         }),
         unify_enclosed_type(
             crate::type_::int(),
             crate::type_::float(),
             Err(UnifyError::CouldNotUnify {
                 expected: crate::type_::string(),
                 given: crate::type_::bits(),
-                situation: Some(UnifyErrorSituation::CaseClauseMismatch)
+                situation: Some(UnifyErrorSituation::CaseClauseMismatch {
+                    clause_location: SrcSpan::default()
+                })
             })
         )
     );
@@ -1437,7 +1445,9 @@ pub fn unify_wrong_returns(
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum UnifyErrorSituation {
     /// Clauses in a case expression were found to return different types.
-    CaseClauseMismatch,
+    CaseClauseMismatch {
+        clause_location: SrcSpan,
+    },
 
     /// A function was found to return a value that did not match its return
     /// annotation.
@@ -1480,7 +1490,7 @@ pub enum FunctionsMismatchReason {
 impl UnifyErrorSituation {
     pub fn description(&self) -> Option<&'static str> {
         match self {
-            Self::CaseClauseMismatch => Some(
+            Self::CaseClauseMismatch { clause_location: _ } => Some(
                 "This case clause was found to return a different type than the previous
 one, but all case clauses must return the same type.",
             ),
@@ -1545,8 +1555,8 @@ impl UnifyError {
         }
     }
 
-    pub fn case_clause_mismatch(self) -> Self {
-        self.with_unify_error_situation(UnifyErrorSituation::CaseClauseMismatch)
+    pub fn case_clause_mismatch(self, clause_location: SrcSpan) -> Self {
+        self.with_unify_error_situation(UnifyErrorSituation::CaseClauseMismatch { clause_location })
     }
 
     pub fn list_element_mismatch(self) -> Self {

compiler-core/src/type_/expression.rs

@@ -1469,9 +1469,10 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
             let typed_clause = self.infer_clause(clause, &typed_subjects);
             all_clauses_panic = all_clauses_panic && self.previous_panics;
 
-            if let Err(e) = unify(return_type.clone(), typed_clause.then.type_())
-                .map_err(|e| e.case_clause_mismatch().into_error(typed_clause.location()))
-            {
+            if let Err(e) = unify(return_type.clone(), typed_clause.then.type_()).map_err(|e| {
+                e.case_clause_mismatch(typed_clause.location)
+                    .into_error(typed_clause.then.type_defining_location())
+            }) {
                 self.problems.error(e);
             }
             typed_clauses.push(typed_clause);

compiler-core/src/type_/printer.rs

@@ -3,7 +3,7 @@ use ecow::EcoString;
 use im::HashMap;
 use std::{collections::HashSet, sync::Arc};
 
-use crate::type_::{collapse_links, Type, TypeVar};
+use crate::type_::{Type, TypeVar};
 
 /// This class keeps track of what names are used for modules in the current
 /// scope, so they can be printed in errors, etc.
@@ -126,9 +126,7 @@ fn compare_arguments(arguments: &[Arc<Type>], parameters: &[Arc<Type>]) -> bool
     arguments
         .iter()
         .zip(parameters)
-        .all(|(argument, parameter)| {
-            collapse_links(argument.clone()) == collapse_links(parameter.clone())
-        })
+        .all(|(argument, parameter)| argument.same_as(parameter))
 }
 
 impl Names {