u.nvim/lua/u/tracker.lua

local M = {}

M.debug = false

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

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

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

  -- We don't handle cyclic updates:
  if self.changing then return end

  local prev_changing = self.changing
  self.changing = true
  for _, cb in ipairs(self.subscribers) do
    pcall(cb, value)
  end
  self.changing = prev_changing
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

--- @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>
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)
  local wrapped_callback = function() callback() end
  for signal in pairs(self.signals) do
    signal:subscribe(wrapped_callback)
  end

  return function()
    for signal in pairs(self.signals) do
      signal:unsubscribe(wrapped_callback)
    end
  end
end

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

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

--- @param fn function
--- @param name? string
--- @return Signal
function M.create_memo(fn, name)
  local result
  M.create_effect(function()
    local value = fn()
    if name and M.debug then vim.notify(name) end
    if result then
      result:set(value)
    else
      result = M.create_signal(value, name and ('m.s:' .. name) or nil)
    end
  end, name)
  return result
end

--- @param fn function
--- @param name? string
function M.create_effect(fn, name)
  local ctx = M.ExecutionContext:new()
  M.ExecutionContext:run(fn, ctx)
  return ctx:subscribe(function()
    if name and M.debug then
      local deps = vim
        .iter(vim.tbl_keys(ctx.signals))
        :map(function(s) return s.name end)
        :filter(function(nm) return nm ~= nil end)
        :join ','
      vim.notify(name .. ':=>' .. deps)
    end
    fn()
  end)
end

return M