snacks.nvim/lua/snacks/util/spawn.lua

---@class snacks.spawn
local M = {}

local uv = vim.uv or vim.loop

---@class snacks.spawn.Config: uv.spawn.options,{}
---@field cmd string
---@field args? (string|number)[]
---@field timeout? number
---@field run? boolean
---@field debug? boolean
---@field on_stdout? fun(proc: snacks.spawn.Proc, data: string)
---@field on_stderr? fun(proc: snacks.spawn.Proc, data: string)
---@field on_exit? fun(proc: snacks.spawn.Proc, err: boolean)

---@class snacks.spawn.Multi: snacks.spawn.Config,{}
---@field cmd? nil
---@field on_exit? fun(procs: snacks.spawn.Proc[], err: boolean)

---@class snacks.spawn.Proc
---@field opts snacks.spawn.Config
---@field handle? uv.uv_process_t
---@field stdout uv.uv_pipe_t
---@field stderr uv.uv_pipe_t
---@field code? number
---@field signal? number
---@field timer? uv.uv_timer_t
---@field aborted? boolean
---@field data table<uv.uv_pipe_t, string[]>
local Proc = {}
Proc.__index = Proc

---@param handle uv.uv_handle_t?
local function close(handle)
  if handle and not handle:is_closing() then
    handle:close()
  end
end

---@param opts snacks.spawn.Config
function Proc.new(opts)
  local self = setmetatable({}, Proc)
  self.opts = opts
  self.code, self.signal = 0, 0
  self.data = {}
  if opts.run ~= false then
    self:run()
  end
  return self
end

function Proc:running()
  return self.handle and not self.handle:is_closing()
end

---@param signal? string|number
function Proc:kill(signal)
  close(self.stdout)
  close(self.stderr)
  if self:running() then
    self.aborted = true
    self.handle:kill(signal or "sigterm")
  end
end

function Proc:failed()
  if self:running() then
    return false
  end
  return self.code ~= 0 or self.signal ~= 0
end

function Proc:run()
  assert(not self.handle, "already running")
  self.stdout = assert(uv.new_pipe())
  self.stderr = assert(uv.new_pipe())
  self.data = { [self.stdout] = {}, [self.stderr] = {} }
  if self.opts.debug then
    vim.schedule(function()
      Snacks.debug.cmd({ cmd = self.opts.cmd, args = self.opts.args, cwd = self.opts.cwd })
    end)
  end
  local opts = vim.tbl_deep_extend("force", self.opts, {
    stdio = { nil, self.stdout, self.stderr },
    hide = true,
    args = vim.tbl_map(tostring, self.opts.args or {}),
  })
  self.handle = uv.spawn(self.opts.cmd, opts, function(code, signal)
    self.code = code
    self.signal = signal
    self:on_exit()
  end)
  if not self.handle then
    self.code = 1
    self.data[self.stderr] = { "Failed to spawn " .. self.opts.cmd }
    close(self.stdout)
    close(self.stderr)
    return self:on_exit()
  end
  if self.opts.timeout then
    self.timer = assert(uv.new_timer())
    self.timer:start(self.opts.timeout, 0, function()
      self:kill("sigterm")
    end)
  end
  for _, handle in ipairs({ self.stdout, self.stderr }) do
    handle:read_start(function(err, data)
      assert(not err, err)
      if data then
        self:on_data(data, handle)
      else
        close(handle)
      end
    end)
  end
end

function Proc:out()
  return table.concat(self.data[self.stdout] or {})
end

function Proc:err()
  return table.concat(self.data[self.stderr] or {})
end

function Proc:lines()
  return vim.split(self:out(), "\n", { plain = true })
end

---@param data string
---@param handle uv.uv_pipe_t
function Proc:on_data(data, handle)
  table.insert(self.data[handle], data)
  if self.opts.on_stdout and handle == self.stdout then
    self.opts.on_stdout(self, data)
  elseif self.opts.on_stderr and handle == self.stderr then
    self.opts.on_stderr(self, data)
  end
end

function Proc:on_exit()
  close(self.timer)
  close(self.handle)
  local check = assert(uv.new_check())
  check:start(function()
    for _, handle in ipairs({ self.stdout, self.stderr }) do
      if handle and not handle:is_closing() then
        return
      end
    end
    check:stop()
    close(check)
    close(self.stdout)
    close(self.stderr)
    if self.opts.on_exit then
      self.opts.on_exit(self, self.code ~= 0 or self.signal ~= 0)
    end
  end)
end

---@param procs snacks.spawn.Proc[]
---@param opts? snacks.spawn.Multi
function M.multi(procs, opts)
  if #procs == 0 then
    return
  end
  opts = opts or {}
  local current = 0

  local function done()
    if opts.on_exit then
      opts.on_exit(procs, procs[current]:failed())
    end
  end

  local function next()
    current = current + 1
    assert(current <= #procs, "current > #procs")
    local proc = procs[current]
    proc.opts = Snacks.config.merge(vim.deepcopy(opts), proc.opts, {
      on_exit = function(_, err)
        if err or current == #procs then
          done()
        else
          next()
        end
      end,
    })
    proc:run()
  end

  ---@type snacks.spawn.Proc|{procs: snacks.spawn.Proc[]}
  local ret = setmetatable({
    procs = procs,
    run = next,
  }, {
    __index = function(_, k)
      return procs[current][k]
    end,
  })

  if opts.run ~= false then
    next()
  end
  return ret
end

M.new = Proc.new

return M