u.nvim/lua/u/signal.lua

local M = {}

--- @class Signal
--- @field value any
--- @field subscribers table<function, boolean>
local Signal = {}
M.Signal = Signal
Signal.__index = Signal

--- @param value any
--- @return Signal
function Signal:new(value)
  local obj = setmetatable({
    value = value,
    subscribers = {},
  }, self)
  return obj
end

--- @param value any
function Signal:set(value)
  self.value = value
  self:_notify()
end

--- @return any
function Signal:get()
  local ctx = M.ExecutionContext.current()
  if ctx then ctx:track(self) end
  return self.value
end

--- @param fn function
function Signal:update(fn) self:set(fn(self.value)) end

--- @private
function Signal:_notify()
  for _, cb in ipairs(self.subscribers) do
    cb(self.value)
  end
end

--- @param callback function
function Signal:subscribe(callback)
  table.insert(self.subscribers, callback)
  return function() self:unsubscribe(callback) end
end

--- @param callback function
function Signal:unsubscribe(callback)
  for i, cb in ipairs(self.subscribers) do
    if cb == callback then
      table.remove(self.subscribers, i)
      break
    end
  end
end

function Signal:dispose() self.subscribers = {} end

CURRENT_CONTEXT = nil

--- @class ExecutionContext
--- @field signals table<Signal, boolean>
--- @field subscribers table<function, boolean>
local ExecutionContext = {}
M.ExecutionContext = ExecutionContext
ExecutionContext.__index = ExecutionContext

--- @return ExecutionContext
function ExecutionContext:new()
  return setmetatable({
    signals = {},
    subscribers = {},
  }, ExecutionContext)
end

function ExecutionContext.current() return CURRENT_CONTEXT end

--- @param fn function
--- @param ctx ExecutionContext
function ExecutionContext:run(fn, ctx)
  local oldCtx = CURRENT_CONTEXT
  CURRENT_CONTEXT = ctx
  local result
  local success, err = pcall(function() result = fn() end)

  CURRENT_CONTEXT = oldCtx

  if not success then error(err) end

  return result
end

function ExecutionContext:track(signal) self.signals[signal] = true end

--- @param callback function
function ExecutionContext:subscribe(callback)
  self.subscribers[callback] = true
  for signal in pairs(self.signals) do
    signal:subscribe(function() callback() end)
  end
  return function() self:unsubscribe(callback) end
end

--- @param callback function
function ExecutionContext:unsubscribe(callback) self.subscribers[callback] = nil end

function ExecutionContext:dispose()
  for signal, _ in pairs(self.signals) do
    signal:dispose()
  end
  self.signals = {}
  self.subscribers = {}
end

--- @param value any
--- @return Signal
function M.create_signal(value) return Signal:new(value) end

--- @param fn function
--- @return Signal
function M.create_memo(fn)
  local result
  M.create_effect(function()
    local value = fn()
    if result then
      result:set(value)
    else
      result = M.create_signal(value)
    end
  end)
  return result
end

--- @param fn function
function M.create_effect(fn)
  local ctx = M.ExecutionContext:new()
  M.ExecutionContext:run(fn, ctx)
  return ctx:subscribe(fn)
end

return M