From b9d93a0e35c2014224032e55c80cadf7fa0e5b1d Mon Sep 17 00:00:00 2001
From: Anna Henningsen <anna.henningsen@mongodb.com>
Date: Thu, 8 Apr 2021 10:38:18 +0200
Subject: [PATCH] chore(async-rewriter2): add handling for
 Array.prototype.{sort,flatMap}

---
 .../src/async-writer-babel.spec.ts            | 44 +++++++++++++++++++
 .../src/runtime-support.nocov.js              | 26 ++++++++++-
 2 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/packages/async-rewriter2/src/async-writer-babel.spec.ts b/packages/async-rewriter2/src/async-writer-babel.spec.ts
index 5d951bec79..32fe8d8461 100644
--- a/packages/async-rewriter2/src/async-writer-babel.spec.ts
+++ b/packages/async-rewriter2/src/async-writer-babel.spec.ts
@@ -386,6 +386,14 @@ describe('AsyncWriter', () => {
           for (const value of gen) return value;
         })()`)).to.throw('[ASYNC-10012] Result of expression "implicitlyAsyncFn()" cannot be used in this context');
       });
+
+      it('cannot implicitly await inside of array.sort() callback', () => {
+        implicitlyAsyncFn.callsFake((x, y) => x.a - y.a);
+        expect(() => runTranspiledCode(`
+        const arr = [{ a: 2 }, { a : 1 }];
+        arr.sort((x, y) => implicitlyAsyncFn(x, y));
+        `)).to.throw('[ASYNC-10012] Result of expression "compareFn(...args)" cannot be used in this context');
+      });
     });
   });
 
@@ -523,6 +531,15 @@ describe('AsyncWriter', () => {
         expect(implicitlyAsyncFn).to.have.been.calledWith(4, 4, set);
         expect(implicitlyAsyncFn).to.have.been.calledWith(6, 6, set);
       });
+
+      it('supports Array.prototype.flatMap', async() => {
+        implicitlyAsyncFn.callsFake(x => [ x - 1, x ]);
+        const arr = await runTranspiledCode(`
+          const arr = [ 2, 4, 6, 8 ];
+          arr.flatMap(implicitlyAsyncFn)
+        `);
+        expect(arr).to.deep.equal([1, 2, 3, 4, 5, 6, 7, 8]);
+      });
     });
 
     context('synchronous', () => {
@@ -630,6 +647,33 @@ describe('AsyncWriter', () => {
         expect(plainFn).to.have.been.calledWith(4, 4, set);
         expect(plainFn).to.have.been.calledWith(6, 6, set);
       });
+
+      it('supports Array.prototype.flatMap', () => {
+        plainFn.callsFake(x => [ x - 1, x ]);
+        const arr = runTranspiledCode(`
+          const arr = [ 2, 4, 6, 8 ];
+          arr.flatMap(plainFn)
+        `);
+        expect(arr).to.deep.equal([1, 2, 3, 4, 5, 6, 7, 8]);
+      });
+
+      it('supports Array.prototype.sort', () => {
+        plainFn.callsFake((x, y) => x.a - y.a);
+        const arr = runTranspiledCode(`
+          const arr = [ { a: 1 }, { a: 9 }, { a: 4 }, { a: 16 } ];
+          arr.sort(plainFn)
+        `);
+        expect(arr).to.deep.equal([ { a: 1 }, { a: 4 }, { a: 9 }, { a: 16 } ]);
+      });
+
+      it('supports TypedArray.prototype.sort', () => {
+        plainFn.callsFake((x, y) => x - y);
+        const arr = runTranspiledCode(`
+          const arr = new Uint8Array([1, 9, 4, 16]);
+          arr.sort(plainFn)
+        `);
+        expect(arr).to.deep.equal(new Uint8Array([1, 4, 9, 16]));
+      });
     });
 
     context('Function.prototype.toString', () => {
diff --git a/packages/async-rewriter2/src/runtime-support.nocov.js b/packages/async-rewriter2/src/runtime-support.nocov.js
index 6716eb64b8..9fb33c60f6 100644
--- a/packages/async-rewriter2/src/runtime-support.nocov.js
+++ b/packages/async-rewriter2/src/runtime-support.nocov.js
@@ -451,8 +451,30 @@ module.exports = '(' + function() {
     });
   };
 
-  // Currently Missing: (Typed)Array.prototype.sort
-  // Currently Missing: Array.prototype.flatMap
+  const origArraySort = Array.prototype.sort;
+  Array.prototype.sort = function(compareFn) {
+    return origArraySort.call(this, function(...args) {
+      // (Ab-)use a generator function as one of the places where using
+      // implicit async expression results in an error.
+      return [...(function*() {
+        yield compareFn(...args);
+      })()][0];
+    });
+  };
+  const origTypedArraySort = TypedArray.prototype.sort;
+  TypedArray.prototype.sort = function(compareFn) {
+    return origTypedArraySort.call(this, function(...args) {
+      // (Ab-)use a generator function as one of the places where using
+      // implicit async expression results in an error.
+      return [...(function*() {
+        yield compareFn(...args);
+      })()][0];
+    });
+  };
+
+  Array.prototype.flatMap = function(...args) {
+    return Array.prototype.map.call(this, ...args).flat();
+  };
 
   TypedArray.prototype.reduce = Array.prototype.reduce;
   TypedArray.prototype.reduceRight = Array.prototype.reduceRight;