From fb04d897b335f166e1f3fb42f65deada5dc9a435 Mon Sep 17 00:00:00 2001
From: JothamWong <45916998+JothamWong@users.noreply.github.com>
Date: Tue, 16 Apr 2024 03:45:45 +0800
Subject: [PATCH] Optimise waitgroup

---
 src/parser/ooga.pegjs        |  3 +--
 src/vm/oogavm-machine.ts     | 15 ++++++++++++++-
 src/vm/oogavm-typechecker.ts |  1 +
 std/ooga-std.ooga            |  3 ++-
 4 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/parser/ooga.pegjs b/src/parser/ooga.pegjs
index 51fcea3..ff9b983 100644
--- a/src/parser/ooga.pegjs
+++ b/src/parser/ooga.pegjs
@@ -1005,7 +1005,6 @@ FunctionDeclaration
       };
     }
 
-// Arnav: Not sure what this is for!!
 FunctionExpression
   = FunctionToken __ id:(Identifier __)?
     "(" __ params:(FormalParameterList __)? ")" __
@@ -1325,4 +1324,4 @@ BreakpointStatement
     return {
       tag: "BreakpointStatement"
     };
-  }
\ No newline at end of file
+  }
diff --git a/src/vm/oogavm-machine.ts b/src/vm/oogavm-machine.ts
index 06442c1..307be51 100644
--- a/src/vm/oogavm-machine.ts
+++ b/src/vm/oogavm-machine.ts
@@ -245,6 +245,9 @@ function timeoutThread() {
 // Handbook for developing new built-in functions
 // Builtins need to return heap addresses, if they dont return any
 // useful value, return either heap.True or heap.Undefined
+
+// Used to either block the thread or voluntarily yield thread without signalling blocked
+let blockedThreadState = false;
 let yieldThreadState = false;
 
 export const builtinMappings = {
@@ -277,6 +280,10 @@ export const builtinMappings = {
     getThreadID: () => {
         return TSValueToAddress(currentThreadId);
     },
+    blockThread: () => {
+        blockedThreadState = true;
+        return null;
+    },
     yieldThread: () => {
         yieldThreadState = true;
         return null;
@@ -1207,9 +1214,15 @@ function runInstruction() {
     const instr = instrs[PC++];
     log(instr);
     microcode[instr.tag](instr);
+
+    if (blockedThreadState) {
+        blockedThreadState = false;
+        blockThread();
+    }
+
     if (yieldThreadState) {
         yieldThreadState = false;
-        blockThread();
+        timeoutThread();
     }
 
     if (!isAtomicSection) {
diff --git a/src/vm/oogavm-typechecker.ts b/src/vm/oogavm-typechecker.ts
index 591a597..b23939f 100644
--- a/src/vm/oogavm-typechecker.ts
+++ b/src/vm/oogavm-typechecker.ts
@@ -117,6 +117,7 @@ const global_type_frame = {
     startAtomic: new FunctionType([], new NullType()),
     endAtomic: new FunctionType([], new NullType()),
     yieldThread: new FunctionType([], new NullType()),
+    blockThread: new FunctionType([], new NullType()),
     getThreadID: new FunctionType([], new IntegerType()),
     oogaError: new FunctionType([], new NullType()),
 };
diff --git a/std/ooga-std.ooga b/std/ooga-std.ooga
index 64dfe3f..f450c1a 100644
--- a/std/ooga-std.ooga
+++ b/std/ooga-std.ooga
@@ -24,6 +24,7 @@ func (wg *WaitGroup) Done() {
 
 func (wg *WaitGroup) Wait() {
     for wg.counter > 0 {
+        yieldThread(); // yield to avoid wasteful loop
     }
 }
 
@@ -45,7 +46,7 @@ func (m *Mutex) Lock() {
         }
         // else
         endAtomic();
-        yieldThread(); // block and loop until can pick up lock
+        blockThread(); // block and loop until can pick up lock
     }
     m.currentThread = getThreadID();
     endAtomic();