From 00cc01b34f9de1fd5b053edf581418267ed1886a Mon Sep 17 00:00:00 2001
From: BlackMATov <blackmat@mail.ru>
Date: Tue, 3 Dec 2024 23:36:26 +0700
Subject: [PATCH] add registry.batch_set

---
 README.md                            |  2 ++
 ROADMAP.md                           |  2 +-
 develop/untests/registry_untests.lua | 42 ++++++++++++++++++++++++++++
 evolved/registry.lua                 | 26 +++++++++++++++++
 4 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 462f4f5..a4603c9 100644
--- a/README.md
+++ b/README.md
@@ -55,6 +55,7 @@ registry.entity -> (entity)
 registry.guid -> entity -> (id)
 registry.alive -> entity -> (boolean)
 registry.set -> entity -> entity -> any -> (entity)
+registry.batch_set -> query -> entity -> any -> (integer, integer)
 registry.get -> entity -> entity... -> (any...)
 registry.has -> entity -> entity -> (boolean)
 registry.has_all -> entity -> entity... -> (boolean)
@@ -104,6 +105,7 @@ entity:destroy -> (boolean)
 query:include -> entity... -> query
 query:exclude -> entity... -> query
 query:execute -> ({execution_state? -> chunk?}, execution_state?)
+query:batch_set -> entity -> any -> (integer, integer)
 query:batch_apply -> {any -> any} -> entity -> (integer)
 query:batch_assign -> entity -> any -> (integer)
 query:batch_insert -> entity -> any -> (integer)
diff --git a/ROADMAP.md b/ROADMAP.md
index 9210428..c2b1278 100644
--- a/ROADMAP.md
+++ b/ROADMAP.md
@@ -19,4 +19,4 @@
 - [ ] add inplace vector operations
 - [x] cache chunk lists in batch operations
 - [ ] impl compat.move for 5.1 vanilla lua
-- [ ] add registry.batch_set
\ No newline at end of file
+- [x] add registry.batch_set
\ No newline at end of file
diff --git a/develop/untests/registry_untests.lua b/develop/untests/registry_untests.lua
index c78615f..2635154 100644
--- a/develop/untests/registry_untests.lua
+++ b/develop/untests/registry_untests.lua
@@ -395,6 +395,48 @@ do
     end
 end
 
+do
+    local f1, f2, f3 = evo.registry.entity(), evo.registry.entity(), evo.registry.entity()
+
+    local e1 = evo.registry.entity():set(f1, 10)
+    local e2 = evo.registry.entity():set(f1, 15)
+    local e3 = evo.registry.entity():set(f1, 20):set(f2, 40)
+    local e4 = evo.registry.entity():set(f1, 25):set(f2, 45)
+
+    do
+        local q = evo.registry.query(f2)
+        local assigned, inserted = q:batch_set(f1, 42)
+        assert(assigned == 2 and inserted == 0)
+        assert(e1:get(f1) == 10 and e2:get(f1) == 15 and e3:get(f1) == 42 and e4:get(f1) == 42)
+        assert(e1:get(f2) == nil and e2:get(f2) == nil and e3:get(f2) == 40 and e4:get(f2) == 45)
+    end
+
+    do
+        local q = evo.registry.query(f1)
+        local assigned, inserted = q:batch_set(f1, 21)
+        assert(assigned == 4 and inserted == 0)
+        assert(e1:get(f1) == 21 and e2:get(f1) == 21 and e3:get(f1) == 21 and e4:get(f1) == 21)
+        assert(e1:get(f2) == nil and e2:get(f2) == nil and e3:get(f2) == 40 and e4:get(f2) == 45)
+    end
+
+    do
+        local q = evo.registry.query(f1)
+        local assigned, inserted = q:batch_set(f2, 84)
+        assert(assigned == 2 and inserted == 2)
+        assert(e1:get(f1) == 21 and e2:get(f1) == 21 and e3:get(f1) == 21 and e4:get(f1) == 21)
+        assert(e1:get(f2) == 84 and e2:get(f2) == 84 and e3:get(f2) == 84 and e4:get(f2) == 84)
+    end
+
+    do
+        local q = evo.registry.query(f1, f2)
+        local assigned, inserted = q:batch_set(f3, 22)
+        assert(assigned == 0 and inserted == 4)
+        assert(e1:get(f1) == 21 and e2:get(f1) == 21 and e3:get(f1) == 21 and e4:get(f1) == 21)
+        assert(e1:get(f2) == 84 and e2:get(f2) == 84 and e3:get(f2) == 84 and e4:get(f2) == 84)
+        assert(e1:get(f3) == 22 and e2:get(f3) == 22 and e3:get(f3) == 22 and e4:get(f3) == 22)
+    end
+end
+
 do
     local f1, f2 = evo.registry.entity(), evo.registry.entity()
 
diff --git a/evolved/registry.lua b/evolved/registry.lua
index 47a9fc2..37b339c 100644
--- a/evolved/registry.lua
+++ b/evolved/registry.lua
@@ -698,6 +698,31 @@ function registry.set(entity, fragment, component)
     return entity
 end
 
+---@param query evolved.query
+---@param fragment evolved.entity
+---@param component any
+---@return integer assigned_count
+---@return integer inserted_count
+function registry.batch_set(query, fragment, component)
+    local to_assign_chunks = __execution_stack_acquire()
+    local to_insert_chunks = __execution_stack_acquire()
+
+    for chunk in registry.execute(query) do
+        if chunk.__components[fragment] ~= nil then
+            to_assign_chunks[#to_assign_chunks + 1] = chunk
+        else
+            to_insert_chunks[#to_insert_chunks + 1] = chunk
+        end
+    end
+
+    local assigned_count = __batch_assign(to_assign_chunks, fragment, component)
+    local inserted_count = __batch_insert(to_insert_chunks, fragment, component)
+
+    __execution_stack_release(to_assign_chunks)
+    __execution_stack_release(to_insert_chunks)
+    return assigned_count, inserted_count
+end
+
 ---@param entity evolved.entity
 ---@param ... evolved.entity fragments
 ---@return any ... components
@@ -1284,6 +1309,7 @@ end
 evolved_query_mt.include = registry.include
 evolved_query_mt.exclude = registry.exclude
 evolved_query_mt.execute = registry.execute
+evolved_query_mt.batch_set = registry.batch_set
 evolved_query_mt.batch_apply = registry.batch_apply
 evolved_query_mt.batch_assign = registry.batch_assign
 evolved_query_mt.batch_insert = registry.batch_insert