fix Array.chain to handle empty arrays, following implementation of A…

…rray.flatten.
dfinity · Nov 14, 2023 · ec125bc · ec125bc
1 parent 46232dd
commit ec125bc
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 8 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
@@ -286,9 +286,39 @@ let suite = Suite.suite(
     ),
     Suite.test(
       "chain mix",
-      Array.chain<Nat, Nat>([0, 1, 2, 0, 1, 2],
+      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,1,0,1]))
+      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(
       "foldLeft",
@@ -331,7 +361,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, []))
     ),