Add context()

test/benchmark.luau

@@ -4,11 +4,14 @@ local BENCH, START = testkit.benchmark()
 local vide = require "src/init"
 local source = vide.source
 local derive = vide.derive
+local effect = vide.effect
 local indexes = vide.indexes
 local values = vide.values
 local batch = vide.batch
 local cleanup = vide.cleanup
+local untrack = vide.untrack
 local create = vide.create
+local context = vide.context
 
 assert(not vide.strict)
 
@@ -446,6 +449,61 @@ ROOT_BENCH("values() all remove", function()
     src(data)
 end)
 
+TITLE "context()"
+
+ROOT_BENCH("set context", function()
+    local ctx = context()
+
+    for i = 1, START(N) do
+        ctx(i, function() end)
+    end
+end)
+
+ROOT_BENCH("get context (depth=1)", function()
+    local ctx = context()
+
+    local function run()
+        for i = 1, START(N) do
+            ctx()
+        end
+    end
+
+    ctx(1, function()
+        run()
+    end)
+end)
+
+local depth = 10
+ROOT_BENCH(`get context (depth={depth})`, function()
+
+    local ctx = context()
+
+    local function run()
+        for i = 1, START(N) do
+            ctx()
+        end
+    end
+
+    local function nest_effect(fn)
+        untrack(function()
+            effect(fn)
+            return nil
+        end)
+    end
+    
+    local f = run
+    for i = 1, depth - 1 do
+        local f_inner = f
+        f = function()
+            nest_effect(f_inner)
+        end
+    end
+
+    ctx(1, function()
+        f()
+    end)
+end)
+
 N *= 1024
 
 TITLE "aggregate"

test/tests.luau

@@ -2155,6 +2155,108 @@ TEST("read()", wrap_root(function()
     end
 end))
 
+TEST("context()", function()
+    local root = vide.root
+    local context = vide.context
+    local effect = vide.effect
+    local untrack = vide.untrack
+    local show = vide.show
+
+    do CASE "set context"
+        local ctx = context()
+
+        root(function()
+            ctx(1, function()
+                CHECK(ctx() == 1)
+
+                effect(function()
+                    CHECK(ctx() == 1)
+                end)
+            end)
+        end)
+    end
+
+    do CASE "set context outside of scope"
+        local ctx = context()
+
+        local ok = pcall(function()
+            ctx(1, function() end)
+        end)
+        
+        CHECK(not ok)
+    end
+
+    do CASE "get default context"
+        local ctx = context(1)
+
+        CHECK(ctx() == 1)
+
+        root(function()
+            ctx(2, function()
+                CHECK(ctx() == 2)
+            end)
+
+            CHECK(ctx() == 1)
+        end)
+    end
+
+    do CASE "context cascade"
+        local ctx = context(1)
+        local ctx2 = context()
+
+        root(function()
+            ctx(2, function()
+                ctx2(true, function()
+                    show(function() return true end, function()
+                        ctx(3, function()
+                            effect(function()
+                                CHECK(ctx() == 3)
+                                untrack(function()
+                                    effect(function()
+                                        CHECK(ctx() == 3)
+                                    end)
+                                    CHECK(ctx2() == true)
+                                    return {}
+                                end)
+                            end)
+                            CHECK(ctx() == 3)
+                        end)
+                        CHECK(ctx() == 2)
+                        return {}
+                    end)
+                end)
+                CHECK(ctx() == 2)
+            end)
+
+            CHECK(ctx() == 1)
+        end)
+    end
+
+    do CASE "nil context"
+        local ctx = context(nil)
+
+        root(function()
+            CHECK(ctx() == nil)
+            ctx(nil, function()
+                CHECK(ctx() == nil)
+                ctx(true :: any, function()
+                    CHECK(ctx() == true)
+                    ctx(nil, function()
+                        effect(function()
+                            untrack(function()
+                                effect(function()
+                                    CHECK(ctx() == nil)
+                                end)
+                                return {}
+                            end)
+                        end)
+                    end)
+                end)
+            end)
+        end)
+    end
+end)
+
 TEST("nested effects cases", function()
     local vide = require "src/init"
     local source = vide.source