feat: subprocess parsing (#119)

Massively reduces load on editor during parsing by performing parsing in separate Neovim instance and receiving results over RPC

lua/neotest/client/events/init.lua

@@ -45,7 +45,7 @@ function NeotestEventProcessor:emit(event, ...)
   async.run(function()
     logger.info("Emitting", event, "event")
     for name, listener in pairs(self.listeners[event] or {}) do
-      logger.info("Calling listener", name, "for event", event)
+      logger.debug("Calling listener", name, "for event", event)
       listener(unpack(args))
     end
   end)

lua/neotest/client/init.lua

@@ -331,6 +331,9 @@ function NeotestClient:_start(args)
   if self._started and not args.force then
     return
   end
+  if not lib.subprocess.enabled() then
+    lib.subprocess.init()
+  end
   local process_tracker = NeotestProcessTracker()
   self._runner = NeotestRunner(process_tracker)
   self._state = NeotestState(self._events)

lua/neotest/client/strategies/init.lua

@@ -1,12 +1,11 @@
 local lib = require("neotest.lib")
 local async = require("neotest.async")
 local logger = require("neotest.logging")
-local fu = lib.func_util
 
 ---@return neotest.Strategy
-local get_strategy = fu.memoize(function(name)
+local get_strategy = function(name)
   return require("neotest.client.strategies." .. name)
-end)
+end
 
 ---@class neotest.ProcessTracker
 ---@field _instances table<integer, neotest.Process>

lua/neotest/consumers/summary/init.lua

@@ -64,7 +64,7 @@ local function render(expanded)
 end
 
 async.run(function()
-  local _, err = pcall(function()
+  xpcall(function()
     while true do
       if not pending_render then
         render_cond:wait()
@@ -109,10 +109,9 @@ async.run(function()
       async.api.nvim_exec("redraw", false)
       async.util.sleep(100)
     end
+  end, function(msg)
+    logger.error("Error in summary consumer", debug.traceback(msg, 2))
   end)
-  if err then
-    logger.error("Error in summary consumer", err)
-  end
 end)
 
 local function expand(pos_id, recursive, focus)

lua/neotest/lib/file/find.lua

@@ -1,5 +1,6 @@
 local uv = vim.loop
 local M = {}
+local lib = require("neotest.lib")
 
 local ignored = vim.tbl_add_reverse_lookup({
   "node_modules",
@@ -13,7 +14,7 @@ local ignored = vim.tbl_add_reverse_lookup({
 function M.find(root, opts)
   opts = opts or {}
   local filter_dir = opts.filter_dir
-  local sep = require("neotest.lib").files.sep
+  local sep = lib.files.sep
   local dirs_to_scan = {}
 
   local paths = {}

lua/neotest/lib/file/init.lua

@@ -146,7 +146,14 @@ function M.is_dir(path)
   return M.exists(path .. M.sep)
 end
 
-M.find = require("neotest.lib.file.find").find
+--- Find all files under the given directory.
+--- Does not search hidden directories.
+---@async
+---@param root string
+---@return string[] @Absolute paths of all files within directories to search
+M.find = function(root, opts)
+  return require("neotest.lib.file.find").find(root, opts)
+end
 
 function M.parent(path)
   local elems = vim.split(path, M.sep, { plain = true })

lua/neotest/lib/init.lua

@@ -1,8 +1,10 @@
-local xml = require("neotest.lib.xml")
-local xml_tree = require("neotest.lib.xml.tree")
-
 local M = {}
 
+local lazy_require = require("neotest.lib.require")
+
+local xml = lazy_require("neotest.lib.xml")
+local xml_tree = lazy_require("neotest.lib.xml.tree")
+
 M.xml = {
   ---@param xml_data string
   ---@return table
@@ -14,11 +16,14 @@ M.xml = {
   end,
 }
 
-M.files = require("neotest.lib.file")
+---@module 'neotest.lib.file'
+M.files = lazy_require("neotest.lib.file")
 
-M.func_util = require("neotest.lib.func_util")
+---@module 'neotest.lib.func_util'
+M.func_util = lazy_require("neotest.lib.func_util")
 
-M.treesitter = require("neotest.lib.treesitter")
+---@module 'neotest.lib.treesitter''
+M.treesitter = lazy_require("neotest.lib.treesitter")
 
 M.notify = function(msg, level, opts)
   vim.schedule(function()
@@ -33,12 +38,22 @@ M.notify = function(msg, level, opts)
   end)
 end
 
-M.vim_test = require("neotest.lib.vim_test")
+---@module 'neotest.lib.vim_test''
+M.vim_test = lazy_require("neotest.lib.vim_test")
+
+---@module 'neotest.lib.ui''
+M.ui = lazy_require("neotest.lib.ui")
 
-M.ui = require("neotest.lib.ui")
+---@module 'neotest.lib.positions''
+M.positions = lazy_require("neotest.lib.positions")
 
-M.positions = require("neotest.lib.positions")
+---@module 'neotest.lib.process''
+M.process = lazy_require("neotest.lib.process")
 
-M.process = require("neotest.lib.process")
+---Module to interact with a child Neovim instance.
+---This can be used for CPU intensive work like treesitter parsing.
+---All usage should be guarded by checking that the subprocess has been started using the `enabled` function.
+---@module 'neotest.lib.subprocess''
+M.subprocess = lazy_require("neotest.lib.subprocess")
 
 return M

lua/neotest/lib/require.lua

@@ -0,0 +1,13 @@
+return function(module)
+  return setmetatable({}, {
+    __index = function(_, key)
+      return require(module)[key]
+    end,
+    __newindex = function(_, key, value)
+      require(module)[key] = value
+    end,
+    __call = function(_, ...)
+      return require(module)(...)
+    end,
+  })
+end

lua/neotest/lib/subprocess.lua

@@ -0,0 +1,135 @@
+local async = require("neotest.async")
+local logger = require("neotest.logging")
+
+local child_chan, parent_chan
+local callbacks = {}
+local next_cb_id = 1
+local enabled = false
+
+local M = {}
+
+---Initialize the subprocess module.
+---Do not call this, neotest core will initialize.
+function M.init()
+  logger.info("Starting child process")
+  local parent_address = async.fn.serverstart()
+  local success
+  success, child_chan = pcall(async.fn.jobstart, { vim.loop.exepath(), "--embed", "--headless" }, {
+    rpc = true,
+    on_exit = function()
+      logger.info("Child process exited")
+      enabled = false
+    end,
+  })
+  if not success then
+    logger.error("Failed to start child process", child_chan)
+    return
+  end
+  local mode = async.fn.rpcrequest(child_chan, "nvim_get_mode")
+  if mode.blocking then
+    logger.error("Child process is waiting for input at startup. Aborting.")
+  end
+  xpcall(function()
+    -- Trigger lazy loading of neotest
+    async.fn.rpcrequest(child_chan, "nvim_exec_lua", "return require('neotest') and 0", {})
+    async.fn.rpcrequest(
+      child_chan,
+      "nvim_exec_lua",
+      "return require('neotest.lib').subprocess._set_parent_address(...)",
+      { parent_address }
+    )
+    enabled = true
+  end, function(msg)
+    logger.error("Failed to initialize child process", debug.traceback(msg, 2))
+    child_chan = nil
+  end)
+end
+
+function M._set_parent_address(parent_address)
+  _G._NEOTEST_IS_CHILD = true
+  parent_chan = vim.fn.sockconnect("pipe", parent_address, { rpc = true })
+  logger.info("Connected to parent instance")
+end
+
+function M._register_result(callback_id, res, err)
+  logger.debug("Result registed for callback", callback_id)
+  local cb = callbacks[callback_id]
+  callbacks[callback_id] = nil
+  cb(res, err)
+end
+
+local function get_chan()
+  if M.is_child() then
+    return parent_chan
+  else
+    return child_chan
+  end
+end
+
+---@async
+---Wrapper around vim.fn.rpcrequest that will automatically select the channel for the child or parent process,
+---depending on if the current instance is the child or parent.
+---See `:help rpcrequest` for more information.
+function M.request(method, ...)
+  async.fn.rpcrequest(get_chan(), method, ...)
+end
+
+---@async
+---Wrapper around vim.fn.rpcnotify that will automatically select the channel for the child or parent process,
+---depending on if the current instance is the child or parent.
+---See `:help rpcnotify` for more information.
+function M.notify(method, ...)
+  async.fn.rpcnotify(get_chan(), method, ...)
+end
+
+---@async
+---Call a lua function in the other process with the given argument list, returning the result.
+---The function will be called in async context.
+---@param func string A globally accessible function in the other process. e.g. `"require('neotest.lib').files.read"`
+---@param args? any[] Arguments to pass to the function
+---@return any, string?: Result or error message if call failed
+function M.call(func, args)
+  local send_result, await_result = async.control.channel.oneshot()
+  local cb_id = next_cb_id
+  next_cb_id = next_cb_id + 1
+  callbacks[cb_id] = send_result
+  logger.debug("Waiting for result", cb_id)
+  M.notify(
+    "nvim_exec_lua",
+    "return require('neotest.lib.subprocess')._remote_call(" .. func .. ", ...)",
+    { cb_id, args or {} }
+  )
+  return await_result()
+end
+
+function M._remote_call(func, cb_id, args)
+  logger.debug("Received remote call", cb_id, func)
+  async.run(function()
+    xpcall(function()
+      local res = func(unpack(args))
+      M.notify(
+        "nvim_exec_lua",
+        "return require('neotest.lib.subprocess')._register_result(...)",
+        { cb_id, res }
+      )
+    end, function(msg)
+      local err = debug.traceback(msg, 2)
+      logger.warn("Error in remote call", err)
+      M.notify(
+        "nvim_exec_lua",
+        "return require('neotest.lib.subprocess')._register_result(...)",
+        { cb_id, nil, err }
+      )
+    end)
+  end)
+end
+
+function M.enabled()
+  return enabled
+end
+
+function M.is_child()
+  return parent_chan ~= nil
+end
+
+return M