fix(async-rewriter2): properly stringify functions MONGOSH-664 (#753)

Add a copy of the original source code at the beginning of transformed
functions, so that `Function.prototype.toString` works properly
for them.

Author: addaleax

packages/async-rewriter2/README.md

@@ -92,6 +92,8 @@ made for readability).
 
 ```js
 (() => {
+  // Keep a copy of the original source code for Function.prototype.toString.
+  '<async_rewriter>(() => {\n  return db.test.find().toArray();\n})</>';
   const _syntheticPromise = Symbol.for("@@mongosh.syntheticPromise");
 
   function _markSyntheticPromise(p) {

packages/async-rewriter2/src/async-writer-babel.spec.ts

@@ -578,6 +578,30 @@ describe('AsyncWriter', () => {
         expect(plainFn).to.have.been.calledWith(6, 6, set);
       });
     });
+
+    context('Function.prototype.toString', () => {
+      it('returns the original function source', () => {
+        expect(runTranspiledCode('Function.prototype.toString.call(() => {})'))
+          .to.equal('() => {}');
+        expect(runTranspiledCode('Function.prototype.toString.call(function () {})'))
+          .to.equal('function () {}');
+        expect(runTranspiledCode('Function.prototype.toString.call(async function () {})'))
+          .to.equal('async function () {}');
+        expect(runTranspiledCode('Function.prototype.toString.call(function* () {})'))
+          .to.equal('function* () {}');
+        expect(runTranspiledCode('Function.prototype.toString.call(async function* () {})'))
+          .to.equal('async function* () {}');
+        expect(runTranspiledCode('Function.prototype.toString.call((class { method() {} }).prototype.method)'))
+          .to.equal('method() {}');
+      });
+
+      it('lets us not worry about special characters', () => {
+        expect(runTranspiledCode('Function.prototype.toString.call(() => {\n  method();\n})'))
+          .to.equal('() => {\n    method();\n  }');
+        expect(runTranspiledCode('Function.prototype.toString.call(() => { const 八 = 8; })'))
+          .to.equal('() => {\n    const 八 = 8;\n  }');
+      });
+    });
   });
 
   context('error messages', () => {

packages/async-rewriter2/src/async-writer-babel.ts

@@ -362,6 +362,16 @@ export const makeMaybeAsyncFunctionPlugin = ({ types: t }: { types: typeof Babel
           return;
         }
 
+        const originalSource = path.parent.start !== undefined ?
+          (this.file as any).code.slice(path.parent.start, path.parent.end) :
+          'function () { [unknown code] }';
+        // Encode using UTF-16 + hex encoding so we don't have to worry about
+        // special characters.
+        const encodedOriginalSource =
+          [...originalSource].map(char => char.charCodeAt(0).toString(16).padStart(4, '0')).join('');
+        const originalSourceNode = t.expressionStatement(
+          t.stringLiteral(`<async_rewriter>${encodedOriginalSource}</>`));
+
         // A parent function might have a set of existing helper methods.
         // If it does, we re-use the functionally equivalent ones.
         const existingIdentifiers: AsyncFunctionIdentifiers | null =
@@ -436,6 +446,7 @@ export const makeMaybeAsyncFunctionPlugin = ({ types: t }: { types: typeof Babel
           // a re-throwing try/catch around the body so that we can perform
           // error message adjustment through the CatchClause handler below.
           path.replaceWith(t.blockStatement([
+            originalSourceNode,
             ...promiseHelpers,
             rethrowTemplate({
               ORIGINAL_CODE: path.node.body
@@ -464,6 +475,7 @@ export const makeMaybeAsyncFunctionPlugin = ({ types: t }: { types: typeof Babel
 
         // Generate the wrapper function. See the README for a full code snippet.
         path.replaceWith(t.blockStatement([
+          originalSourceNode,
           ...promiseHelpers,
           ...wrapperFunction
         ]));

packages/async-rewriter2/src/runtime-support.nocov.js

@@ -468,4 +468,17 @@ module.exports = '(' + function() {
     const array = Array.prototype.filter.call(this, func, thisArg);
     return new (this.constructor)(array);
   };
+
+  // Special addition: Function.prototype.toString!
+  const origFptS = Function.prototype.toString;
+  Function.prototype.toString = function() {
+    const source = origFptS.call(this, arguments);
+    const match = source.match(/^[^"]*"<async_rewriter>(?<encoded>[a-z0-9]+)<\/>";/);
+    if (match) {
+      // Decode using hex + UTF-16
+      return String.fromCharCode(
+        ...match.groups.encoded.match(/.{4}/g).map(hex => parseInt(hex, 16)));
+    }
+    return source;
+  };
 } + ')();';

packages/cli-repl/test/e2e.spec.ts

@@ -342,6 +342,28 @@ describe('e2e', function() {
       shell.assertNoErrors();
     });
 
+    it('rewrites async properly for mapReduce', async() => {
+      // This is being run under the new async rewriter because the old one
+      // did not support mapReduce at all (because of needing 'this').
+      // Once the new async rewriter is the default, this block can be removed.
+      shell = TestShell.start({
+        args: [ await testServer.connectionString() ],
+        env: { ...process.env, MONGOSH_ASYNC_REWRITER2: '1' }
+      });
+      await shell.waitForPrompt();
+      shell.assertNoErrors();
+
+      await shell.executeLine(`use ${dbName}`);
+      await shell.executeLine('db.test.insertMany([{i:1},{i:2},{i:3},{i:4}]);');
+      const result = await shell.executeLine(`db.test.mapReduce(function() {
+        emit(this.i % 2, this.i);
+      }, function(key, values) {
+        return Array.sum(values);
+      }, { out: { inline: 1 } }).results`);
+      expect(result).to.include('{ _id: 0, value: 6 }');
+      expect(result).to.include('{ _id: 1, value: 4 }');
+    });
+
     it('expands explain output indefinitely', async() => {
       await shell.executeLine('explainOutput = db.test.find().explain()');
       await shell.executeLine('explainOutput.a = {b:{c:{d:{e:{f:{g:{h:{i:{j:{}}}}}}}}}}');