From ff5491caa068cfead52d559408aba116af7defcb Mon Sep 17 00:00:00 2001
From: rvcas <x@rvcas.dev>
Date: Thu, 29 Feb 2024 11:19:26 -0500
Subject: [PATCH] fix(check): only disallow ml_result in data

---
 crates/aiken-lang/src/builtins.rs    |  2 +-
 crates/aiken-lang/src/tests/check.rs | 27 +++++++++++++++++++++++++++
 crates/aiken-lang/src/tipo/infer.rs  |  2 +-
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/crates/aiken-lang/src/builtins.rs b/crates/aiken-lang/src/builtins.rs
index c9016593..c3b0d662 100644
--- a/crates/aiken-lang/src/builtins.rs
+++ b/crates/aiken-lang/src/builtins.rs
@@ -153,7 +153,7 @@ pub fn prelude(id_gen: &IdGenerator) -> TypeInfo {
         MILLER_LOOP_RESULT.to_string(),
         TypeConstructor {
             parameters: vec![],
-            tipo: int(),
+            tipo: miller_loop_result(),
             location: Span::empty(),
             module: "".to_string(),
             public: true,
diff --git a/crates/aiken-lang/src/tests/check.rs b/crates/aiken-lang/src/tests/check.rs
index d7bc8eb7..1c978c6d 100644
--- a/crates/aiken-lang/src/tests/check.rs
+++ b/crates/aiken-lang/src/tests/check.rs
@@ -50,6 +50,33 @@ fn check_validator(
     check_module(ast, ModuleKind::Validator)
 }
 
+#[test]
+fn bls12_381_elements_in_data_type() {
+    let source_code = r#"
+        type Datum {
+          D0(G1Element)
+          D1(G2Element)
+        }
+    "#;
+
+    assert!(check(parse(source_code)).is_ok())
+}
+
+#[test]
+fn bls12_381_ml_result_in_data_type() {
+    let source_code = r#"
+        type Datum {
+          thing: MillerLoopResult
+        }
+    "#;
+
+    let res = check(parse(source_code));
+
+    dbg!(&res);
+
+    assert!(matches!(res, Err((_, Error::IllegalTypeInData { .. }))))
+}
+
 #[test]
 fn validator_illegal_return_type() {
     let source_code = r#"
diff --git a/crates/aiken-lang/src/tipo/infer.rs b/crates/aiken-lang/src/tipo/infer.rs
index 6adf8da9..9b94181b 100644
--- a/crates/aiken-lang/src/tipo/infer.rs
+++ b/crates/aiken-lang/src/tipo/infer.rs
@@ -546,7 +546,7 @@ fn infer_definition(
                         });
                     }
 
-                    if tipo.is_bls381_12_g1() || tipo.is_bls381_12_g2() || tipo.is_ml_result() {
+                    if tipo.is_ml_result() {
                         return Err(Error::IllegalTypeInData {
                             location: *location,
                             tipo: tipo.clone(),