bugfix: Array.chain(as,f) incorrectly trapping when f(a) is an empty …

…array. (#599) Implemented by following same pattern in Array.flatten, adding tests.
dfinity · Nov 14, 2023 · 3d0bced · 3d0bced
1 parent 8028da2
commit 3d0bced
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 6 deletions.
diff --git a/src/Array.mo b/src/Array.mo
@@ -520,28 +520,28 @@ module {
   /// *Runtime and space assumes that `k` runs in O(1) time and space.
   public func chain<X, Y>(array : [X], k : X -> [Y]) : [Y] {
     var flatSize = 0;
-    let subArrays = Prim.Array_tabulate<[Y]>(
+    let arrays = Prim.Array_tabulate<[Y]>(
       array.size(),
       func i {
         let subArray = k(array[i]);
         flatSize += subArray.size();
         subArray
       }
     );
+
     // could replace with a call to flatten,
     // but it would require an extra pass (to compute `flatSize`)
     var outer = 0;
     var inner = 0;
     Prim.Array_tabulate<Y>(
       flatSize,
       func _ {
-        let subArray = subArrays[outer];
-        let element = subArray[inner];
-        inner += 1;
-        if (inner == subArray.size()) {
+        while (inner == arrays[outer].size()) {
           inner := 0;
           outer += 1
         };
+        let element = arrays[outer][inner];
+        inner += 1;
         element
       }
     )

diff --git a/test/Array.test.mo b/test/Array.test.mo
@@ -284,6 +284,54 @@ let suite = Suite.suite(
       Array.chain<Int, Int>([], func x = [x, -x]),
       M.equals(T.array<Int>(T.intTestable, []))
     ),
+    Suite.test(
+      "chain mix",
+      Array.chain<Nat, Nat>([1, 2, 1, 2, 3],
+        func n = Array.tabulate<Nat>(n, func i = i)),
+      M.equals(T.array<Nat>(T.natTestable, [0,0,1,0,0,1,0,1,2]))
+    ),
+    Suite.test(
+      "chain mix empty right",
+      Array.chain<Nat, Nat>([0, 1, 2, 0, 1, 2, 3, 0],
+        func n = Array.tabulate<Nat>(n, func i = i)),
+      M.equals(T.array<Nat>(T.natTestable, [0,0,1,0,0,1,0,1,2]))
+    ),
+    Suite.test(
+      "chain mix empties right",
+      Array.chain<Nat, Nat>([0, 1, 2, 0, 1, 2, 3, 0, 0, 0],
+        func n = Array.tabulate<Nat>(n, func i = i)),
+      M.equals(T.array<Nat>(T.natTestable, [0,0,1,0,0,1,0,1,2]))
+    ),
+    Suite.test(
+      "chain mix empty left",
+      Array.chain<Nat, Nat>([0, 1, 2, 0, 1, 2, 3],
+        func n = Array.tabulate<Nat>(n, func i = i)),
+      M.equals(T.array<Nat>(T.natTestable, [0,0,1,0,0,1,0,1,2]))
+    ),
+    Suite.test(
+      "chain mix empties left",
+      Array.chain<Nat, Nat>([0, 0, 0, 1, 2, 0, 1, 2, 3],
+        func n = Array.tabulate<Nat>(n, func i = i)),
+      M.equals(T.array<Nat>(T.natTestable, [0,0,1,0,0,1,0,1,2]))
+    ),
+    Suite.test(
+      "chain mix empties middle",
+      Array.chain<Nat, Nat>([0, 1, 2, 0, 0, 0, 1, 2, 3],
+        func n = Array.tabulate<Nat>(n, func i = i)),
+      M.equals(T.array<Nat>(T.natTestable, [0,0,1,0,0,1,0,1,2]))
+    ),
+    Suite.test(
+      "chain mix empties",
+      Array.chain<Nat, Nat>([0, 0, 0],
+        func n = Array.tabulate<Nat>(n, func i = i)),
+      M.equals(T.array<Nat>(T.natTestable, []))
+    ),
+    Suite.test(
+      "chain mix empty",
+      Array.chain<Nat, Nat>([],
+        func n = Array.tabulate<Nat>(n, func i = i)),
+      M.equals(T.array<Nat>(T.natTestable, []))
+    ),
     Suite.test(
       "foldLeft",
       Array.foldLeft<Text, Text>(["a", "b", "c"], "", Text.concat),
@@ -325,7 +373,7 @@ let suite = Suite.suite(
       M.equals(T.array<Int>(T.intTestable, [1, 2, 3]))
     ),
     Suite.test(
-      "flatten empty",
+      "flatten singleton empty",
       Array.flatten<Int>([[]]),
       M.equals(T.array<Int>(T.intTestable, []))
     ),