Scott encoding alternative for dict functions

lib/aiken/collection/dict.ak

@@ -13,6 +13,7 @@
 //// > [`dict.from_ascending_list`](#from_ascending_list).
 
 use aiken/builtin
+use aiken/collection/dict/union.{UnionStrategy}
 
 /// An opaque `Dict`. The type is opaque because the module maintains some
 /// invariant, namely: there's only one occurrence of a given key in the dictionary.
@@ -160,7 +161,8 @@ pub fn from_pairs(self: Pairs<ByteArray, value>) -> Dict<key, value> {
 fn do_from_pairs(xs: Pairs<ByteArray, value>) -> Pairs<ByteArray, value> {
   when xs is {
     [] -> []
-    [Pair(k, v), ..rest] -> do_insert(do_from_pairs(rest), k, v)
+    [Pair(k, v), ..rest] ->
+      do_insert_with(do_from_pairs(rest), k, v, union.keep_left())
   }
 }
 
@@ -728,14 +730,13 @@ test insert_2() {
 /// to the merge function, and the new value is passed as the third argument.
 ///
 /// ```aiken
-/// let sum =
-///   fn (_k, a, b) { Some(a + b) }
+/// use aiken/collection/dict/union
 ///
 /// let result =
 ///   dict.empty
-///     |> dict.insert_with(key: "a", value: 1, with: sum)
-///     |> dict.insert_with(key: "b", value: 2, with: sum)
-///     |> dict.insert_with(key: "a", value: 3, with: sum)
+///     |> dict.insert_with(key: "a", value: 1, with: union.sum())
+///     |> dict.insert_with(key: "b", value: 2, with: union.sum())
+///     |> dict.insert_with(key: "a", value: 3, with: union.sum())
 ///     |> dict.to_pairs()
 ///
 /// result == [Pair("a", 4), Pair("b", 2)]
@@ -744,47 +745,46 @@ pub fn insert_with(
   self: Dict<key, value>,
   key k: ByteArray,
   value v: value,
-  with: fn(ByteArray, value, value) -> Option<value>,
+  with: UnionStrategy<ByteArray, value>,
 ) -> Dict<key, value> {
   Dict {
-    inner: do_insert_with(self.inner, k, v, fn(k, v1, v2) { with(k, v2, v1) }),
+    inner: do_insert_with(
+      self.inner,
+      k,
+      v,
+      fn(k, v1, v2, some, none) { with(k, v2, v1, some, none) },
+    ),
   }
 }
 
 test insert_with_1() {
-  let sum =
-    fn(_k, a, b) { Some(a + b) }
-
   let result =
     empty
-      |> insert_with(key: "foo", value: 1, with: sum)
-      |> insert_with(key: "bar", value: 2, with: sum)
+      |> insert_with(key: "foo", value: 1, with: union.sum())
+      |> insert_with(key: "bar", value: 2, with: union.sum())
       |> to_pairs()
 
   result == [Pair("bar", 2), Pair("foo", 1)]
 }
 
 test insert_with_2() {
-  let sum =
-    fn(_k, a, b) { Some(a + b) }
-
   let result =
     empty
-      |> insert_with(key: "foo", value: 1, with: sum)
-      |> insert_with(key: "bar", value: 2, with: sum)
-      |> insert_with(key: "foo", value: 3, with: sum)
+      |> insert_with(key: "foo", value: 1, with: union.sum())
+      |> insert_with(key: "bar", value: 2, with: union.sum())
+      |> insert_with(key: "foo", value: 3, with: union.sum())
       |> to_pairs()
 
   result == [Pair("bar", 2), Pair("foo", 4)]
 }
 
 test insert_with_3() {
   let with =
-    fn(k, a, _b) {
+    fn(k, a, _b, keep, discard) {
       if k == "foo" {
-        Some(a)
+        keep(a)
       } else {
-        None
+        discard()
       }
     }
 
@@ -1011,15 +1011,15 @@ test union_4() {
 pub fn union_with(
   left: Dict<key, value>,
   right: Dict<key, value>,
-  with: fn(ByteArray, value, value) -> Option<value>,
+  with: UnionStrategy<ByteArray, value>,
 ) -> Dict<key, value> {
   Dict { inner: do_union_with(left.inner, right.inner, with) }
 }
 
 fn do_union_with(
   left: Pairs<ByteArray, value>,
   right: Pairs<ByteArray, value>,
-  with: fn(ByteArray, value, value) -> Option<value>,
+  with: UnionStrategy<ByteArray, value>,
 ) -> Pairs<ByteArray, value> {
   when left is {
     [] -> right
@@ -1032,7 +1032,7 @@ fn do_insert_with(
   self: Pairs<ByteArray, value>,
   key k: ByteArray,
   value v: value,
-  with: fn(ByteArray, value, value) -> Option<value>,
+  with: UnionStrategy<ByteArray, value>,
 ) -> Pairs<ByteArray, value> {
   when self is {
     [] -> [Pair(k, v)]
@@ -1041,10 +1041,13 @@ fn do_insert_with(
         [Pair(k, v), ..self]
       } else {
         if k == k2 {
-          when with(k, v, v2) is {
-            Some(combined) -> [Pair(k, combined), ..rest]
-            None -> rest
-          }
+          with(
+            k,
+            v,
+            v2,
+            fn(combined) { [Pair(k, combined), ..rest] },
+            fn() { rest },
+          )
         } else {
           [Pair(k2, v2), ..do_insert_with(rest, k, v, with)]
         }
@@ -1062,7 +1065,7 @@ test union_with_1() {
       |> insert(bar, 42)
       |> insert(foo, 1337)
 
-  let result = union_with(left, right, with: fn(_, l, r) { Some(l + r) })
+  let result = union_with(left, right, with: union.sum())
 
   result == from_pairs([Pair(foo, 1351), Pair(bar, 42)])
 }

lib/aiken/collection/dict/union.ak

@@ -0,0 +1,47 @@
+/// A strategy for combining two values in a dictionnary that belong to the same key.
+pub type UnionStrategy<key, value> =
+  fn(key, value, value, KeepValue<key, value>, DiscardValue<key, value>) ->
+    Pairs<key, value>
+
+/// A callback to keep a combined value at a given key
+pub type KeepValue<key, value> =
+  fn(value) -> Pairs<key, value>
+
+/// A callback to discard a value at a given key
+pub type DiscardValue<key, value> =
+  fn() -> Pairs<key, value>
+
+/// A strategy which always fail, enforcing the dict contains no duplicate.
+pub fn expect_no_duplicate() -> UnionStrategy<key, value> {
+  fn(_, _, _, _, _) {
+    fail @"unexpected duplicate key found in dict."
+  }
+}
+
+/// Combine values by keeping the values present in the left-most dict.
+pub fn keep_left() -> UnionStrategy<key, value> {
+  fn(_key, left, _right, keep, _discard) { keep(left) }
+}
+
+/// Combine values by keeping the values present in the left-most dict.
+pub fn keep_right() -> UnionStrategy<key, value> {
+  fn(_key, _left, right, keep, _discard) { keep(right) }
+}
+
+/// Combine values by taking their sum.
+pub fn sum() -> UnionStrategy<key, Int> {
+  fn(_key, left, right, keep, _discard) { keep(left + right) }
+}
+
+/// Combine values by taking their sum, only if it is non-null. Otherwise,
+/// discard the key/value entirely.
+pub fn sum_if_non_zero() -> UnionStrategy<key, Int> {
+  fn(_key, left, right, keep, discard) {
+    let value = left + right
+    if value == 0 {
+      discard()
+    } else {
+      keep(value)
+    }
+  }
+}

lib/cardano/assets.ak

@@ -1,8 +1,10 @@
 use aiken/builtin
 use aiken/collection/dict.{Dict, from_ascending_pairs_with}
+use aiken/collection/dict/union as dict_union
 use aiken/collection/list
 use aiken/crypto.{Blake2b_224, Hash, Script}
 use aiken/option
+use cardano/assets/union
 
 /// Lovelace is now a type wrapper for Int.
 pub type Lovelace =
@@ -77,14 +79,7 @@ pub fn from_asset_list(xs: Pairs<PolicyId, Pairs<AssetName, Int>>) -> Value {
           expect Pair(p, [_, ..] as x) = inner
           x
             |> from_ascending_pairs_with(fn(v) { v != 0 })
-            |> dict.insert_with(
-                acc,
-                p,
-                _,
-                fn(_, _, _) {
-                  fail @"Duplicate policy in the asset list."
-                },
-              )
+            |> dict.insert_with(acc, p, _, dict_union.expect_no_duplicate())
         },
       )
     |> Value
@@ -555,25 +550,19 @@ pub fn add(
     self
   } else {
     let helper =
-      fn(_, left, _right) {
+      fn(_, left, _right, keep, discard) {
         let inner_result =
           dict.insert_with(
             left,
             asset_name,
             quantity,
-            fn(_k, ql, qr) {
-              let q = ql + qr
-              if q == 0 {
-                None
-              } else {
-                Some(q)
-              }
-            },
+            dict_union.sum_if_non_zero(),
           )
+
         if dict.is_empty(inner_result) {
-          None
+          discard()
         } else {
-          Some(inner_result)
+          keep(inner_result)
         }
       }
 
@@ -633,24 +622,12 @@ pub fn merge(left v0: Value, right v1: Value) -> Value {
     dict.union_with(
       v0.inner,
       v1.inner,
-      fn(_, a0, a1) {
-        let result =
-          dict.union_with(
-            a0,
-            a1,
-            fn(_, q0, q1) {
-              let q = q0 + q1
-              if q == 0 {
-                None
-              } else {
-                Some(q)
-              }
-            },
-          )
+      fn(_, a0, a1, keep, discard) {
+        let result = dict.union_with(a0, a1, dict_union.sum_if_non_zero())
         if dict.is_empty(result) {
-          None
+          discard()
         } else {
-          Some(result)
+          keep(result)
         }
       },
     ),
@@ -818,7 +795,7 @@ pub fn flatten(self: Value) -> List<(PolicyId, AssetName, Int)> {
 /// When the transform function returns `None`, the result is discarded altogether.
 pub fn flatten_with(
   self: Value,
-  with: fn(PolicyId, AssetName, Int) -> Option<result>,
+  with: union.UnionStrategy<result>,
 ) -> List<result> {
   dict.foldr(
     self.inner,
@@ -828,18 +805,21 @@ pub fn flatten_with(
         asset_list,
         value,
         fn(asset_name, quantity, xs) {
-          when with(policy_id, asset_name, quantity) is {
-            None -> xs
-            Some(x) -> [x, ..xs]
-          }
+          with(
+            policy_id,
+            asset_name,
+            quantity,
+            fn(x) { [x, ..xs] },
+            fn() { xs },
+          )
         },
       )
     },
   )
 }
 
 test flatten_with_1() {
-  flatten_with(zero, fn(p, a, q) { Some((p, a, q)) }) == []
+  flatten_with(zero, union.triple()) == []
 }
 
 test flatten_with_2() {
@@ -851,11 +831,11 @@ test flatten_with_2() {
 
   flatten_with(
     v,
-    fn(p, a, q) {
+    fn(p, a, q, keep, discard) {
       if q == 42 {
-        Some((p, a))
+        keep((p, a))
       } else {
-        None
+        discard()
       }
     },
   ) == [("a", "2"), ("b", "")]

lib/cardano/assets/union.ak

@@ -0,0 +1,27 @@
+use aiken/crypto.{Blake2b_224, Hash, Script}
+
+/// A strategy for combining two values in a dictionnary that belong to the same key.
+pub type UnionStrategy<result> =
+  fn(
+    Hash<Blake2b_224, Script>,
+    ByteArray,
+    Int,
+    KeepResult<result>,
+    DiscardResult<result>,
+  ) ->
+    List<result>
+
+/// A callback to keep result at a given key
+pub type KeepResult<result> =
+  fn(result) -> List<result>
+
+/// A callback to discard result at a given key
+pub type DiscardResult<result> =
+  fn() -> List<result>
+
+/// Combine values by keeping them in a 3-tuple.
+pub fn triple() -> UnionStrategy<(Hash<Blake2b_224, Script>, ByteArray, Int)> {
+  fn(policy_id, asset_name, quantity, keep, _discard) {
+    keep((policy_id, asset_name, quantity))
+  }
+}