From d07c705b64c25416b1ff0080cd27b30c2dc2acb7 Mon Sep 17 00:00:00 2001
From: Aaron Smith <83140718+centau@users.noreply.github.com>
Date: Wed, 27 Sep 2023 10:02:09 +0100
Subject: [PATCH] Fix `switch()` not working in strict mode

Fixes #20
---
 src/graph.luau  | 35 ++++++++++++++---------------------
 src/switch.luau |  2 +-
 test/tests.luau | 34 +++++++++++++++++++++++++++++++---
 3 files changed, 46 insertions(+), 25 deletions(-)

diff --git a/src/graph.luau b/src/graph.luau
index 6917462..bb42fbb 100644
--- a/src/graph.luau
+++ b/src/graph.luau
@@ -23,26 +23,17 @@ export type Node<T> =  {
 -- reactive scope stack
 local scopes = { n = 0 } :: { [number]: Node<any>, n: number }
 
--- runs a given callback in a context that Luau does not allow yielding in
-local check_for_yield: <T...>(fn: (T...) -> (), T...) -> (boolean, string?) do
-    local t = { __mode = "kv" }
-    setmetatable(t, t)
+local function ycall<T, U>(fn: (T) -> U, arg: T): (boolean, string|U)
+    local thread = coroutine.create(pcall)
+    local resume_ok, run_ok, result = coroutine.resume(thread, fn, arg)
 
-    check_for_yield = function(fn, ...: any)
-        local args = { ... }
-
-        t.__unm = function(_)
-            fn(unpack(args))
-        end
-
-        local ok, err: string? = pcall(function()
-            local _ = -t
-        end)
-
-        return ok, if err == "attempt to yield across metamethod/C-call boundary"
-            or err == "thread is not yieldable" then "yield occured"
-            else err
+    assert(resume_ok)
+        
+    if coroutine.status(thread) ~= "dead" then
+        return false, "attempt to yield in reactive scope"
     end
+
+    return run_ok, result
 end
 
 local function get_scope(): Node<unknown>?
@@ -161,11 +152,13 @@ local function evaluate_node<T>(node: Node<T>)
 
         open_scope(node)
 
-        local ok, err = check_for_yield(node.effect :: (T) -> T, cur_value)
+        local ok, new_value = ycall(node.effect :: (T) -> T, cur_value)
         
         close_scope()
         
-        if not ok then throw(err :: string) end
+        if not ok then throw(new_value :: string) end
+
+        node.cache = new_value :: T
     end
 
     run_cleanups(node)
@@ -173,7 +166,7 @@ local function evaluate_node<T>(node: Node<T>)
 
     open_scope(node)
 
-    local ok, new_value = pcall(node.effect :: (T) -> T, cur_value)
+    local ok, new_value = pcall(node.effect :: (T) -> T, node.cache)
 
     close_scope()
 
diff --git a/src/switch.luau b/src/switch.luau
index 421d583..1ddc8e6 100644
--- a/src/switch.luau
+++ b/src/switch.luau
@@ -53,7 +53,7 @@ local function switch<T, U>(source: () -> T): (map: Map<T, ((() -> U)?)>) -> ()
             return result
         end
 
-        local node = create_node(nil :: any, update)
+        local node = create_node(nil :: U?, update)
 
         set_owner(node, owner)
         evaluate_node(node)
diff --git a/test/tests.luau b/test/tests.luau
index 73ea0b7..eb5581c 100644
--- a/test/tests.luau
+++ b/test/tests.luau
@@ -1067,9 +1067,6 @@ TEST("show()", wrap_root(function()
         CHECK(outer == 2)
         CHECK(inner == 4)
         CHECK(destroyed == 3)
-
-
-
     end
 end))
 
@@ -1181,6 +1178,22 @@ TEST("switch()", wrap_root(function()
 
         CHECK(n0 == n1)
     end
+
+    do CASE "strict"
+        vide.strict = true
+
+        local input = source(0)
+        local output = switch(input) {
+            [0] = function() return 0 end,
+            [1] = function() return 1 end,
+        }
+
+        CHECK(output() == 0)
+        input(1)
+        CHECK(output() == 1)
+
+        vide.strict = false
+    end
 end))
 
 TEST("indexes()", wrap_root(function()
@@ -1364,6 +1377,21 @@ TEST("indexes()", wrap_root(function()
         CHECK(updated[3] == 2)
         CHECK(updated[4] == 2)
     end
+
+    do CASE "strict"
+        vide.strict = true
+
+        local input = source{1}
+        local output = indexes(input, function(v)
+            return { v }
+        end)
+
+        CHECK(output()[1][1]() == 1)
+        input{2}
+        CHECK(output()[1][1]() == 2)
+
+        vide.strict = false
+    end
 end))
 
 TEST("values()", wrap_root(function()