From 583e34f71a7c2d7db87fd870532abac983f749e8 Mon Sep 17 00:00:00 2001
From: Dmitry Vasilev <dmitry.vasil@gmail.com>
Date: Mon, 5 Jun 2023 15:53:08 +0300
Subject: [PATCH] fix deferred calls bug

---
 src/effects.js   |  2 +-
 src/eval.js      |  2 +-
 src/record_io.js |  2 +-
 src/runtime.js   | 14 +++++++-------
 test/test.js     | 25 +++++++++++++++++++++++++
 5 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/src/effects.js b/src/effects.js
index 0bd5e50..6b1c9f5 100644
--- a/src/effects.js
+++ b/src/effects.js
@@ -183,7 +183,7 @@ export const render_common_side_effects = (prev, next, command, ui) => {
     next.eval_modules_state != null
   ) {
     const s = next.eval_modules_state
-    s.promise.then(result => {
+    s.promise.__original_then(result => {
       exec('eval_modules_finished', 
         next, /* becomes prev_state */
         result, 
diff --git a/src/eval.js b/src/eval.js
index 4283531..ba8a68c 100644
--- a/src/eval.js
+++ b/src/eval.js
@@ -362,7 +362,7 @@ export const eval_modules = (
   })
 
   if(is_async) {
-    return result.then(make_result)
+    return result.__original_then(make_result)
   } else {
     return make_result(result)
   }
diff --git a/src/record_io.js b/src/record_io.js
index a715bb6..ca1eaea 100644
--- a/src/record_io.js
+++ b/src/record_io.js
@@ -77,7 +77,7 @@ const io_patch = (path, use_context = false) => {
           : original.apply(this, args)
 
         if(value instanceof cxt.window.Promise) {
-          // TODO use cxt.promise_then, not finally which calls
+          // TODO use __original_then, not finally which calls
           // patched 'then'?
           value = value.finally(() => {
             if(cxt_copy != cxt) {
diff --git a/src/runtime.js b/src/runtime.js
index ba71120..4d8d6cf 100644
--- a/src/runtime.js
+++ b/src/runtime.js
@@ -16,8 +16,7 @@ const gen_to_promise = gen_fn => {
       } else {
         // If promise
         if(result.value?.then != null) {
-          return result.value.then.__original.call(
-            result.value,
+          return result.value.__original_then(
             value => next(gen.next(value)),
             error => next(gen.throw(error)),
           )
@@ -113,6 +112,10 @@ const apply_promise_patch = cxt => {
 
   cxt.promise_then = cxt.window.Promise.prototype.then
 
+  if(cxt.window.Promise.prototype.__original_then == null) {
+    cxt.window.Promise.prototype.__original_then = cxt.window.Promise.prototype.then
+  }
+
   cxt.window.Promise.prototype.then = function then(on_resolve, on_reject) {
 
     if(cxt.children == null) {
@@ -135,14 +138,11 @@ const apply_promise_patch = cxt => {
           }
         }
 
-    return cxt.promise_then.call(
-      this,
+    return this.__original_then(
       make_callback(on_resolve, true),
       make_callback(on_reject, false),
     )
   }
-
-  cxt.window.Promise.prototype.then.__original = cxt.promise_then
 }
 
 const remove_promise_patch = cxt => {
@@ -266,7 +266,7 @@ const __do_await = async (cxt, value) => {
   }
   const children_copy = cxt.children
   if(value instanceof cxt.window.Promise) {
-    cxt.promise_then.call(value,
+    value.__original_then(
       v => {
         value.status = {ok: true, value: v}
       },
diff --git a/test/test.js b/test/test.js
index 290919b..f5d338b 100644
--- a/test/test.js
+++ b/test/test.js
@@ -3050,6 +3050,31 @@ const y = x()`
     assert_equal(moved_state.active_calltree_node.fn.name, 'fn2')
   }),
   
+  test('async/await async deferred call', async () => {
+    const code = `
+      await new Object()
+      export const fn = () => 1
+    `
+    const {state: i, on_deferred_call} = test_deferred_calls_state(code)
+
+    await i.eval_modules_state.promise.__original_then(result => {
+      const s = COMMANDS.eval_modules_finished(
+        i, 
+        i,
+        result, 
+        i.eval_modules_state.node, 
+        i.eval_modules_state.toplevel
+      )
+
+      // Make deferred call
+      s.modules[''].fn()
+      const state = on_deferred_call(s)
+      assert_equal(get_deferred_calls(state).length, 1)
+      assert_equal(get_deferred_calls(state)[0].value, 1)
+    })
+
+  }),
+
   test('async/await await argument bug', async () => {
     await assert_code_evals_to_async(
       `