From 9df03992e9b0e03366d66c48a98d100c6ed22707 Mon Sep 17 00:00:00 2001
From: Pavel Lechenko <1770098+pavel-lechenko@users.noreply.github.com>
Date: Wed, 8 Nov 2023 11:30:22 +0300
Subject: [PATCH] fix(msgpack): encode huge objects (#3698)

Co-authored-by: Asher Gomez <ashersaupingomez@gmail.com>
---
 bytes/concat.ts        | 33 +++++++++++++++++++++++++----
 bytes/concat_test.ts   | 48 +++++++++++++++++++++++++++++++++++++++++-
 msgpack/encode.ts      |  2 +-
 msgpack/encode_test.ts | 21 ++++++++++++++++++
 4 files changed, 98 insertions(+), 6 deletions(-)

diff --git a/bytes/concat.ts b/bytes/concat.ts
index d45e5a35c..3ad30a19a 100644
--- a/bytes/concat.ts
+++ b/bytes/concat.ts
@@ -8,18 +8,43 @@
  * const a = new Uint8Array([0, 1, 2]);
  * const b = new Uint8Array([3, 4, 5]);
  * console.log(concat(a, b)); // [0, 1, 2, 3, 4, 5]
+ * ```
  */
-export function concat(...buf: Uint8Array[]): Uint8Array {
+export function concat(buf: Uint8Array[]): Uint8Array;
+export function concat(...buf: Uint8Array[]): Uint8Array;
+export function concat(...buf: (Uint8Array | Uint8Array[])[]): Uint8Array {
+  // No need to concatenate if there is only one element in array or sub-array
+  if (buf.length === 1) {
+    if (!Array.isArray(buf[0])) {
+      return buf[0];
+    } else if (buf[0].length === 1) {
+      return buf[0][0];
+    }
+  }
+
   let length = 0;
   for (const b of buf) {
-    length += b.length;
+    if (Array.isArray(b)) {
+      for (const b1 of b) {
+        length += b1.length;
+      }
+    } else {
+      length += b.length;
+    }
   }
 
   const output = new Uint8Array(length);
   let index = 0;
   for (const b of buf) {
-    output.set(b, index);
-    index += b.length;
+    if (Array.isArray(b)) {
+      for (const b1 of b) {
+        output.set(b1, index);
+        index += b1.length;
+      }
+    } else {
+      output.set(b, index);
+      index += b.length;
+    }
   }
 
   return output;
diff --git a/bytes/concat_test.ts b/bytes/concat_test.ts
index b7e9b97bb..d3ab80e51 100644
--- a/bytes/concat_test.ts
+++ b/bytes/concat_test.ts
@@ -21,7 +21,7 @@ Deno.test("[bytes] concat empty arrays", () => {
   assert(u2 !== joined);
 });
 
-Deno.test("[bytes] concat multiple arrays", () => {
+Deno.test("[bytes] concat multiple Uint8Array", () => {
   const encoder = new TextEncoder();
   const u1 = encoder.encode("Hello ");
   const u2 = encoder.encode("W");
@@ -34,3 +34,49 @@ Deno.test("[bytes] concat multiple arrays", () => {
   assert(u1 !== joined);
   assert(u2 !== joined);
 });
+
+Deno.test("[bytes] concat an array of Uint8Array", () => {
+  const a = [
+    new Uint8Array([0, 1, 2, 3]),
+    new Uint8Array([4, 5, 6]),
+    new Uint8Array([7, 8, 9]),
+  ];
+  const joined = concat(a);
+  const expected = new Uint8Array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]);
+  assertEquals(joined, expected);
+});
+
+Deno.test("[bytes] concat multiple arrays of Uint8Array using spread operator", () => {
+  const a = [new Uint8Array([0, 1, 2, 3]), new Uint8Array([4, 5, 6, 7, 8, 9])];
+  const b = [
+    new Uint8Array([10, 11]),
+    new Uint8Array([12, 13]),
+    new Uint8Array([14, 15]),
+    new Uint8Array([16]),
+    new Uint8Array([17, 18, 19]),
+  ];
+  const joined = concat(...a, ...b);
+  const expected = new Uint8Array([
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+  ]);
+  assertEquals(joined, expected);
+});
diff --git a/msgpack/encode.ts b/msgpack/encode.ts
index cce432b26..8e62aa112 100644
--- a/msgpack/encode.ts
+++ b/msgpack/encode.ts
@@ -50,7 +50,7 @@ const encoder = new TextEncoder();
 export function encode(object: ValueType) {
   const byteParts: Uint8Array[] = [];
   encodeSlice(object, byteParts);
-  return concat(...byteParts);
+  return concat(byteParts);
 }
 
 function encodeFloat64(num: number) {
diff --git a/msgpack/encode_test.ts b/msgpack/encode_test.ts
index cb149df03..6ece7beb6 100644
--- a/msgpack/encode_test.ts
+++ b/msgpack/encode_test.ts
@@ -162,3 +162,24 @@ Deno.test("maps", () => {
   const nestedMap = { "a": -1, "b": 2, "c": "three", "d": null, "e": map1 };
   assertEquals(decode(encode(nestedMap)), nestedMap);
 });
+
+Deno.test("huge array with 100k objects", () => {
+  const bigArray = [];
+  for (let i = 0; i < 100000; i++) {
+    bigArray.push({ a: { i: `${i}` }, i: i });
+  }
+  const bigObject = { a: bigArray };
+
+  assertEquals(decode(encode(bigObject)), bigObject);
+});
+
+Deno.test("huge object with 100k properties", () => {
+  const bigObject = {};
+  for (let i = 0; i < 100000; i++) {
+    const _ = Object.defineProperty(bigObject, `prop_${i}`, {
+      value: i,
+      enumerable: true,
+    });
+  }
+  assertEquals(decode(encode(bigObject)), bigObject);
+});