From 460da20e4a3ed83d12991de5a5a276f588c689b8 Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Thu, 9 Feb 2023 02:10:22 -0500
Subject: [PATCH] feat: warn when using expect in a completely safe way

---
 crates/aiken-lang/src/ast.rs        |  2 +-
 crates/aiken-lang/src/tipo/error.rs | 15 +++++++++
 crates/aiken-lang/src/tipo/expr.rs  | 48 +++++++++++++++++++++++------
 3 files changed, 55 insertions(+), 10 deletions(-)

diff --git a/crates/aiken-lang/src/ast.rs b/crates/aiken-lang/src/ast.rs
index 51ba9e2d..54c944e9 100644
--- a/crates/aiken-lang/src/ast.rs
+++ b/crates/aiken-lang/src/ast.rs
@@ -774,7 +774,7 @@ impl AssignmentKind {
         matches!(self, AssignmentKind::Let)
     }
 
-    pub fn is_assert(&self) -> bool {
+    pub fn is_expect(&self) -> bool {
         matches!(self, AssignmentKind::Expect)
     }
 }
diff --git a/crates/aiken-lang/src/tipo/error.rs b/crates/aiken-lang/src/tipo/error.rs
index 9be75612..5a24446b 100644
--- a/crates/aiken-lang/src/tipo/error.rs
+++ b/crates/aiken-lang/src/tipo/error.rs
@@ -1076,6 +1076,21 @@ pub enum Warning {
         sample: UntypedExpr,
     },
 
+    #[error("I found an {} trying to match a type with one constructor", "expect".purple())]
+    #[diagnostic(
+        code("single_constructor_expect"),
+        help("If your type has one constructor, unless you are casting {} {}, you can\nprefer using a {} binding like so...\n\n{}", "FROM".bold(), "Data".purple(), "let".purple(), format_suggestion(sample))
+    )]
+    SingleConstructorExpect {
+        #[label("use let")]
+        location: Span,
+        #[label("only one constructor")]
+        pattern_location: Span,
+        #[label("is not {}", "Data".purple())]
+        value_location: Span,
+        sample: UntypedExpr,
+    },
+
     #[error("I found a todo left in the code.\n")]
     #[diagnostic(help("You probably want to replace that one with real code... eventually."))]
     #[diagnostic(code("todo"))]
diff --git a/crates/aiken-lang/src/tipo/expr.rs b/crates/aiken-lang/src/tipo/expr.rs
index fd0915fb..919be51f 100644
--- a/crates/aiken-lang/src/tipo/expr.rs
+++ b/crates/aiken-lang/src/tipo/expr.rs
@@ -817,15 +817,18 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
 
     fn infer_assignment(
         &mut self,
-        pattern: UntypedPattern,
-        value: UntypedExpr,
+        untyped_pattern: UntypedPattern,
+        untyped_value: UntypedExpr,
         kind: AssignmentKind,
         annotation: &Option<Annotation>,
         location: Span,
     ) -> Result<TypedExpr, Error> {
-        let typed_value = self.in_new_scope(|value_typer| value_typer.infer(value.clone()))?;
+        let typed_value =
+            self.in_new_scope(|value_typer| value_typer.infer(untyped_value.clone()))?;
         let mut value_typ = typed_value.tipo();
 
+        let value_is_data = value_typ.is_data();
+
         // Check that any type annotation is accurate.
         let pattern = if let Some(ann) = annotation {
             let ann_typ = self
@@ -836,25 +839,25 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                 ann_typ.clone(),
                 value_typ.clone(),
                 typed_value.type_defining_location(),
-                (kind.is_let() && ann_typ.is_data()) || (kind.is_assert() && value_typ.is_data()),
+                (kind.is_let() && ann_typ.is_data()) || (kind.is_expect() && value_is_data),
             )?;
 
             value_typ = ann_typ.clone();
 
             // Ensure the pattern matches the type of the value
             PatternTyper::new(self.environment, &self.hydrator).unify(
-                pattern,
+                untyped_pattern.clone(),
                 value_typ.clone(),
                 Some(ann_typ),
             )?
         } else {
-            if value_typ.is_data() && !pattern.is_var() && !pattern.is_discard() {
+            if value_is_data && !untyped_pattern.is_var() && !untyped_pattern.is_discard() {
                 return Err(Error::CastDataNoAnn {
                     location,
                     value: UntypedExpr::Assignment {
                         location,
-                        value: value.into(),
-                        pattern,
+                        value: untyped_value.into(),
+                        pattern: untyped_pattern,
                         kind,
                         annotation: Some(Annotation::Constructor {
                             location: Span::empty(),
@@ -868,7 +871,7 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
 
             // Ensure the pattern matches the type of the value
             PatternTyper::new(self.environment, &self.hydrator).unify(
-                pattern,
+                untyped_pattern.clone(),
                 value_typ.clone(),
                 None,
             )?
@@ -888,6 +891,33 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                     unmatched,
                 });
             }
+        } else if !value_is_data
+            && self
+                .environment
+                .check_exhaustiveness(
+                    vec![pattern.clone()],
+                    collapse_links(value_typ.clone()),
+                    location,
+                )
+                .is_ok()
+        {
+            self.environment
+                .warnings
+                .push(Warning::SingleConstructorExpect {
+                    location: Span {
+                        start: location.start,
+                        end: location.start + 6,
+                    },
+                    pattern_location: dbg!(untyped_pattern.location()),
+                    value_location: dbg!(untyped_value.location()),
+                    sample: UntypedExpr::Assignment {
+                        location: Span::empty(),
+                        value: Box::new(untyped_value),
+                        pattern: untyped_pattern,
+                        kind: AssignmentKind::Let,
+                        annotation: None,
+                    },
+                })
         }
 
         Ok(TypedExpr::Assignment {