Move collectUint8Arrays() to util/async.ts, clean up, fix bugs, add tests

Original: https://github.com/denoland/deno_std/commit/dcad420b92f79311b386ef524bd1487c04400cc0

2 files changed, 66 insertions, 2 deletions

diff --git a/util/async.ts b/util/async.ts
index f9f2477d0..15edb1b42 100644
--- a/util/async.ts
+++ b/util/async.ts
@@ -83,3 +83,28 @@ export class MuxAsyncIterator<T> implements AsyncIterable<T> {
     return this.iterate();
   }
 }
+
+/** Collects all Uint8Arrays from an AsyncIterable and retuns a single
+ * Uint8Array with the concatenated contents of all the collected arrays.
+ */
+export async function collectUint8Arrays(
+  it: AsyncIterable<Uint8Array>
+): Promise<Uint8Array> {
+  const chunks = [];
+  let length = 0;
+  for await (const chunk of it) {
+    chunks.push(chunk);
+    length += chunk.length;
+  }
+  if (chunks.length === 1) {
+    // No need to copy.
+    return chunks[0];
+  }
+  const collected = new Uint8Array(length);
+  let offset = 0;
+  for (let chunk of chunks) {
+    collected.set(chunk, offset);
+    offset += chunk.length;
+  }
+  return collected;
+}
diff --git a/util/async_test.ts b/util/async_test.ts
index c704002d4..adaac1e22 100644
--- a/util/async_test.ts
+++ b/util/async_test.ts
@@ -1,7 +1,7 @@
 // Copyright 2018-2019 the Deno authors. All rights reserved. MIT license.
 import { test, runIfMain } from "../testing/mod.ts";
-import { assertEquals } from "../testing/asserts.ts";
-import { MuxAsyncIterator, deferred } from "./async.ts";
+import { assert, assertEquals, assertStrictEq } from "../testing/asserts.ts";
+import { collectUint8Arrays, deferred, MuxAsyncIterator } from "./async.ts";
 
 test(async function asyncDeferred(): Promise<void> {
   const d = deferred<number>();
@@ -31,4 +31,43 @@ test(async function asyncMuxAsyncIterator(): Promise<void> {
   assertEquals(results.size, 6);
 });
 
+test(async function collectUint8Arrays0(): Promise<void> {
+  async function* gen(): AsyncIterableIterator<Uint8Array> {}
+  const result = await collectUint8Arrays(gen());
+  assert(result instanceof Uint8Array);
+  assertEquals(result.length, 0);
+});
+
+test(async function collectUint8Arrays0(): Promise<void> {
+  async function* gen(): AsyncIterableIterator<Uint8Array> {}
+  const result = await collectUint8Arrays(gen());
+  assert(result instanceof Uint8Array);
+  assertStrictEq(result.length, 0);
+});
+
+test(async function collectUint8Arrays1(): Promise<void> {
+  const buf = new Uint8Array([1, 2, 3]);
+  async function* gen(): AsyncIterableIterator<Uint8Array> {
+    yield buf;
+  }
+  const result = await collectUint8Arrays(gen());
+  assertStrictEq(result, buf);
+  assertStrictEq(result.length, 3);
+});
+
+test(async function collectUint8Arrays4(): Promise<void> {
+  async function* gen(): AsyncIterableIterator<Uint8Array> {
+    yield new Uint8Array([1, 2, 3]);
+    yield new Uint8Array([]);
+    yield new Uint8Array([4, 5]);
+    yield new Uint8Array([6]);
+  }
+  const result = await collectUint8Arrays(gen());
+  assert(result instanceof Uint8Array);
+  assertStrictEq(result.length, 6);
+  for (let i = 0; i < 6; i++) {
+    assertStrictEq(result[i], i + 1);
+  }
+});
+
 runIfMain(import.meta);