Merge pull request #962 from grilledham/task+priority_queue_refactor

Task+ PriorityQueue clean up + micro optimisations

utils/priority_queue.lua

@@ -1,20 +1,46 @@
+local Debug = require 'utils.debug'
+local is_closure = Debug.is_closure
+local floor = math.floor
+
 local PriorityQueue = {}
 
-function PriorityQueue.new()
-    return {}
-end
-
-local function default_comp(a, b)
+local function default_comparator(a, b)
     return a < b
 end
 
-local function HeapifyFromEndToStart(queue, comp)
-    comp = comp or default_comp
-    local pos = #queue
+--- Min heap implementation of a priority queue. Smaller elements, as determined by the comparator,
+-- have a higher priority.
+-- @param comparator <function|nil> the comparator function used to compare elements, if nil the
+-- deafult comparator is used.
+-- @usage
+-- local PriorityQueue = require 'utils.priority_queue'
+--
+-- local queue = PriorityQueue.new()
+-- PriorityQueue.push(queue, 4)
+-- PriorityQueue.push(queue, 7)
+-- PriorityQueue.push(queue, 2)
+--
+-- game.print(PriorityQueue.pop(queue)) -- 2
+-- game.print(PriorityQueue.pop(queue)) -- 4
+-- game.print(PriorityQueue.pop(queue)) -- 7
+function PriorityQueue.new(comparator)
+    if comparator == nil then
+        comparator = default_comparator
+    elseif is_closure(comparator) then
+        error('comparator cannot be a closure.', 2)
+    end
+
+    return {_comparator = comparator}
+end
+
+local function heapify_from_end_to_start(self)
+    local comparator = self._comparator
+    local pos = #self
     while pos > 1 do
-        local parent = bit32.rshift(pos, 1) -- integer division by 2
-        if comp(queue[pos], queue[parent]) then
-            queue[pos], queue[parent] = queue[parent], queue[pos]
+        local parent = floor(pos * 0.5)
+        local a, b = self[pos], self[parent]
+        if comparator(a, b) then
+            self[pos], self[parent] = b, a
             pos = parent
         else
             break
@@ -22,25 +48,26 @@ local function HeapifyFromEndToStart(queue, comp)
     end
 end
 
-local function HeapifyFromStartToEnd(queue, comp)
-    comp = comp or default_comp
+local function heapify_from_start_to_end(self)
+    local comparator = self._comparator
     local parent = 1
     local smallest = 1
+    local count = #self
     while true do
         local child = parent * 2
-        if child > #queue then
+        if child > count then
             break
         end
-        if comp(queue[child], queue[parent]) then
+        if comparator(self[child], self[parent]) then
             smallest = child
         end
         child = child + 1
-        if child <= #queue and comp(queue[child], queue[smallest]) then
+        if child <= count and comparator(self[child], self[smallest]) then
             smallest = child
         end
 
         if parent ~= smallest then
-            queue[parent], queue[smallest] = queue[smallest], queue[parent]
+            self[parent], self[smallest] = self[smallest], self[parent]
             parent = smallest
         else
             break
@@ -48,27 +75,33 @@ local function HeapifyFromStartToEnd(queue, comp)
     end
 end
 
-function PriorityQueue.size(queue)
-    return #queue
+--- Returns the number of the number of elements in the priority queue.
+function PriorityQueue.size(self)
+    return #self
 end
 
-function PriorityQueue.push(queue, element, comp)
-    table.insert(queue, element)
-    HeapifyFromEndToStart(queue, comp)
+-- Inserts an element into the priority queue.
+function PriorityQueue.push(self, element)
+    self[#self + 1] = element
+    heapify_from_end_to_start(self)
 end
 
-function PriorityQueue.pop(queue, comp)
-    local element = queue[1]
+-- Removes and returns the highest priority element from the priority queue.
+-- If the priority queue is empty returns nil.
+function PriorityQueue.pop(self)
+    local element = self[1]
 
-    queue[1] = queue[#queue]
-    queue[#queue] = nil
-    HeapifyFromStartToEnd(queue, comp)
+    self[1] = self[#self]
+    self[#self] = nil
+    heapify_from_start_to_end(self)
 
     return element
 end
 
-function PriorityQueue.peek(queue)
-    return queue[1]
+-- Returns, without removing, the highest priority element from the priority queue.
+-- If the priority queue is empty returns nil.
+function PriorityQueue.peek(self)
+    return self[1]
 end
 
 return PriorityQueue

utils/task.lua

@@ -9,38 +9,65 @@ local PriorityQueue = require 'utils.priority_queue'
 local Event = require 'utils.event'
 local Token = require 'utils.token'
 local ErrorLogging = require 'utils.error_logging'
+local Global = require 'utils.global'
+
+local floor = math.floor
+local log10 = math.log10
+local Token_get = Token.get
+local pcall = pcall
+local Queue_peek = Queue.peek
+local Queue_pop = Queue.pop
+local Queue_push = Queue.push
+local PriorityQueue_peek = PriorityQueue.peek
+local PriorityQueue_pop = PriorityQueue.pop
+local PriorityQueue_push = PriorityQueue.push
 
 local Task = {}
 
-global.callbacks = global.callbacks or PriorityQueue.new()
-global.next_async_callback_time = -1
-global.task_queue = global.task_queue or Queue.new()
-global.total_task_weight = 0
-global.task_queue_speed = 1
-
-local function comp(a, b)
+local function comparator(a, b)
     return a.time < b.time
 end
 
-global.tpt = global.task_queue_speed
-local function get_task_per_tick()
-    if game.tick % 300 == 0 then
-        local size = global.total_task_weight
-        global.tpt = math.floor(math.log10(size + 1)) * global.task_queue_speed
-        if global.tpt < 1 then
-            global.tpt = 1
-        end
+local callbacks = PriorityQueue.new(comparator)
+local task_queue = Queue.new()
+local primitives = {
+    next_async_callback_time = -1,
+    total_task_weight = 0,
+    task_queue_speed = 1,
+    task_per_tick = 1
+}
+
+Global.register(
+    {callbacks = callbacks, task_queue = task_queue, primitives = primitives},
+    function(tbl)
+        callbacks = tbl.callbacks
+        task_queue = tbl.task_queue
+        primitives = tbl.primitives
     end
-    return global.tpt
+)
+
+local function get_task_per_tick(tick)
+    if tick % 300 == 0 then
+        local size = primitives.total_task_weight
+        local task_per_tick = floor(log10(size + 1)) * primitives.task_queue_speed
+        if task_per_tick < 1 then
+            task_per_tick = 1
+        end
+
+        primitives.task_per_tick = task_per_tick
+        return task_per_tick
+    end
+    return primitives.task_per_tick
 end
 
 local function on_tick()
-    local queue = global.task_queue
-    for i = 1, get_task_per_tick() do
-        local task = Queue.peek(queue)
+    local tick = game.tick
+
+    for i = 1, get_task_per_tick(tick) do
+        local task = Queue_peek(task_queue)
         if task ~= nil then
             -- result is error if not success else result is a boolean for if the task should stay in the queue.
-            local success, result = pcall(Token.get(task.func_token), task.params)
+            local success, result = pcall(Token_get(task.func_token), task.params)
             if not success then
                 if _DEBUG then
                     error(result)
@@ -48,19 +75,18 @@ local function on_tick()
                     log(result)
                     ErrorLogging.generate_error_report(result)
                 end
-                Queue.pop(queue)
-                global.total_task_weight = global.total_task_weight - task.weight
+                Queue_pop(task_queue)
+                primitives.total_task_weight = primitives.total_task_weight - task.weight
             elseif not result then
-                Queue.pop(queue)
-                global.total_task_weight = global.total_task_weight - task.weight
+                Queue_pop(task_queue)
+                primitives.total_task_weight = primitives.total_task_weight - task.weight
             end
         end
     end
 
-    local callbacks = global.callbacks
-    local callback = PriorityQueue.peek(callbacks)
-    while callback ~= nil and game.tick >= callback.time do
-        local success, result = pcall(Token.get(callback.func_token), callback.params)
+    local callback = PriorityQueue_peek(callbacks)
+    while callback ~= nil and tick >= callback.time do
+        local success, result = pcall(Token_get(callback.func_token), callback.params)
         if not success then
             if _DEBUG then
                 error(result)
@@ -69,8 +95,8 @@ local function on_tick()
                 ErrorLogging.generate_error_report(result)
             end
         end
-        PriorityQueue.pop(callbacks, comp)
-        callback = PriorityQueue.peek(callbacks)
+        PriorityQueue_pop(callbacks)
+        callback = PriorityQueue_peek(callbacks)
     end
 end
 
@@ -85,7 +111,7 @@ function Task.set_timeout_in_ticks(ticks, func_token, params)
     end
     local time = game.tick + ticks
     local callback = {time = time, func_token = func_token, params = params}
-    PriorityQueue.push(global.callbacks, callback, comp)
+    PriorityQueue_push(callbacks, callback)
 end
 
 --- Allows you to set a timer (in seconds) after which the tokened function will be run with params given as an argument
@@ -109,8 +135,21 @@ end
 -- Ex. if the task is expected to repeat multiple times (ie. the function returns true and loops several ticks)
 function Task.queue_task(func_token, params, weight)
     weight = weight or 1
-    global.total_task_weight = global.total_task_weight + weight
-    Queue.push(global.task_queue, {func_token = func_token, params = params, weight = weight})
+    primitives.total_task_weight = primitives.total_task_weight + weight
+    Queue_push(task_queue, {func_token = func_token, params = params, weight = weight})
+end
+
+function Task.get_queue_speed()
+    return primitives.task_queue_speed
+end
+
+function Task.set_queue_speed(value)
+    value = value or 1
+    if value < 0 then
+        value = 0
+    end
+
+    primitives.task_queue_speed = value
 end
 
 Event.add(defines.events.on_tick, on_tick)