From ff35ff329df13f7e0000b7e033b82a68ac3b070f Mon Sep 17 00:00:00 2001
From: Ahmed Yasin Koculu <koculu@gmail.com>
Date: Mon, 1 Jan 2024 02:32:22 +0100
Subject: [PATCH 1/2] Add await expression handler interface.

---
 src/Topaz.Test/AwaitTests.cs                  | 48 +++++++++++++++++++
 .../Expressions/AwaitExpressionHandler.cs     |  7 ++-
 .../Expressions/AwaitExpressionHandler.cs     | 11 ++++-
 src/Topaz/Directory.Build.props               |  4 +-
 src/Topaz/Interop/IAwaitExpressionHandler.cs  |  8 ++++
 src/Topaz/TopazEngine.cs                      |  3 ++
 src/Topaz/TopazEngineSetup.cs                 |  2 +
 7 files changed, 79 insertions(+), 4 deletions(-)
 create mode 100644 src/Topaz/Interop/IAwaitExpressionHandler.cs

diff --git a/src/Topaz.Test/AwaitTests.cs b/src/Topaz.Test/AwaitTests.cs
index 430e9e2..d8575e1 100644
--- a/src/Topaz.Test/AwaitTests.cs
+++ b/src/Topaz.Test/AwaitTests.cs
@@ -150,6 +150,54 @@ public void AsyncEnumerator()
 ");
         Assert.That(model.result2, Is.EqualTo(5));
     }
+
+    [Test]
+    public void CustomAwaitHandler()
+    {
+        var engine = new TopazEngine(new TopazEngineSetup { AwaitExpressionHandler = new CustomAwaitExpressionHandler() });
+        dynamic model = new JsObject();
+        engine.SetValue("test", this);
+        engine.SetValue("model", model);
+        engine.ExecuteScriptAsync(@"
+model.result1 = await test.GenericTask();
+").Wait();
+        Assert.That(model.result1, Is.EqualTo(33));
+        engine.ExecuteScript(@"
+model.result2 = await test.GenericTask();
+");
+        Assert.That(model.result2, Is.EqualTo(33));
+    }
+
+    [Test]
+    public void CustomAwaitHandler2()
+    {
+        var engine = new TopazEngine(new TopazEngineSetup { AwaitExpressionHandler = new CustomAwaitExpressionHandler() });
+        dynamic model = new JsObject();
+        engine.SetValue("test", this);
+        engine.SetValue("model", model);
+        engine.ExecuteScriptAsync(@"
+var t = test.GenericTask();
+model.result1 = await t;
+").Wait();
+        Assert.That(model.result1, Is.EqualTo(33));
+        engine.ExecuteScript(@"
+var t = test.GenericTask();
+model.result2 = await t;
+");
+        Assert.That(model.result2, Is.EqualTo(33));
+    }
+}
+
+class CustomAwaitExpressionHandler : IAwaitExpressionHandler
+{
+    public async Task<object> HandleAwaitExpression(object awaitObject)
+    {
+        if (awaitObject is Task<int> task)
+        {
+            return await task;
+        }
+        return awaitObject;
+    }
 }
 
 class CustomMemberInfoProvider : IMemberInfoProvider
diff --git a/src/Topaz/AstHandlers/AsyncHandlers/Expressions/AwaitExpressionHandler.cs b/src/Topaz/AstHandlers/AsyncHandlers/Expressions/AwaitExpressionHandler.cs
index af9d14e..04ed3b3 100644
--- a/src/Topaz/AstHandlers/AsyncHandlers/Expressions/AwaitExpressionHandler.cs
+++ b/src/Topaz/AstHandlers/AsyncHandlers/Expressions/AwaitExpressionHandler.cs
@@ -10,9 +10,14 @@ internal static partial class AwaitExpressionHandler
     internal async static ValueTask<object> ExecuteAsync(ScriptExecutor scriptExecutor, Node expression, CancellationToken token)
     {
         var expr = (AwaitExpression)expression;
-        var result = await scriptExecutor.ExecuteStatementAsync(expr.Argument, token);
+        var result = await scriptExecutor.ExecuteExpressionAndGetValueAsync(expr.Argument, token);
         if (result == null)
             return null;
+        var awaitHandler = scriptExecutor.TopazEngine.AwaitExpressionHandler;
+        if (awaitHandler != null)
+        {
+            return await awaitHandler.HandleAwaitExpression(result);
+        }
         if (result is Task task)
         {
             var type = task.GetType();
diff --git a/src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs b/src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs
index 40404ea..27dadc5 100644
--- a/src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs
+++ b/src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs
@@ -11,7 +11,16 @@ internal static partial class AwaitExpressionHandler
     internal static object Execute(ScriptExecutor scriptExecutor, Node expression, CancellationToken token)
     {
         var expr = (AwaitExpression)expression;
-        var result = scriptExecutor.ExecuteStatement(expr.Argument, token);
+        var result = scriptExecutor.ExecuteExpressionAndGetValue(expr.Argument, token);
+        var awaitHandler = scriptExecutor.TopazEngine.AwaitExpressionHandler;
+        if (result == null)
+            return null;
+        if (awaitHandler != null)
+        {
+            var task1 = awaitHandler.HandleAwaitExpression(result);
+            task1.Wait(token);
+            return task1.Result;
+        }
         if (result is Task task)
         {
             task.Wait(token);
diff --git a/src/Topaz/Directory.Build.props b/src/Topaz/Directory.Build.props
index 885604e..28d21f0 100644
--- a/src/Topaz/Directory.Build.props
+++ b/src/Topaz/Directory.Build.props
@@ -5,8 +5,8 @@
     <Authors>Ahmed Yasin Koculu</Authors>
     <PackageId>Topaz</PackageId>
     <Title>Topaz</Title>
-    <ProductVersion>1.3.8.0</ProductVersion>
-    <Version>1.3.8.0</Version>
+    <ProductVersion>1.3.9.0</ProductVersion>
+    <Version>1.3.9.0</Version>
     <Authors>Ahmed Yasin Koculu</Authors>
     <AssemblyTitle>Topaz</AssemblyTitle>
     <Description>Multithreaded Javascript Engine for .NET</Description>
diff --git a/src/Topaz/Interop/IAwaitExpressionHandler.cs b/src/Topaz/Interop/IAwaitExpressionHandler.cs
new file mode 100644
index 0000000..d48c786
--- /dev/null
+++ b/src/Topaz/Interop/IAwaitExpressionHandler.cs
@@ -0,0 +1,8 @@
+using System.Threading.Tasks;
+
+namespace Tenray.Topaz.Interop;
+
+public interface IAwaitExpressionHandler
+{
+    Task<object> HandleAwaitExpression(object awaitObject);
+}
\ No newline at end of file
diff --git a/src/Topaz/TopazEngine.cs b/src/Topaz/TopazEngine.cs
index 992f086..e20d10c 100644
--- a/src/Topaz/TopazEngine.cs
+++ b/src/Topaz/TopazEngine.cs
@@ -37,6 +37,8 @@ public sealed class TopazEngine : ITopazEngine
 
     public IMemberInfoProvider MemberInfoProvider { get; }
 
+    public IAwaitExpressionHandler AwaitExpressionHandler { get; }
+
     internal ScriptExecutorPool ScriptExecutorPool = new();
 
     public TopazEngine(TopazEngineSetup setup = null)
@@ -52,6 +54,7 @@ public TopazEngine(TopazEngineSetup setup = null)
         ObjectProxyRegistry = setup.ObjectProxyRegistry ?? new DictionaryObjectProxyRegistry();
         extensionMethodRegistry = new();
         MemberInfoProvider = setup.MemberInfoProvider ?? new MemberInfoProvider();
+        AwaitExpressionHandler = setup.AwaitExpressionHandler;
         ValueConverter = setup.ValueConverter ?? new DefaultValueConverter();
         DefaultObjectProxy = setup.DefaultObjectProxy ?? new ObjectProxyUsingReflection(null, extensionMethodRegistry, ValueConverter, MemberInfoProvider);
         DelegateInvoker = setup.DelegateInvoker ?? new DelegateInvoker(ValueConverter);
diff --git a/src/Topaz/TopazEngineSetup.cs b/src/Topaz/TopazEngineSetup.cs
index 42e6b9c..e0384bf 100644
--- a/src/Topaz/TopazEngineSetup.cs
+++ b/src/Topaz/TopazEngineSetup.cs
@@ -24,4 +24,6 @@ public sealed class TopazEngineSetup
     public IValueConverter ValueConverter { get; set; }
 
     public IMemberInfoProvider MemberInfoProvider { get; set; }
+
+    public IAwaitExpressionHandler AwaitExpressionHandler { get; set; }
 }

From 44df8aad4903eca1e9feb7fc4de4538436229915 Mon Sep 17 00:00:00 2001
From: Ahmed Yasin Koculu <koculu@gmail.com>
Date: Mon, 1 Jan 2024 02:38:02 +0100
Subject: [PATCH 2/2] Add cancellation token.

---
 src/Topaz.Test/AwaitTests.cs                                 | 3 ++-
 .../AsyncHandlers/Expressions/AwaitExpressionHandler.cs      | 2 +-
 src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs  | 2 +-
 src/Topaz/Interop/IAwaitExpressionHandler.cs                 | 5 +++--
 4 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/Topaz.Test/AwaitTests.cs b/src/Topaz.Test/AwaitTests.cs
index d8575e1..c1ca6cc 100644
--- a/src/Topaz.Test/AwaitTests.cs
+++ b/src/Topaz.Test/AwaitTests.cs
@@ -4,6 +4,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Reflection;
+using System.Threading;
 using System.Threading.Tasks;
 using Tenray.Topaz.API;
 using Tenray.Topaz.Interop;
@@ -190,7 +191,7 @@ public void CustomAwaitHandler2()
 
 class CustomAwaitExpressionHandler : IAwaitExpressionHandler
 {
-    public async Task<object> HandleAwaitExpression(object awaitObject)
+    public async Task<object> HandleAwaitExpression(object awaitObject, CancellationToken token)
     {
         if (awaitObject is Task<int> task)
         {
diff --git a/src/Topaz/AstHandlers/AsyncHandlers/Expressions/AwaitExpressionHandler.cs b/src/Topaz/AstHandlers/AsyncHandlers/Expressions/AwaitExpressionHandler.cs
index 04ed3b3..7c265fb 100644
--- a/src/Topaz/AstHandlers/AsyncHandlers/Expressions/AwaitExpressionHandler.cs
+++ b/src/Topaz/AstHandlers/AsyncHandlers/Expressions/AwaitExpressionHandler.cs
@@ -16,7 +16,7 @@ internal async static ValueTask<object> ExecuteAsync(ScriptExecutor scriptExecut
         var awaitHandler = scriptExecutor.TopazEngine.AwaitExpressionHandler;
         if (awaitHandler != null)
         {
-            return await awaitHandler.HandleAwaitExpression(result);
+            return await awaitHandler.HandleAwaitExpression(result, token);
         }
         if (result is Task task)
         {
diff --git a/src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs b/src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs
index 27dadc5..7c943a2 100644
--- a/src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs
+++ b/src/Topaz/AstHandlers/Expressions/AwaitExpressionHandler.cs
@@ -17,7 +17,7 @@ internal static object Execute(ScriptExecutor scriptExecutor, Node expression, C
             return null;
         if (awaitHandler != null)
         {
-            var task1 = awaitHandler.HandleAwaitExpression(result);
+            var task1 = awaitHandler.HandleAwaitExpression(result, token);
             task1.Wait(token);
             return task1.Result;
         }
diff --git a/src/Topaz/Interop/IAwaitExpressionHandler.cs b/src/Topaz/Interop/IAwaitExpressionHandler.cs
index d48c786..f08506d 100644
--- a/src/Topaz/Interop/IAwaitExpressionHandler.cs
+++ b/src/Topaz/Interop/IAwaitExpressionHandler.cs
@@ -1,8 +1,9 @@
-using System.Threading.Tasks;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Tenray.Topaz.Interop;
 
 public interface IAwaitExpressionHandler
 {
-    Task<object> HandleAwaitExpression(object awaitObject);
+    Task<object> HandleAwaitExpression(object awaitObject, CancellationToken token);
 }
\ No newline at end of file