Merge pull request #421 from aiken-lang/monomorph-panic-fix

Functions with only a generic return weren't being properly monomorph…

crates/aiken-lang/src/builder.rs

@@ -1200,7 +1200,7 @@ pub fn find_and_replace_generics(tipo: &mut Arc<Type>, mono_types: &IndexMap<u64
     }
 }
 
-pub fn get_generics_and_type(tipo: &Type, param: &Type) -> Vec<(u64, Arc<Type>)> {
+pub fn get_generic_id_and_type(tipo: &Type, param: &Type) -> Vec<(u64, Arc<Type>)> {
     let mut generics_ids = vec![];
 
     if let Some(id) = tipo.get_generic() {
@@ -1213,7 +1213,7 @@ pub fn get_generics_and_type(tipo: &Type, param: &Type) -> Vec<(u64, Arc<Type>)>
         .iter()
         .zip(param.get_inner_types().iter())
     {
-        generics_ids.append(&mut get_generics_and_type(tipo, param_type));
+        generics_ids.append(&mut get_generic_id_and_type(tipo, param_type));
     }
     generics_ids
 }
@@ -1990,7 +1990,7 @@ pub fn replace_opaque_type(t: &mut Arc<Type>, data_types: IndexMap<DataTypeKey,
 
         for (tipo, param) in new_type_fields.iter().zip(t.arg_types().unwrap()) {
             let mut map = mono_types.into_iter().collect_vec();
-            map.append(&mut get_generics_and_type(tipo, &param));
+            map.append(&mut get_generic_id_and_type(tipo, &param));
             mono_types = map.into_iter().collect();
         }
 

crates/aiken-lang/src/uplc.rs

@@ -26,7 +26,7 @@ use crate::{
     builder::{
         check_replaceable_opaque_type, check_when_pattern_needs, constants_ir,
         convert_constants_to_data, convert_data_to_type, convert_type_to_data, get_common_ancestor,
-        get_generics_and_type, handle_clause_guard, handle_func_dependencies_ir,
+        get_generic_id_and_type, handle_clause_guard, handle_func_dependencies_ir,
         handle_recursion_ir, list_access_to_uplc, lookup_data_type_by_tipo, monomorphize,
         rearrange_clauses, replace_opaque_type, wrap_validator_args, AssignmentProperties,
         ClauseProperties, DataTypeKey, FuncComponents, FunctionAccessKey,
@@ -3275,18 +3275,24 @@ impl<'a> CodeGenerator<'a> {
                         let param_types = constructor.tipo.arg_types().unwrap();
 
                         let mut mono_types: IndexMap<u64, Arc<Type>> = IndexMap::new();
+                        let mut map = mono_types.into_iter().collect_vec();
 
                         for (index, arg) in function.arguments.iter().enumerate() {
                             if arg.tipo.is_generic() {
-                                let mut map = mono_types.into_iter().collect_vec();
                                 let param_type = &param_types[index];
 
-                                map.append(&mut get_generics_and_type(&arg.tipo, param_type));
-
-                                mono_types = map.into_iter().collect();
+                                map.append(&mut get_generic_id_and_type(&arg.tipo, param_type));
                             }
                         }
 
+                        if function.return_type.is_generic() {
+                            if let Type::Fn { ret, .. } = &*constructor.tipo {
+                                map.append(&mut get_generic_id_and_type(&function.return_type, ret))
+                            }
+                        }
+
+                        mono_types = map.into_iter().collect();
+
                         let (variant_name, func_ir) =
                             monomorphize(func_ir, mono_types, &constructor.tipo);
 
@@ -3347,23 +3353,32 @@ impl<'a> CodeGenerator<'a> {
                                     } else if let (Some(function), Type::Fn { .. }) =
                                         (function, &*tipo)
                                     {
+                                        let param_types = tipo.arg_types().unwrap();
+
                                         let mut mono_types: IndexMap<u64, Arc<Type>> =
                                             IndexMap::new();
-
-                                        let param_types = tipo.arg_types().unwrap();
+                                        let mut map = mono_types.into_iter().collect_vec();
 
                                         for (index, arg) in function.arguments.iter().enumerate() {
                                             if arg.tipo.is_generic() {
-                                                let mut map = mono_types.into_iter().collect_vec();
-                                                map.append(&mut get_generics_and_type(
-                                                    &arg.tipo,
-                                                    &param_types[index],
-                                                ));
+                                                let param_type = &param_types[index];
 
-                                                mono_types = map.into_iter().collect();
+                                                map.append(&mut get_generic_id_and_type(
+                                                    &arg.tipo, param_type,
+                                                ));
                                             }
                                         }
 
+                                        if function.return_type.is_generic() {
+                                            if let Type::Fn { ret, .. } = &*constructor.tipo {
+                                                map.append(&mut get_generic_id_and_type(
+                                                    &function.return_type,
+                                                    ret,
+                                                ))
+                                            }
+                                        }
+
+                                        mono_types = map.into_iter().collect();
                                         let mut func_ir = vec![];
 
                                         self.build_ir(&function.body, &mut func_ir, scope.to_vec());

examples/acceptance_tests/074/aiken.lock

@@ -0,0 +1,13 @@
+# This file was generated by Aiken
+# You typically do not need to edit this file
+
+[[requirements]]
+name = "aiken-lang/stdlib"
+version = "main"
+source = "github"
+
+[[packages]]
+name = "aiken-lang/stdlib"
+version = "main"
+requirements = []
+source = "github"

examples/acceptance_tests/074/aiken.toml

@@ -0,0 +1,7 @@
+name = "aiken-lang/acceptance_test_074"
+version = "0.0.0"
+
+[[dependencies]]
+name = 'aiken-lang/stdlib'
+version = 'main'
+source = 'github'

examples/acceptance_tests/074/lib/tests.ak

@@ -0,0 +1,282 @@
+use aiken/bytearray.{from_string}
+use aiken/hash.{Hash, Sha2_256, sha2_256}
+use aiken/list
+use aiken/option.{choice, is_none}
+
+/// Variant of MerkleTree with only hash but without actual value
+pub type MerkleTree<a> {
+  Empty
+  Leaf { hash: Hash<Sha2_256, ByteArray> }
+  Node {
+    hash: Hash<Sha2_256, ByteArray>,
+    left: MerkleTree<a>,
+    right: MerkleTree<a>,
+  }
+}
+
+pub type Proof =
+  List<ProofItem>
+
+pub type ProofItem {
+  Left { hash: Hash<Sha2_256, ByteArray> }
+  Right { hash: Hash<Sha2_256, ByteArray> }
+}
+
+// Function returning a hash of a given Merkle Tree element
+pub fn root_hash(self: MerkleTree<a>) -> Hash<Sha2_256, ByteArray> {
+  when self is {
+    Empty -> #""
+    Leaf { hash } -> hash
+    Node { hash, .. } -> hash
+  }
+}
+
+/// Function atests whether two Merkle Tress are equal, this is the case when their root hashes match.
+pub fn is_equal(left: MerkleTree<a>, right: MerkleTree<a>) -> Bool {
+  root_hash(left) == root_hash(right)
+}
+
+/// Function returns a total numbers of leaves in the tree.
+pub fn size(self: MerkleTree<a>) -> Int {
+  when self is {
+    Empty -> 0
+    Leaf{..} -> 1
+    Node { left, right, .. } -> size(left) + size(right)
+  }
+}
+
+fn combine_hash(
+  left: Hash<Sha2_256, a>,
+  right: Hash<Sha2_256, a>,
+) -> Hash<Sha2_256, a> {
+  sha2_256(bytearray.concat(left, right))
+}
+
+/// Function that returns whether merkle tree has any elements
+pub fn is_empty(self: MerkleTree<a>) -> Bool {
+  when self is {
+    Empty -> True
+    _ -> False
+  }
+}
+
+fn do_proof(
+  self: MerkleTree<a>,
+  item_hash: Hash<Sha2_256, ByteArray>,
+  proof: Proof,
+  serialise_fn: fn(a) -> ByteArray,
+) -> Option<Proof> {
+  when self is {
+    Empty -> None
+    Leaf { hash } ->
+      if hash == item_hash {
+        Some(proof)
+      } else {
+        None
+      }
+    Node { left, right, .. } -> {
+      let rh = root_hash(right)
+      let lh = root_hash(left)
+      let go_left: Option<Proof> =
+        do_proof(
+          left,
+          item_hash,
+          list.push(proof, Right { hash: rh }),
+          serialise_fn,
+        )
+      let go_right: Option<Proof> =
+        do_proof(
+          right,
+          item_hash,
+          list.push(proof, Left { hash: lh }),
+          serialise_fn,
+        )
+      choice([go_left, go_right])
+    }
+  }
+}
+
+/// Construct a membership 'Proof' from an element and a 'MerkleTree'. Returns
+/// 'None' if the element isn't a member of the tree to begin with.
+/// Note function will return Some([]) in case root of the tree is also it's only one and only element
+pub fn get_proof(
+  self: MerkleTree<a>,
+  item: a,
+  serialise_fn: fn(a) -> ByteArray,
+) -> Option<Proof> {
+  let empty: Proof = []
+
+  do_proof(self, sha2_256(serialise_fn(item)), empty, serialise_fn)
+}
+
+fn do_from_list(
+  items: List<a>,
+  len: Int,
+  serialise_fn: fn(a) -> ByteArray,
+) -> MerkleTree<a> {
+  when items is {
+    [] -> Empty
+    [item] -> {
+      let hashed_item = sha2_256(serialise_fn(item))
+      Leaf { hash: hashed_item }
+    }
+    all -> {
+      let cutoff: Int = len / 2
+      let left =
+        all
+        |> list.take(cutoff)
+        |> do_from_list(cutoff, serialise_fn)
+      let right =
+        all
+        |> list.drop(cutoff)
+        |> do_from_list(len - cutoff, serialise_fn)
+      let hash = combine_hash(root_hash(left), root_hash(right))
+      Node { hash, left, right }
+    }
+  }
+}
+
+/// Construct a 'MerkleTree' from a list of elements.
+/// Note that, while this operation is doable on-chain, it is expensive and
+/// preferably done off-chain.
+pub fn from_list(
+  items: List<a>,
+  serialise_fn: fn(a) -> ByteArray,
+) -> MerkleTree<a> {
+  do_from_list(items, list.length(items), serialise_fn)
+}
+
+fn do_from_hashes_list(
+  items: List<Hash<Sha2_256, a>>,
+  len: Int,
+) -> MerkleTree<a> {
+  when items is {
+    [] -> Empty
+    [hashed_item] -> Leaf { hash: hashed_item }
+    all -> {
+      let cutoff: Int = len / 2
+      let left =
+        all
+        |> list.take(cutoff)
+        |> do_from_hashes_list(cutoff)
+      let right =
+        all
+        |> list.drop(cutoff)
+        |> do_from_hashes_list(len - cutoff)
+      let hash = combine_hash(root_hash(left), root_hash(right))
+      Node { hash, left, right }
+    }
+  }
+}
+
+/// Construct a 'MerkleTree' from a list of hashes.
+/// Note that, while this operation is doable on-chain, it is expensive and
+/// preferably done off-chain.
+pub fn from_hashes_list(items: List<Hash<Sha2_256, a>>) -> MerkleTree<a> {
+  do_from_hashes_list(items, list.length(items))
+}
+
+// Check whether a hashed element is part of a 'MerkleTree' using only its root hash
+// and a 'Proof'. The proof is guaranteed to be in log(n) of the size of the
+// tree, which is why we are interested in such data-structure in the first
+// place.
+pub fn member_from_hash(
+  item_hash: Hash<Sha2_256, a>,
+  root_hash: Hash<Sha2_256, a>,
+  proof: Proof,
+  serialise_fn: fn(a) -> ByteArray,
+) -> Bool {
+  when proof is {
+    [] -> root_hash == item_hash
+    [head, ..tail] ->
+      when head is {
+        Left { hash: l } ->
+          member_from_hash(
+            combine_hash(l, item_hash),
+            root_hash,
+            tail,
+            serialise_fn,
+          )
+        Right { hash: r } ->
+          member_from_hash(
+            combine_hash(item_hash, r),
+            root_hash,
+            tail,
+            serialise_fn,
+          )
+      }
+  }
+}
+
+// Check whether an element is part of a 'MerkleTree' using only its root hash
+// and a 'Proof'.
+pub fn member(
+  item: a,
+  root_hash: Hash<Sha2_256, ByteArray>,
+  proof: Proof,
+  serialise_fn: fn(a) -> ByteArray,
+) -> Bool {
+  let item_hash = sha2_256(serialise_fn(item))
+  member_from_hash(item_hash, root_hash, proof, serialise_fn)
+}
+
+pub fn member_from_tree(
+  tree: MerkleTree<a>,
+  item: a,
+  serialise_fn: fn(a) -> ByteArray,
+) -> Bool {
+  let proof: Option<Proof> = get_proof(tree, item, serialise_fn)
+  let rh = root_hash(tree)
+
+  when proof is {
+    Some(p) -> member(item, rh, p, serialise_fn)
+    None -> False
+  }
+}
+
+// needed only for tests
+fn create_string_item_serialise_fn() -> fn(String) -> ByteArray {
+  fn(x: String) { from_string(x) }
+}
+
+test from_hashes_list_5() {
+  let dog = @"dog"
+  let cat = @"cat"
+  let mouse = @"mouse"
+  let horse = @"horse"
+
+  let serialise_fn = create_string_item_serialise_fn()
+
+  let items = [dog, cat, mouse, horse]
+  let hashes_items = list.map(items, fn(item) { sha2_256(serialise_fn(item)) })
+
+  let mt = from_hashes_list(hashes_items)
+
+  let left_node_hash =
+    sha2_256(
+      bytearray.concat(sha2_256(serialise_fn(dog)), sha2_256(serialise_fn(cat))),
+    )
+  let right_node_hash =
+    sha2_256(
+      bytearray.concat(
+        sha2_256(serialise_fn(mouse)),
+        sha2_256(serialise_fn(horse)),
+      ),
+    )
+
+  let root_hash = sha2_256(bytearray.concat(left_node_hash, right_node_hash))
+
+  Node {
+    hash: root_hash,
+    left: Node {
+      hash: left_node_hash,
+      left: Leaf { hash: sha2_256(serialise_fn(dog)) },
+      right: Leaf { hash: sha2_256(serialise_fn(cat)) },
+    },
+    right: Node {
+      hash: right_node_hash,
+      left: Leaf { hash: sha2_256(serialise_fn(mouse)) },
+      right: Leaf { hash: sha2_256(serialise_fn(horse)) },
+    },
+  } == mt
+}