From de2d030aa97a40491643e45e2bf332d0a4e2e335 Mon Sep 17 00:00:00 2001
From: Dmitry Vasilev <dmitry.vasil@gmail.com>
Date: Thu, 8 Dec 2022 09:42:42 +0800
Subject: [PATCH] refactor children -> __children

---
 src/eval.js  | 24 ++++++++++++++++++++++--
 test/test.js | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+), 2 deletions(-)

diff --git a/src/eval.js b/src/eval.js
index e315e9d..3b8e576 100644
--- a/src/eval.js
+++ b/src/eval.js
@@ -216,7 +216,11 @@ const codegen = (node, cxt, parent) => {
       + do_codegen(node.property)  
       + ']'
   } else if(node.type == 'unary') {
-    return '(' + node.operator + ' ' + do_codegen(node.expr) + ')'
+    if(node.operator == 'await') {
+      return `(await __with_restore_children(${do_codegen(node.expr)}))`
+    } else {
+      return '(' + node.operator + ' ' + do_codegen(node.expr) + ')'
+    }
   } else if(node.type == 'binary'){
     return ''
       + do_codegen(node.args[0])
@@ -267,7 +271,7 @@ export const eval_modules = (
   calltree_changed_token,
   location
 ) => {
-  // TODO gensym __modules, __exports
+  // TODO gensym __modules, __exports, children
 
   // TODO bug if module imported twice, once as external and as regular
 
@@ -365,6 +369,22 @@ export const eval_modules = (
       }
     }
 
+    const __with_restore_children = async value => {
+      // 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
+      // which will be used after await
+      if(children == null) {
+        children = []
+      }
+      const children_copy = children
+      try {
+        return await value
+      } finally {
+        children = children_copy
+      }
+    }
+
     const trace = (fn, name, argscount, __location, get_closure) => {
       const result = (...args) => {
         if(result.__closure == null) {
diff --git a/test/test.js b/test/test.js
index 6144cc4..5ff6ec6 100644
--- a/test/test.js
+++ b/test/test.js
@@ -2589,6 +2589,15 @@ const y = x()`
     assert_equal(get_deferred_calls(result), null)
   }),
 
+  test('async/await await non promise', async () => {
+    await assert_code_evals_to_async(
+      `
+        await 1
+      `,
+      1
+    )
+  }),
+
   test('async/await return from async function', async () => {
     await assert_code_evals_to_async(
       `
@@ -2688,4 +2697,38 @@ const y = x()`
     )
   }),
 
+  //test('async/await calltree', async () => {
+  //  const i = await test_initial_state_async(`
+  //    const x = () => 1
+  //    const delay = async time => {
+  //      await 1 //Promise.resolve()
+  //      x()
+  //    }
+  //    await delay(3)
+  //    /* TODO
+  //    await Promise.all([
+  //      delay(3),
+  //    ])
+  //    */
+  //  `)
+  //  log(pp_calltree(root_calltree_node(i)))
+  //  assert_equal(root_calltree_node(i).children.length, 1)
+  //}),
+
+  test('async/await logs out of order', async () => {
+    const i = await test_initial_state_async(`
+      const delay = async time => {
+        await new Promise(res => globalThis.setTimeout(res, time*10))
+        console.log(time)
+      }
+      await Promise.all([
+        delay(3),
+        delay(2),
+        delay(1),
+      ])
+    `)
+    const logs = i.logs.logs.map(l => l.args[0])
+    assert_equal(logs, [1,2,3])
+  })
+
 ]