fix await bug

src/eval.js

@@ -80,6 +80,11 @@ const codegen_function_expr = (node, node_cxt) => {
   // TODO gensym __obj, __fn, __call_id, __let_vars, __literals
   const prolog = 
     '{const __call_id = __cxt.call_counter;' 
+    + (
+        node.is_async
+        ? 'let __await_state;'
+        : ''
+      )
     + (
         node.has_versioned_let_vars
         ? 'const __let_vars = __cxt.let_vars;'
@@ -346,7 +351,9 @@ const codegen = (node, node_cxt) => {
       + ']'
   } else if(node.type == 'unary') {
     if(node.operator == 'await') {
-      return `(await __do_await(__cxt, ${do_codegen(node.expr)}))`
+      return `(__await_state = __await_start(__cxt, ${do_codegen(node.expr)}),` +
+             `await __await_state.promise,` +
+             `__await_finish(__cxt, __await_state))`
     } else {
       return '(' + node.operator + ' ' + do_codegen(node.expr) + ')'
     }
@@ -416,7 +423,7 @@ export const eval_modules = (
   io_trace,
 ) => {
   // TODO gensym __cxt, __trace, __trace_call, __calltree_node_by_loc,
-  // __do_await, __Multiversion, __create_array, __create_object
+  // __await_start, __await_finish, __Multiversion, __create_array, __create_object
 
   // TODO bug if module imported twice, once as external and as regular
 
@@ -446,7 +453,8 @@ export const eval_modules = (
         '__calltree_node_by_loc',
         '__trace',
         '__trace_call',
-        '__do_await',
+        '__await_start',
+        '__await_finish',
         '__save_ct_node_for_path',
         '__Multiversion',
         '__create_array',
@@ -457,6 +465,7 @@ export const eval_modules = (
          * because we dont want to codegen differently for if statements in
          * toplevel and if statements within functions*/
         'const __call_id = __cxt.call_counter;' + 
+        'let __await_state;' +
         code
       )
     }

src/runtime/runtime.js

@@ -99,7 +99,8 @@ const do_run = function*(module_fns, cxt, io_trace){
         calltree_node_by_loc.get(module),
         __trace, 
         __trace_call, 
-        __do_await, 
+        __await_start, 
+        __await_finish,
         __save_ct_node_for_path,
         LetMultiversion,
         create_array,
@@ -250,9 +251,7 @@ export const do_eval_expand_calltree_node = (cxt, node) => {
   }
 }
 
-
-
-const __do_await = async (cxt, value) => {
+const __await_start = (cxt, promise) => {
   // children is an array of child calls for current function call. But it
   // can be null to save one empty array allocation in case it has no child
   // calls. Allocate array now, so we can have a reference to this array
@@ -261,20 +260,33 @@ const __do_await = async (cxt, value) => {
     cxt.children = []
   }
   const children_copy = cxt.children
-  if(value instanceof cxt.window.Promise) {
-    value.__original_then(
-      v => {
-        value.status = {ok: true, value: v}
+  const result = {children_copy, promise}
+
+  if(promise instanceof cxt.window.Promise) {
+    result.promise = promise.then(
+      (value) => {
+        result.status = {ok: true, value}
+        // We do not return value on purpose - it will be return in
+        // __await_finish
+      },
+      (error) => {
+        result.status = {ok: false, error}
+        // We do not throw error on purpose
       },
-      e => {
-        value.status = {ok: false, error: e}
-      }
     )
+  } else {
+    result.status = {ok: true, value: promise}
   }
-  try {
-    return await value
-  } finally {
-    cxt.children = children_copy
+
+  return result
+}
+
+const __await_finish = (__cxt, await_state) => {
+  __cxt.children = await_state.children_copy
+  if(await_state.status.ok) {
+    return await_state.status.value
+  } else {
+    throw await_state.status.error
   }
 }
 

test/test.js

@@ -3695,6 +3695,18 @@ const y = x()`
     )
   }),
 
+  test('async/await await rejected Promise fn call', async () => {
+    await assert_code_error_async(
+      `
+        async function test() {
+          await Promise.reject('boom')
+        }
+        await test()
+      `,
+      'boom'
+    )
+  }),
+
   test('async/await promise rejected with null', async () => {
     await assert_code_error_async(
       `await Promise.reject()`,
@@ -3872,6 +3884,19 @@ const y = x()`
     )
   }),
 
+  test('async/await await bug', async () => {
+    const code = `
+      const x = () => {}
+      const test = async () => {
+        await 1
+        x()
+      }
+      await Promise.all([test(), test()])
+    `
+    const i = await test_initial_state_async(code, code.indexOf('await 1'))
+    assert_value_explorer(i ,1)
+  }),
+
   test('async/await edit', async () => {
     const code = `
       const f = async () => {