task spreader for factorio

PocketWatch.lua

--[[
General Purpose Tick Spreader for Factorio. Completely agnostic about *what* it's doing, 
and should be compatible with any project that can use some distribution of code across ticks.

Usage:
local PocketWatch = require"path.to.PocketWatch" 
-- defines the PocketWatch class. The path is relative, as always with require.

PocketWatch.taskMap[id] = <function>    
-- teach PocketWatch how to perform a task. Since functions with upvalues cannot be serialized currently, we can't simply drop function objects (especially closures) into the global table. 
-- I recommend that you structure your recursive task functions like so (this creates behavior that is somewhat similar to how coroutines work)
PocketWatch.taskMap['someTask'] = function(state, ...)
  -- perform some small manipulation to state
  return watch:Do('someTask', state, ...) -- where watch is a PocketWatch Object. Be sure to return, so that it is a proper tail call.
end

timers = PocketWatch:Init()
-- intended for on_init. Sets up global to accept timers

timers = PocketWatch:Load()
-- intended for on_load. Subscribes to defines.events.on_tick, and restores metatables.

watch = PocketWatch:New(id) 
-- create a new PocketWatch Object. Usually should only be used in on_init.

watch:Do(id, ...) 
-- will attempt to call PocketWatch.taskMap[id](...). If the # of tasks performed in the current tick exceeds the 
-- task limit, then this will Schedule the task for the next tick, and allow the simulation to continue.

watch.isBlocking = <bool> 
-- if isBlocking is true, then PocketWatch will act synchronously and immediately perform all tasks assigned, instead of spreading them across ticks.

watch.taskLimit = <uint> 
-- maximum # of tasks that a single watch can perform, before it relinquishes control and allows the simulation to continue.

watch:Schedule(task, isOld, sleepTime) 
-- schedules a task for later execution (by default on the next tick). 
-- The task object must be an array whose first element (the 1 index) maps to a valid function from PocketWatch.taskMap, 
-- and the extra information in the task object is what is passed to the task function. 
-- The isOld flag causes the task to be placed last in the queue, instead of in the front.

data = watch:Dump() 
-- dumps stats about the watch into a human-readable format. return type is string.
-- Intended to be used for debugging/statistics

watch:DoTasks(time) 
-- Performs tasks scheduled for the specified time. Should not be necessary to explicitly call - on_tick will handle that.

watch.ContinueWork 
PocketWatch.setmetatable(watch) 

-- on_tick handler and metatable repair tool. These should not be touched, save for on_load, where they should be used like so:
--in control.lua:
local watch
script.on_load(function()
  --misc code, not related to PocketWatch
  watch = PocketWatch.setmetatable(global.watch)
  script.on_event(defines.events.on_tick, watch.ContinueWork)
end)
-- **NOTE** PocketWatch:Load() performs this work for you, and it's recommended that you simply use this and not touch PocketWatch.ContinueWork/setmetatable
--]]

local PocketWatch = {}
local watchMt = { __index = PocketWatch }
local inspect = require 'inspect'

PocketWatch.taskMap = {}
PocketWatch.timers = {}
PocketWatch.mt = watchMt

function PocketWatch:Init()
  global.timers = self.timers
  return timers
end

function PocketWatch:Load()
  local isWorking
  self.timers = global.timers
  for _, timer in pairs(global.timers) do
    self.setmetatable(timer)
    isWorking = isWorking or timer.working
  end
  if isWorking then 
    script.on_event(defines.events.on_tick, self.ContinueWork)
  end
  return global.timers
end

function PocketWatch:New(id)
  --creates a new Pocketwatch object.
  self.timers[id] = self.timers[id] or self.setmetatable({
    taskList = {},
    id = id,
    isBlocking = nil,
    taskLimit = 1,
    taskCount = 0,
    globalCount = 0,
    ticksWorked = 0,
    futureTasks = 0,
    emptyTicks = 0,
    now = 0,
    type = "PocketWatch",
  })
  return timer
end

function PocketWatch:Do(id, ...)
  local task = self.taskMap[id]
  if task then
    if self.isBlocking or self.isEarly then
      return task(...)
    elseif self.taskCount < self.taskLimit then
      self.taskCount = self.taskCount + 1
      return task(...)
    else --too much! relinquish control and let the simulation tick
      if not self.working then
        self.working = true
        script.on_event(defines.events.on_tick, self.ContinueWork)
      end
      self:Schedule({id, ...})
    end
  end
end

function PocketWatch:Schedule(task, old, sleepTime)
  local nextTick, taskList = game.tick + (sleepTime or 1)
  if not self.taskList[nextTick] then
    self.taskList[nextTick] = {}
  end
  taskList = self.taskList[nextTick]
  if old then
    table.insert(taskList, 1, task)
  else
    self.futureTasks = self.futureTasks + 1
    table.insert(taskList, task)
  end
  if not self.working then
    script.on_event(defines.events.on_tick, self.ContinueWork)
  end
end

function PocketWatch:DoTasks()
  local now = game.tick
  local tooBusy
  self.globalCount = self.globalCount + self.taskCount
  self.taskCount = 0
  local tasks = self.taskList[now]
  if tasks then
    for _, task in ipairs(tasks) do
      if self.taskCount < self.taskLimit then
        self.futureTasks = self.futureTasks - 1
        self:Do(unpack(task))
      else
        tooBusy = true
        self:Schedule(task, true)
      end
    end
  end
  if self.taskCount > 0 then
    self.ticksWorked = self.ticksWorked + 1
  else
    self.emptyTicks = self.emptyTicks + 1
  end
  self.working = tooBusy
  self.taskList[now] = nil
end

function PocketWatch.ContinueWork(event)
  local futureTasks = 0
  for _, timer in pairs(PocketWatch.timers) do
    timer:DoTasks(game.tick)
    futureTasks = futureTasks + timer.futureTasks
  end
  if futureTasks == 0 then
    script.on_event(defines.events.on_tick, nil)
  end
end

function PocketWatch:Dump()
  local toLog = ([[
PocketWatch Dump at: %d
* ID: %s
* Unfinished tasks:                %d
* Total Tasks completed:           %d
* Ticks Spent working:             %d
* Ticks Spent Registered but idle: %d
]]):format(
    self.now or 0,
    self.id or "unkown",
    self.futureTasks or 0,
    self.globalCount or 0,
    self.ticksWorked or 0,
    self.emptyTicks or 0
  )
  return toLog
end

function PocketWatch.setmetatable(watch)
  setmetatable(watch, watchMt)
  return watch
end

return PocketWatch