fix TLA last expression return, add tests

brigand · Aug 29, 2019 · 216bf9b · 216bf9b
1 parent c8e5c47
commit 216bf9b
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 13 deletions.
diff --git a/src/plugins/js-eval/__tests__/jsEvalPlugin-test.js b/src/plugins/js-eval/__tests__/jsEvalPlugin-test.js
@@ -90,14 +90,28 @@ describe('jsEvalPlugin', () => {
     expect(output).toEqual(`(okay) undefined`);
   });
 
-  it('handles top-level await', async () => {
-    const output = await testEval('n> let x=await `wat`; x // test');
-    expect(output).toEqual(`(okay) 'wat'`);
-  });
+  describe('top-level-await', () => {
+    it('works', async () => {
+      expect([
+        await testEval('n> var x = await Promise.resolve(2n); x'),
+        await testEval('b> var x = await Promise.resolve(2n); x'),
+        await testEval('n> var x = await Promise.resolve(2n); if (x) {}'),
+        await testEval('b> var x = await Promise.resolve(2n); if (x) {}'),
+      ]).toEqual([
+        '(okay) 2n',
+        '(okay) 2n',
+        '(okay) undefined',
+        '(okay) undefined',
+      ])
+    });
 
-  it('handles top-level await with babel', async () => {
-    const output = await testEval('b> await `wat` // test');
-    expect(output).toEqual(`(okay) 'wat'`);
+    it('works with comments', async () => {
+      const output = await testEval('n> let x=await `wat`; x // test');
+      expect(output).toEqual(`(okay) 'wat'`);
+
+      const output2 = await testEval('b> await `wat` // test');
+      expect(output2).toEqual(`(okay) 'wat'`);
+    });
   });
 
   it('works with engine262', async () => {

diff --git a/src/plugins/js-eval/processTopLevelAwait.js b/src/plugins/js-eval/processTopLevelAwait.js
@@ -58,15 +58,15 @@ function processTopLevelAwait(src) {
   }
 
   let last = root.program.body[root.program.body.length - 1];
+
+  // replace last node with a returnStatement of this node, if the last node is an expression
   if (last.type === 'ExpressionStatement') {
-    last = last.expression;
+    root.program.body[root.program.body.length - 1] = {
+      type: 'ReturnStatement',
+      argument: last.expression
+    };
   }
 
-  // replace last node with a returnStatement of this node
-  root.program.body[root.program.body.length - 1] = {
-    type: 'ReturnStatement',
-    argument: last
-  };
 
   const iiafe = {
     type: 'Program',