control profiling with NVIM_PROFILE env variable or F1 key

init.lua

@@ -1,9 +1,29 @@
-if vim.env["VIRTUAL_ENV"] ~= nil then
-	vim.g.python3_host_prog = vim.fn.system("which -a python3 | sed -n 2p | tr -d '\n'")
-else
-	vim.g.python3_host_prog = vim.fn.system("which python3 | tr -d '\n'")
+local should_profile = os.getenv("NVIM_PROFILE")
+if should_profile then
+  require("profile").instrument_autocmds()
+  if should_profile:lower():match("^start") then
+    require("profile").start("*")
+  else
+    require("profile").instrument("*")
+  end
 end
 
+local function toggle_profile()
+  local prof = require("profile")
+  if prof.is_recording() then
+    prof.stop()
+    vim.ui.input({ prompt = "Save profile to:", completion = "file", default = "profile.json" }, function(filename)
+      if filename then
+        prof.export(filename)
+        vim.notify(string.format("Wrote %s", filename))
+      end
+    end)
+  else
+    prof.start("*")
+  end
+end
+vim.keymap.set("", "<f1>", toggle_profile)
+
 -- lazy.nvim bootstrap
 local lazypath = vim.fn.stdpath("data") .. "/lazy/lazy.nvim"
 if not vim.loop.fs_stat(lazypath) then

lua/profile.lua

@@ -0,0 +1,143 @@
+local autocmd = require("profile.autocmd")
+local clock = require("profile.clock")
+local instrument = require("profile.instrument")
+local util = require("profile.util")
+local M = {}
+
+local event_defaults = {
+  pid = 1,
+  tid = 1,
+}
+
+---Call this at the top of your init.vim to get durations for autocmds. If you
+---don't, autocmds will show up as 'instant' events in the profile
+M.instrument_autocmds = function()
+  autocmd.instrument_start()
+end
+
+---Instrument matching modules
+---@param name string Name of module or glob pattern (e.g. "telescope*")
+M.instrument = function(name)
+  instrument(name)
+end
+
+---Mark matching modules to be ignored by profiling
+---@param name string Name of module or glob pattern (e.g. "telescope*")
+M.ignore = function(name)
+  instrument.ignore(name)
+end
+
+---@param sample_rate number Float between 0 and 1
+M.set_sample_rate = function(sample_rate)
+  instrument.set_sample_rate(sample_rate)
+end
+
+---Start collecting data for a profile
+---@param ... string Names or patterns of modules to instrument (if instrument() not called before this)
+M.start = function(...)
+  for _, pattern in ipairs({ ... }) do
+    instrument(pattern)
+  end
+  autocmd.instrument_auto()
+  instrument.clear_events()
+  clock.reset()
+  instrument.recording = true
+  vim.api.nvim_exec_autocmds("User", { pattern = "ProfileStart", modeline = false })
+end
+
+---@return boolean
+M.is_recording = function()
+  return instrument.recording
+end
+
+---@param filename? string If present, write the profile data to this file
+M.stop = function(filename)
+  instrument.recording = false
+  vim.api.nvim_exec_autocmds("User", { pattern = "ProfileStop", modeline = false })
+  if filename then
+    M.export(filename)
+  end
+end
+
+---@private
+---@param name string Name of the function
+---@param ... any Arguments to the function
+M.log_start = function(name, ...)
+  if not instrument.recording then
+    return
+  end
+  instrument.add_event({
+    name = name,
+    args = util.format_args(...),
+    cat = "function,manual",
+    ph = "B",
+    ts = clock(),
+  })
+end
+
+---@private
+---@param name string Name of the function
+---@param ... any Arguments to the function
+M.log_end = function(name, ...)
+  if not instrument.recording then
+    return
+  end
+  instrument.add_event({
+    name = name,
+    args = util.format_args(...),
+    cat = "function,manual",
+    ph = "E",
+    ts = clock(),
+  })
+end
+
+---@private
+---@param name string Name of the function
+---@param ... any Arguments to the function
+M.log_instant = function(name, ...)
+  if not instrument.recording then
+    return
+  end
+  instrument.add_event({
+    name = name,
+    args = util.format_args(...),
+    cat = "",
+    ph = "i",
+    ts = clock(),
+    s = "g",
+  })
+end
+
+---Print out a list of all modules that have been instrumented
+M.print_instrumented_modules = function()
+  instrument.print_modules()
+end
+
+---Write the trace to a file
+---@param filename string
+M.export = function(filename)
+  local file = assert(io.open(filename, "w"))
+  local events = instrument.get_events()
+  file:write("[")
+  for i, event in ipairs(events) do
+    local e = vim.tbl_extend("keep", event, event_defaults)
+    local ok, jse = pcall(vim.json.encode, e)
+    if not ok and e.args then
+      e.args = nil
+      ok, jse = pcall(vim.json.encode, e)
+    end
+    if ok then
+      file:write(jse)
+      if i < #events then
+        file:write(",\n")
+      end
+    else
+      local err = string.format("Could not encode event: %s\n%s", jse, vim.inspect(e))
+      vim.api.nvim_echo({ { err, "Error" } }, true, {})
+    end
+  end
+  file:write("]")
+  file:close()
+end
+
+return M

lua/profile/autocmd.lua

@@ -0,0 +1,154 @@
+local M = {}
+
+local autocmds = {
+  "BufAdd",
+  "BufDelete",
+  "BufEnter",
+  "BufFilePost",
+  "BufFilePre",
+  "BufHidden",
+  "BufLeave",
+  "BufModifiedSet",
+  "BufNew",
+  "BufNewFile",
+  "BufRead",
+  -- "BufReadCmd",
+  "BufReadPre",
+  "BufUnload",
+  "BufWinEnter",
+  "BufWinLeave",
+  "BufWipeout",
+  "BufWrite",
+  -- "BufWriteCmd",
+  "BufWritePost",
+  "ChanInfo",
+  "ChanOpen",
+  "CmdUndefined",
+  "CmdlineChanged",
+  "CmdlineEnter",
+  "CmdlineLeave",
+  "CmdwinEnter",
+  "CmdwinLeave",
+  "ColorScheme",
+  "ColorSchemePre",
+  "CompleteChanged",
+  "CompleteDonePre",
+  "CompleteDone",
+  "CursorHold",
+  "CursorHoldI",
+  "CursorMoved",
+  "CursorMovedI",
+  "DiffUpdated",
+  "DirChanged",
+  -- "FileAppendCmd",
+  "FileAppendPost",
+  "FileAppendPre",
+  "FileChangedRO",
+  "ExitPre",
+  "FileChangedShell",
+  "FileChangedShellPost",
+  -- "FileReadCmd",
+  "FileReadPost",
+  "FileReadPre",
+  "FileType",
+  -- "FileWriteCmd",
+  "FileWritePost",
+  "FileWritePre",
+  "FilterReadPost",
+  "FilterReadPre",
+  "FilterWritePost",
+  "FilterWritePre",
+  "FocusGained",
+  "FocusLost",
+  "FuncUndefined",
+  "UIEnter",
+  "UILeave",
+  "InsertChange",
+  "InsertCharPre",
+  "TextYankPost",
+  "InsertEnter",
+  "InsertLeavePre",
+  "InsertLeave",
+  "LspAttach",
+  "LspDetach",
+  "LspTokenUpdate",
+  "MenuPopup",
+  "OptionSet",
+  "QuickFixCmdPre",
+  "QuickFixCmdPost",
+  "QuitPre",
+  "RemoteReply",
+  "SessionLoadPost",
+  "ShellCmdPost",
+  "Signal",
+  "ShellFilterPost",
+  "SourcePre",
+  "SourcePost",
+  -- "SourceCmd",
+  "SpellFileMissing",
+  "StdinReadPost",
+  "StdinReadPre",
+  "SwapExists",
+  "Syntax",
+  "TabEnter",
+  "TabLeave",
+  "TabNew",
+  "TabNewEntered",
+  "TabClosed",
+  "TermOpen",
+  "TermEnter",
+  "TermLeave",
+  "TermClose",
+  "TermResponse",
+  "TextChanged",
+  "TextChangedI",
+  "TextChangedP",
+  "User",
+  "VimEnter",
+  "VimLeave",
+  "VimLeavePre",
+  "VimResized",
+  "VimResume",
+  "VimSuspend",
+  "WinClosed",
+  "WinEnter",
+  "WinLeave",
+  "WinNew",
+  "WinScrolled",
+}
+
+---@param groupname string
+---@param fn string
+local function create(groupname, fn)
+  local aug = vim.api.nvim_create_augroup(groupname, {})
+  for _, autocmd in ipairs(autocmds) do
+    vim.api.nvim_create_autocmd(autocmd, {
+      desc = "profile.nvim " .. fn,
+      pattern = "*",
+      group = aug,
+      callback = function(args)
+        require("profile")[fn](autocmd, { match = args.match })
+      end,
+    })
+  end
+end
+
+M.instrument_start = function()
+  create("lua_profile_start", "log_start")
+end
+
+M.instrument_auto = function()
+  if vim.fn.exists("#lua_profile_start") ~= 0 then
+    create("lua_profile_end", "log_end")
+  else
+    create("lua_profile", "log_instant")
+  end
+end
+
+M.clear = function()
+  vim.api.nvim_create_augroup("lua_profile", {})
+  vim.api.nvim_create_augroup("lua_profile_start", {})
+  vim.api.nvim_create_augroup("lua_profile_end", {})
+end
+
+return M

lua/profile/clock.lua

@@ -0,0 +1,14 @@
+local hrtime = vim.loop.hrtime
+
+local start = hrtime()
+
+return setmetatable({
+  reset = function()
+    start = hrtime()
+  end,
+}, {
+  __call = function()
+    -- Microseconds
+    return (hrtime() - start) / 1e3
+  end,
+})

lua/profile/instrument.lua

@@ -0,0 +1,209 @@
+local clock = require("profile.clock")
+local util = require("profile.util")
+local M = {}
+
+local rawrequire = require
+local events = {}
+local ignore_list = {
+  "^_G$",
+  "^bit$",
+  "^coroutine$",
+  "^debug$",
+  "^ffi$",
+  "^io$",
+  "^jit.*$",
+  "^luv$",
+  "^math$",
+  "^os$",
+  "^package$",
+  "^string$",
+  "^table$",
+  "^vim%.inspect$",
+  "^profile.*$",
+  "^lspconfig%.util%.script_path$",
+  "^plenary%.async_lib.*$",
+}
+
+local instrument_list = {}
+local wrapped_modules = {}
+local wrapped_functions = {}
+M.recording = false
+M.sample_rate = 1
+local exposed_globals = {
+  ["vim"] = vim,
+  ["vim.fn"] = vim.fn,
+  ["vim.api"] = vim.api,
+}
+
+local function all_modules()
+  return vim.tbl_extend("keep", package.loaded, exposed_globals)
+end
+
+local function should_instrument(name)
+  -- Don't double-wrap
+  if wrapped_functions[name] then
+    return false
+  elseif wrapped_modules[name] then
+    return false
+  else
+    for _, pattern in ipairs(ignore_list) do
+      if string.match(name, pattern) then
+        return false
+      end
+    end
+    return true
+  end
+end
+
+local function wrap_function(name, fn)
+  return function(...)
+    if M.sample_rate < 1 and math.random() > M.sample_rate then
+      return fn(...)
+    end
+    local arg_string = util.format_args(...)
+    local start = clock()
+    local function handle_result(...)
+      local delta = clock() - start
+      M.add_event({
+        name = name,
+        args = arg_string,
+        cat = "function",
+        ph = "X",
+        ts = start,
+        dur = delta,
+      })
+      return ...
+    end
+    return handle_result(fn(...))
+  end
+end
+
+local function patch_function(modname, mod, name)
+  local fn = mod[name]
+  if type(fn) == "function" then
+    local fnname = string.format("%s.%s", modname, name)
+    if should_instrument(fnname) then
+      wrapped_functions[fnname] = fn
+      mod[name] = wrap_function(fnname, fn)
+    end
+  end
+end
+
+local function maybe_wrap_function(name)
+  local parent_mod_name, fn = util.split_path(name)
+  local mod = M.get_module(parent_mod_name)
+  if mod then
+    patch_function(parent_mod_name, mod, fn)
+  end
+end
+
+local function wrap_module(name, mod)
+  name = util.normalize_module_name(name)
+  if type(mod) ~= "table" or not should_instrument(name) then
+    return
+  end
+  wrapped_modules[name] = true
+  for k in pairs(mod) do
+    -- Do not wrap module functions that start with '_'
+    -- Those have to be explicitly passed to instrument()
+    if type(k) == "string" and not util.startswith(k, "_") then
+      patch_function(name, mod, k)
+    end
+  end
+end
+
+local function should_profile_module(name)
+  for _, pattern in ipairs(instrument_list) do
+    if string.match(name, pattern) then
+      return true
+    end
+  end
+  return false
+end
+
+M.hook_require = function(module_name)
+  if rawrequire ~= require then
+    return
+  end
+  local wrapped_require = wrap_function("require", rawrequire)
+  _G.require = function(name)
+    -- Don't time the require if the file is already loaded
+    if package.loaded[name] or not should_profile_module(name) then
+      return rawrequire(name)
+    end
+    local mod = wrapped_require(name)
+    wrap_module(name, mod)
+    return mod
+  end
+end
+
+M.clear_events = function()
+  events = {}
+end
+
+M.add_event = function(event)
+  if M.recording then
+    table.insert(events, event)
+  end
+end
+
+M.get_module = function(name)
+  local ok, mod = pcall(require, name)
+  if ok then
+    if type(mod) == "table" then
+      return mod
+    else
+      return nil
+    end
+  else
+    mod = _G
+    local paths = vim.split(name, ".", { plain = true })
+    for _, token in ipairs(paths) do
+      mod = mod[token]
+      if not mod then
+        break
+      end
+    end
+    return type(mod) == "table" and mod or nil
+  end
+end
+
+M.get_events = function()
+  return events
+end
+
+M.ignore = function(name)
+  table.insert(ignore_list, util.path_glob_to_regex(name))
+end
+
+M.print_modules = function()
+  for module, _ in pairs(wrapped_modules) do
+    print(module)
+  end
+end
+
+local function instrument(_, name)
+  local pattern = util.path_glob_to_regex(name)
+  if not vim.tbl_contains(instrument_list, pattern) then
+    table.insert(instrument_list, pattern)
+  end
+  M.hook_require(name)
+  for modname, mod in pairs(all_modules()) do
+    if string.match(modname, pattern) then
+      wrap_module(modname, mod)
+    end
+  end
+  maybe_wrap_function(name)
+end
+
+---@param sample_rate number Float between 0 and 1
+M.set_sample_rate = function(sample_rate)
+  if sample_rate <= 0 or sample_rate > 1 then
+    error("sample_rate must be between 0 (exclusive) and 1 (inclusive)")
+  end
+  M.sample_rate = sample_rate
+end
+
+return setmetatable(M, {
+  __call = instrument,
+})

lua/profile/util.lua

@@ -0,0 +1,83 @@
+local M = {}
+
+local MAX_ARG_LEN = 200
+local tbl_isarray = vim.isarray or vim.tbl_isarray or vim.tbl_islist
+local pack_len = vim.F.pack_len
+local split = vim.split
+
+---@param glob string
+---@return string
+M.path_glob_to_regex = function(glob)
+  local pattern = string.gsub(glob, "%.", "[%./]")
+  pattern = string.gsub(pattern, "*", ".*")
+  return "^" .. pattern .. "$"
+end
+
+---@param name string
+---@return string
+M.normalize_module_name = function(name)
+  local ret = string.gsub(name, "/", ".")
+  return ret
+end
+
+---@param haystack string
+---@param prefix string
+---@return boolean
+M.startswith = function(haystack, prefix)
+  return string.find(haystack, prefix) == 1
+end
+
+---@param path string
+---@return string module
+---@return string tail
+M.split_path = function(path)
+  local pieces = split(path, ".", { plain = true })
+  if #pieces == 1 then
+    return "_G", path
+  end
+  local mod = table.concat(pack_len(unpack(pieces, 1, #pieces - 1)), ".")
+  return mod, pieces[#pieces]
+end
+
+local function sanitize(table)
+  local clean = {}
+  local iterfn
+  if tbl_isarray(table) then
+    iterfn = ipairs
+  else
+    iterfn = pairs
+  end
+  for i, v in iterfn(table) do
+    local t = type(v)
+    if t == "string" then
+      if string.len(v) > MAX_ARG_LEN then
+        clean[tostring(i)] = string.sub(v, 1, MAX_ARG_LEN - 3) .. "..."
+      else
+        clean[tostring(i)] = v
+      end
+    elseif t == "nil" or t == "boolean" or t == "number" then
+      clean[tostring(i)] = v
+    end
+  end
+  -- If no args, then return nil
+  if next(clean) == nil then
+    return nil
+  else
+    return clean
+  end
+end
+
+---@param ... any[]
+---@return any
+M.format_args = function(...)
+  local args = pack_len(...)
+  if args.n == 0 then
+    return nil
+  elseif args.n == 1 and type(args[1]) == "table" then
+    return sanitize(args[1])
+  else
+    return sanitize(args)
+  end
+end
+
+return M