Merge pull request #979 from aiken-lang/fix_casting_inference

Fix pattern inference and annotation requirements.

crates/aiken-lang/src/gen_uplc.rs

@@ -1404,7 +1404,10 @@ impl<'a> CodeGenerator<'a> {
                         ),
                     )
                 } else {
-                    assert!(data_type.constructors.len() == 1);
+                    assert!(
+                        data_type.constructors.len() == 1,
+                        "data_type={data_type:#?}"
+                    );
                     then
                 };
 

crates/aiken-lang/src/tests/check.rs

@@ -2024,18 +2024,6 @@ fn forbid_expect_into_opaque_type_from_data() {
     ))
 }
 
-#[test]
-fn allow_expect_into_type_from_data() {
-    let source_code = r#"
-        fn bar(n: Data) {
-          expect a: Option<Int> = n
-          a
-        }
-    "#;
-
-    assert!(check(parse(source_code)).is_ok())
-}
-
 #[test]
 fn forbid_partial_down_casting() {
     let source_code = r#"
@@ -2075,7 +2063,106 @@ fn forbid_partial_up_casting() {
 }
 
 #[test]
-fn allow_expect_into_type_from_data_2() {
+fn allow_expect_into_type_from_data() {
+    let source_code = r#"
+        fn bar(n: Data) {
+          expect a: Option<Int> = n
+          a
+        }
+    "#;
+
+    assert!(check(parse(source_code)).is_ok())
+}
+
+#[test]
+fn forbid_casting_into_type_from_data() {
+    let source_code = r#"
+        fn bar(n: Data) {
+          let a: Option<Int> = n
+          a
+        }
+    "#;
+
+    assert!(matches!(
+        check(parse(source_code)),
+        Err((_, Error::CouldNotUnify { .. }))
+    ))
+}
+
+#[test]
+fn forbid_casting_into_var_from_data_with_ann() {
+    let source_code = r#"
+        fn bar(n: Data) {
+          let a: Option<Int> = n
+          a
+        }
+    "#;
+
+    assert!(matches!(
+        check(parse(source_code)),
+        Err((_, Error::CouldNotUnify { .. }))
+    ))
+}
+
+#[test]
+fn allow_let_rebinding() {
+    let source_code = r#"
+        fn bar(n: Data) {
+          let a = n
+          a
+        }
+    "#;
+
+    assert!(check(parse(source_code)).is_ok())
+}
+
+#[test]
+fn expect_rebinding_requires_annotation() {
+    let source_code = r#"
+        fn bar(n: Data) -> Option<Int> {
+          expect a = n
+          a
+        }
+    "#;
+
+    assert!(matches!(
+        check(parse(source_code)),
+        Err((_, Error::CastDataNoAnn { .. }))
+    ))
+}
+
+#[test]
+fn forbid_casting_into_var_from_data_with_ann_indirect() {
+    let source_code = r#"
+        fn bar(n: Data) -> Option<Int> {
+          let a = n
+          a
+        }
+    "#;
+
+    assert!(matches!(
+        check(parse(source_code)),
+        Err((_, Error::CouldNotUnify { .. }))
+    ))
+}
+
+#[test]
+fn forbid_casting_into_pattern_from_data() {
+    let source_code = r#"
+        fn bar(n: Data) {
+          let Some(a) = n
+          a
+        }
+    "#;
+
+    assert!(matches!(
+        check(parse(source_code)),
+        Err((_, Error::CouldNotUnify { .. }))
+    ))
+}
+
+#[test]
+fn allow_expect_into_monomorphic_type_from_data_with_pattern() {
     let source_code = r#"
         fn bar(n: Data) {
           expect Some(a): Option<Int> = n
@@ -2086,6 +2173,52 @@ fn allow_expect_into_type_from_data_2() {
     assert!(check(parse(source_code)).is_ok())
 }
 
+#[test]
+fn forbid_expect_into_generic_type_from_data_with_pattern() {
+    let source_code = r#"
+        fn bar(n: Data) {
+          expect Some(a) = n
+          a
+        }
+    "#;
+
+    assert!(matches!(
+        check(parse(source_code)),
+        Err((_, Error::CastDataNoAnn { .. }))
+    ))
+}
+
+#[test]
+fn allow_generic_expect_without_typecast() {
+    let source_code = r#"
+        pub fn unwrap(opt: Option<a>) -> a {
+          expect Some(a) = opt
+          a
+        }
+    "#;
+
+    assert!(check(parse(source_code)).is_ok())
+}
+
+#[test]
+fn allow_expect_into_custom_type_from_data_no_annotation() {
+    let source_code = r#"
+        type OrderDatum {
+            requested_handle: ByteArray,
+            amount: Int,
+            other: Bool,
+        }
+
+        fn foo(datum: Data) {
+            expect OrderDatum { requested_handle, .. } = datum
+            requested_handle
+        }
+
+    "#;
+
+    assert!(check(parse(source_code)).is_ok())
+}
+
 #[test]
 fn forbid_expect_from_arbitrary_type() {
     let source_code = r#"
@@ -2110,7 +2243,7 @@ fn forbid_expect_from_arbitrary_type() {
 }
 
 #[test]
-fn forbid_expect_into_opaque_type_constructor_without_typecasting_in_module() {
+fn allow_expect_into_opaque_type_constructor_without_typecasting_in_module() {
     let source_code = r#"
         opaque type Thing {
           Foo(Int)

crates/aiken-lang/src/tipo.rs

@@ -406,6 +406,21 @@ impl Type {
         }
     }
 
+    ///  Check whether a given type is fully specialized and has only one possible
+    ///  form. Said differently, this recursively checks if the type still contains
+    ///  unbound or generic variables.
+    pub fn is_monomorphic(&self) -> bool {
+        match self {
+            Self::App { args, .. } => args.iter().all(|arg| arg.is_monomorphic()),
+            Self::Fn { args, ret, .. } => {
+                args.iter().all(|arg| arg.is_monomorphic()) && ret.is_monomorphic()
+            }
+            Self::Tuple { elems, .. } => elems.iter().all(|arg| arg.is_monomorphic()),
+            Self::Pair { fst, snd, .. } => [fst, snd].iter().all(|arg| arg.is_monomorphic()),
+            Self::Var { tipo, .. } => tipo.borrow().is_monomorphic(),
+        }
+    }
+
     pub fn is_generic(&self) -> bool {
         match self {
             Self::App { args, .. } => {
@@ -931,6 +946,16 @@ pub enum TypeVar {
 }
 
 impl TypeVar {
+    ///  Check whether a given type is fully specialized and has only one possible
+    ///  form. Said differently, this recursively checks if the type still contains
+    ///  unbound or generic variables.
+    pub fn is_monomorphic(&self) -> bool {
+        match self {
+            Self::Link { tipo } => tipo.is_monomorphic(),
+            Self::Unbound { .. } | Self::Generic { .. } => false,
+        }
+    }
+
     pub fn is_unbound(&self) -> bool {
         matches!(self, Self::Unbound { .. })
     }

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

@@ -1169,12 +1169,13 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
         location: Span,
     ) -> Result<TypedExpr, Error> {
         let typed_value = self.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 ann_typ = if let Some(ann) = annotation {
+        let pattern = if let Some(ann) = annotation {
             let ann_typ = self
                 .type_from_annotation(ann)
                 .and_then(|t| self.instantiate(t, &mut HashMap::new(), location))?;
@@ -1188,9 +1189,15 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
 
             value_typ = ann_typ.clone();
 
-            Some(ann_typ)
-        } else {
-            if value_is_data && !untyped_pattern.is_var() && !untyped_pattern.is_discard() {
+            // Ensure the pattern matches the type of the value
+            PatternTyper::new(self.environment, &self.hydrator).unify(
+                untyped_pattern.clone(),
+                value_typ.clone(),
+                Some(ann_typ),
+                kind.is_let(),
+            )
+        } else if value_is_data && !kind.is_let() {
+            let cast_data_no_ann = || {
                 let ann = Annotation::Constructor {
                     location: Span::empty(),
                     module: None,
@@ -1198,28 +1205,57 @@ impl<'a, 'b> ExprTyper<'a, 'b> {
                     arguments: vec![],
                 };
 
-                return Err(Error::CastDataNoAnn {
+                Err(Error::CastDataNoAnn {
                     location,
                     value: UntypedExpr::Assignment {
                         location,
-                        value: untyped_value.into(),
-                        patterns: AssignmentPattern::new(untyped_pattern, Some(ann), Span::empty())
+                        value: untyped_value.clone().into(),
+                        patterns: AssignmentPattern::new(
+                            untyped_pattern.clone(),
+                            Some(ann),
+                            Span::empty(),
+                        )
                         .into(),
                         kind,
                     },
-                });
-            }
-
-            None
+                })
             };
 
+            if !untyped_pattern.is_var() && !untyped_pattern.is_discard() {
+                let ann_typ = self.new_unbound_var();
+
+                match PatternTyper::new(self.environment, &self.hydrator).unify(
+                    untyped_pattern.clone(),
+                    ann_typ.clone(),
+                    None,
+                    false,
+                ) {
+                    Ok(pattern) if ann_typ.is_monomorphic() => {
+                        self.unify(
+                            ann_typ.clone(),
+                            value_typ.clone(),
+                            typed_value.type_defining_location(),
+                            true,
+                        )?;
+
+                        value_typ = ann_typ.clone();
+
+                        Ok(pattern)
+                    }
+                    Ok(..) | Err(..) => cast_data_no_ann(),
+                }
+            } else {
+                cast_data_no_ann()
+            }
+        } else {
             // Ensure the pattern matches the type of the value
-        let pattern = PatternTyper::new(self.environment, &self.hydrator).unify(
+            PatternTyper::new(self.environment, &self.hydrator).unify(
                 untyped_pattern.clone(),
                 value_typ.clone(),
-            ann_typ,
+                None,
                 kind.is_let(),
-        )?;
+            )
+        }?;
 
         // If `expect` is explicitly used, we still check exhaustiveness but instead of returning an
         // error we emit a warning which explains that using `expect` is unnecessary.

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

@@ -141,11 +141,11 @@ impl<'a, 'b> PatternTyper<'a, 'b> {
         pattern: UntypedPattern,
         tipo: Rc<Type>,
         ann_type: Option<Rc<Type>>,
-        is_let: bool,
+        warn_on_discard: bool,
     ) -> Result<TypedPattern, Error> {
         match pattern {
             Pattern::Discard { name, location } => {
-                if is_let {
+                if warn_on_discard {
                     // Register declaration for the unused variable detection
                     self.environment
                         .warnings