From 17a3c3acea400679027e675cc19b738e842a5ea0 Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Tue, 14 Feb 2023 11:00:56 +0100
Subject: [PATCH] perf: more cache optims

---
 .neoconf.json            |   2 +-
 lua/lazy/core/cache.lua  | 241 +++++++++++++++++++++++++++------------
 lua/lazy/core/loader.lua |  41 ++++---
 lua/lazy/init.lua        |  13 ++-
 lua/lazy/view/render.lua |  25 ++--
 5 files changed, 224 insertions(+), 98 deletions(-)

diff --git a/.neoconf.json b/.neoconf.json
index 67828bb..f0c23f8 100644
--- a/.neoconf.json
+++ b/.neoconf.json
@@ -7,7 +7,7 @@
     }
   },
   "lspconfig": {
-    "sumneko_lua": {
+    "lua_ls": {
       "Lua.runtime.version": "LuaJIT",
       "Lua.workspace.checkThirdParty": false
     }
diff --git a/lua/lazy/core/cache.lua b/lua/lazy/core/cache.lua
index adf1bfc..3b3d824 100644
--- a/lua/lazy/core/cache.lua
+++ b/lua/lazy/core/cache.lua
@@ -4,33 +4,44 @@ local uv = vim.loop
 local M = {}
 
 ---@alias CacheHash {mtime: {sec:number, nsec:number}, size:number}
----@alias CacheEntry {hash:CacheHash, modpath:string, chunk:string}
+---@alias CacheEntry {hash:CacheHash, chunk:string}
 
 ---@class CacheFindOpts
 ---@field rtp? boolean Search for modname in the runtime path (defaults to `true`)
 ---@field patterns? string[] Paterns to use (defaults to `{"/init.lua", ".lua"}`)
 ---@field paths? string[] Extra paths to search for modname
 
-M.VERSION = 1
+M.VERSION = 2
 M.path = vim.fn.stdpath("cache") .. "/lazy/luac"
 M.enabled = false
-M.stats = { find = { total = 0, time = 0, index = 0, stat = 0, not_found = 0 } }
+M.stats = {
+  find = { total = 0, time = 0, not_found = 0 },
+}
 
 ---@class ModuleCache
 ---@field _rtp string[]
+---@field _rtp_pure string[]
 ---@field _rtp_key string
 local Cache = {
   ---@type table<string, table<string,true>>
+  _indexed = {},
+  ---@type table<string, string[]>
   _topmods = {},
   _loadfile = loadfile,
 }
 
+function M.track(stat, start)
+  M.stats[stat] = M.stats[stat] or { total = 0, time = 0 }
+  M.stats[stat].total = M.stats[stat].total + 1
+  M.stats[stat].time = M.stats[stat].time + uv.hrtime() - start
+end
+
 -- slightly faster/different version than vim.fs.normalize
 -- we also need to have it here, since the cache will load vim.fs
 ---@private
 function Cache.normalize(path)
   if path:sub(1, 1) == "~" then
-    local home = vim.loop.os_homedir()
+    local home = vim.loop.os_homedir() or "~"
     if home:sub(-1) == "\\" or home:sub(-1) == "/" then
       home = home:sub(1, -2)
     end
@@ -42,28 +53,34 @@ end
 
 ---@private
 function Cache.get_rtp()
+  local start = uv.hrtime()
   if vim.in_fast_event() then
-    return Cache._rtp or {}
+    M.track("get_rtp", start)
+    return (Cache._rtp or {}), false
   end
+  local updated = false
   local key = vim.go.rtp
   if key ~= Cache._rtp_key then
     Cache._rtp = {}
     for _, path in ipairs(vim.api.nvim_get_runtime_file("", true)) do
       path = Cache.normalize(path)
       -- skip after directories
-      if path:sub(-6, -1) ~= "/after" then
+      if path:sub(-6, -1) ~= "/after" and not (Cache._indexed[path] and vim.tbl_isempty(Cache._indexed[path])) then
         Cache._rtp[#Cache._rtp + 1] = path
       end
     end
+    updated = true
     Cache._rtp_key = key
   end
-  return Cache._rtp
+  M.track("get_rtp", start)
+  return Cache._rtp, updated
 end
 
 ---@param name string can be a module name, or a file name
 ---@private
 function Cache.cache_file(name)
-  return M.path .. "/" .. name:gsub("[/\\:]", "%%") .. "c"
+  local ret = M.path .. "/" .. name:gsub("[/\\:]", "%%")
+  return ret:sub(-4) == ".lua" and (ret .. "c") or (ret .. ".luac")
 end
 
 ---@param entry CacheEntry
@@ -76,10 +93,8 @@ function Cache.write(name, entry)
     entry.hash.size,
     entry.hash.mtime.sec,
     entry.hash.mtime.nsec,
-    #entry.modpath,
   }
-  uv.fs_write(f, ffi.string(ffi.new("const uint32_t[5]", header), 20))
-  uv.fs_write(f, entry.modpath)
+  uv.fs_write(f, ffi.string(ffi.new("const uint32_t[4]", header), 16))
   uv.fs_write(f, entry.chunk)
   uv.fs_close(f)
 end
@@ -87,6 +102,7 @@ end
 ---@return CacheEntry?
 ---@private
 function Cache.read(name)
+  local start = uv.hrtime()
   local cname = Cache.cache_file(name)
   local f = uv.fs_open(cname, "r", 438)
   if f then
@@ -95,29 +111,52 @@ function Cache.read(name)
     uv.fs_close(f)
 
     ---@type integer[]|{[0]:integer}
-    local header = ffi.cast("uint32_t*", ffi.new("const char[20]", data:sub(1, 20)))
+    local header = ffi.cast("uint32_t*", ffi.new("const char[16]", data:sub(1, 16)))
     if header[0] ~= M.VERSION then
       return
     end
-    local modpath = data:sub(21, 20 + header[4])
+    M.track("read", start)
     return {
       hash = { size = header[1], mtime = { sec = header[2], nsec = header[3] } },
-      chunk = data:sub(20 + header[4] + 1),
-      modpath = modpath,
+      chunk = data:sub(16 + 1),
     }
   end
+  M.track("read", start)
 end
 
 ---@param modname string
 ---@private
 function Cache.loader(modname)
-  modname = modname:gsub("/", ".")
+  local start = uv.hrtime()
   local modpath, hash = Cache.find(modname)
-  local modpath, hash
-  modpath, hash = Cache.find(modname)
+  ---@type function?, string?
+  local chunk, err
   if modpath then
-    return Cache.load(modname, modpath, { hash = hash })
+    chunk, err = M.load(modpath, { hash = hash })
   end
+  M.track("loader", start)
+  return chunk or err or "module " .. modname .. " not found"
+end
+
+---@param modname string
+---@private
+function Cache.loader_lib(modname)
+  local start = uv.hrtime()
+  local modpath = Cache.find(modname, { patterns = jit.os:find("Windows") and { ".dll" } or { ".so" } })
+  ---@type function?, string?
+  if modpath then
+    -- Making function name in Lua 5.1 (see src/loadlib.c:mkfuncname) is
+    -- a) strip prefix up to and including the first dash, if any
+    -- b) replace all dots by underscores
+    -- c) prepend "luaopen_"
+    -- So "foo-bar.baz" should result in "luaopen_bar_baz"
+    local dash = modname:find("-", 1, true)
+    local funcname = dash and modname:sub(dash + 1) or modname
+    local chunk, err = package.loadlib(modpath, "luaopen_" .. funcname:gsub("%.", "_"))
+    M.track("loader_lib", start)
+    return chunk or err
+  end
+  M.track("loader_lib", start)
   return "module " .. modname .. " not found"
 end
 
@@ -127,8 +166,11 @@ end
 ---@return function?, string?  error_message
 ---@private
 function Cache.loadfile(filename, mode, env)
+  local start = uv.hrtime()
   filename = Cache.normalize(filename)
-  return Cache.load(filename, filename, { mode = mode, env = env })
+  local chunk, err = M.load(filename, { mode = mode, env = env })
+  M.track("loadfile", start)
+  return chunk, err
 end
 
 ---@param h1 CacheHash
@@ -139,44 +181,52 @@ function Cache.eq(h1, h2)
 end
 
 ---@param modpath string
----@param opts? {hash?: CacheHash, mode?: "b"|"t"|"bt", env?:table, entry?: CacheEntry}
+---@param opts? {hash?: CacheHash, mode?: "b"|"t"|"bt", env?:table}
 ---@return function?, string? error_message
 ---@private
-function Cache.load(modkey, modpath, opts)
+function M.load(modpath, opts)
+  local start = uv.hrtime()
+
   opts = opts or {}
   local hash = opts.hash or uv.fs_stat(modpath)
-  if not hash then
-    -- trigger correct error
-    return Cache._loadfile(modpath)
-  end
-
   ---@type function?, string?
   local chunk, err
-  local entry = opts.entry or Cache.read(modkey)
+
+  if not hash then
+    -- trigger correct error
+    chunk, err = Cache._loadfile(modpath, opts.mode, opts.env)
+    M.track("load", start)
+    return chunk, err
+  end
+
+  local entry = Cache.read(modpath)
   if entry and Cache.eq(entry.hash, hash) then
     -- found in cache and up to date
-    chunk, err = load(entry.chunk --[[@as string]], "@" .. entry.modpath)
+    -- selene: allow(incorrect_standard_library_use)
+    chunk, err = load(entry.chunk --[[@as string]], "@" .. modpath, opts.mode, opts.env)
     if not (err and err:find("cannot load incompatible bytecode", 1, true)) then
+      M.track("load", start)
       return chunk, err
     end
   end
   entry = { hash = hash, modpath = modpath }
 
-  chunk, err = Cache._loadfile(entry.modpath)
+  chunk, err = Cache._loadfile(modpath, opts.mode, opts.env)
   if chunk then
     entry.chunk = string.dump(chunk)
-    Cache.write(modkey, entry)
+    Cache.write(modpath, entry)
   end
+  M.track("load", start)
   return chunk, err
 end
 
 ---@param modname string
 ---@param opts? CacheFindOpts
----@return string? modpath, CacheHash? hash
+---@return string? modpath, CacheHash? hash, CacheEntry? entry
 function Cache.find(modname, opts)
-  opts = opts or {}
   local start = uv.hrtime()
-  M.stats.find.total = M.stats.find.total + 1
+  opts = opts or {}
+
   modname = modname:gsub("/", ".")
   local basename = modname:gsub("%.", "/")
   local idx = modname:find(".", 1, true)
@@ -184,40 +234,54 @@ function Cache.find(modname, opts)
 
   -- OPTIM: search for a directory first when topmod == modname
   local patterns = opts.patterns or (topmod == modname and { "/init.lua", ".lua" } or { ".lua", "/init.lua" })
-  local rtp = opts.rtp ~= false and Cache.get_rtp() or {}
-  if opts.paths then
-    rtp = vim.deepcopy(rtp)
-    for _, dir in ipairs(opts.paths) do
-      rtp[#rtp + 1] = Cache.normalize(dir)
-    end
-  end
   for p, pattern in ipairs(patterns) do
     patterns[p] = "/lua/" .. basename .. pattern
   end
 
-  for _, path in ipairs(rtp) do
-    if M.lsmod(path)[topmod] then
-      for _, pattern in ipairs(patterns) do
-        local modpath = path .. pattern
-        M.stats.find.stat = M.stats.find.stat + 1
-        local hash = uv.fs_stat(modpath)
-        if hash then
-          M.stats.find.time = M.stats.find.time + uv.hrtime() - start
-          return modpath, hash
+  ---@param paths string[]
+  local function _find(paths)
+    for _, path in ipairs(paths) do
+      if M.lsmod(path)[topmod] then
+        for _, pattern in ipairs(patterns) do
+          local modpath = path .. pattern
+          M.stats.find.stat = (M.stats.find.stat or 0) + 1
+          local hash = uv.fs_stat(modpath)
+          if hash then
+            return modpath, hash
+          end
         end
       end
     end
   end
 
+  ---@type string, CacheHash
+  local modpath, hash
+
+  if opts.rtp ~= false then
+    modpath, hash = _find(Cache._rtp or {})
+    if not modpath then
+      local rtp, updated = Cache.get_rtp()
+      if updated then
+        modpath, hash = _find(rtp)
+      end
+    end
+  end
+  if (not modpath) and opts.paths then
+    modpath, hash = _find(opts.paths)
+  end
+
+  M.track("find", start)
+  if modpath then
+    return modpath, hash
+  end
   -- module not found
   M.stats.find.not_found = M.stats.find.not_found + 1
-  M.stats.find.time = M.stats.find.time + uv.hrtime() - start
 end
 
 --- Resets the topmods cache for the path
 ---@param path string
 function M.reset(path)
-  Cache._topmods[Cache.normalize(path)] = nil
+  Cache._indexed[Cache.normalize(path)] = nil
 end
 
 function M.enable()
@@ -228,7 +292,17 @@ function M.enable()
   vim.fn.mkdir(vim.fn.fnamemodify(M.path, ":p"), "p")
   -- selene: allow(global_usage)
   _G.loadfile = Cache.loadfile
+  -- add lua loader
   table.insert(package.loaders, 2, Cache.loader)
+  -- add libs loader
+  table.insert(package.loaders, 3, Cache.loader_lib)
+  -- remove Neovim loader
+  for l, loader in ipairs(package.loaders) do
+    if loader == vim._load_package then
+      table.remove(package.loaders, l)
+      break
+    end
+  end
 end
 
 function M.disable()
@@ -240,18 +314,19 @@ function M.disable()
   _G.loadfile = Cache._loadfile
   ---@diagnostic disable-next-line: no-unknown
   for l, loader in ipairs(package.loaders) do
-    if loader == Cache.loader then
+    if loader == Cache.loader or loader == Cache.loader_lib then
       table.remove(package.loaders, l)
     end
   end
+  table.insert(package.loaders, 2, vim._load_package)
 end
 
 -- Return the top-level `/lua/*` modules for this path
 ---@return string[]
 function M.lsmod(path)
-  if not Cache._topmods[path] then
-    M.stats.find.index = M.stats.find.index + 1
-    Cache._topmods[path] = {}
+  if not Cache._indexed[path] then
+    local start = uv.hrtime()
+    Cache._indexed[path] = {}
     local handle = vim.loop.fs_scandir(path .. "/lua")
     while handle do
       local name, t = vim.loop.fs_scandir_next(handle)
@@ -259,7 +334,7 @@ function M.lsmod(path)
         break
       end
       -- HACK: type is not always returned due to a bug in luv
-      t = t or vim.loop.fs_stat(path .. "/" .. name).type
+      t = t or uv.fs_stat(path .. "/" .. name).type
       ---@type string
       local topname
       if name:sub(-4) == ".lua" then
@@ -268,11 +343,16 @@ function M.lsmod(path)
         topname = name
       end
       if topname then
-        Cache._topmods[path][topname] = true
+        Cache._indexed[path][topname] = true
+        Cache._topmods[topname] = Cache._topmods[topname] or {}
+        if not vim.tbl_contains(Cache._topmods[topname], path) then
+          table.insert(Cache._topmods[topname], path)
+        end
       end
     end
+    M.track("lsmod", start)
   end
-  return Cache._topmods[path]
+  return Cache._indexed[path]
 end
 
 ---@param modname string
@@ -283,22 +363,43 @@ function M.find(modname, opts)
   return modpath
 end
 
+function M.profile_loaders()
+  for l, loader in pairs(package.loaders) do
+    local loc = debug.getinfo(loader, "Sn").source:sub(2)
+    package.loaders[l] = function(modname)
+      local start = vim.loop.hrtime()
+      local ret = loader(modname)
+      M.track("loader " .. l .. ": " .. loc, start)
+      M.track("loader_all", start)
+      return ret
+    end
+  end
+end
+
 function M.inspect()
   local function ms(nsec)
     return math.floor(nsec / 1e6 * 1000 + 0.5) / 1000 .. "ms"
   end
-  local props = {
-    { "total", M.stats.find.total, "Number" },
-    { "time", ms(M.stats.find.time), "Bold" },
-    { "avg time", ms(M.stats.find.time / M.stats.find.total), "Bold" },
-    { "index", M.stats.find.index, "Number" },
-    { "fs_stat", M.stats.find.stat, "Number" },
-    { "not found", M.stats.find.not_found, "Number" },
-  }
   local chunks = {} ---@type string[][]
-  for _, prop in ipairs(props) do
-    chunks[#chunks + 1] = { "* " .. prop[1] .. ": " }
-    chunks[#chunks + 1] = { tostring(prop[2]) .. "\n", prop[3] }
+  ---@type string[]
+  local stats = vim.tbl_keys(M.stats)
+  table.sort(stats)
+  for _, stat in ipairs(stats) do
+    vim.list_extend(chunks, {
+      { "\n" .. stat .. "\n", "Title" },
+      { "* total:    " },
+      { tostring(M.stats[stat].total) .. "\n", "Number" },
+      { "* time:     " },
+      { ms(M.stats[stat].time) .. "\n", "Bold" },
+      { "* avg time: " },
+      { ms(M.stats[stat].time / M.stats[stat].total) .. "\n", "Bold" },
+    })
+    for k, v in pairs(M.stats[stat]) do
+      if not vim.tbl_contains({ "time", "total" }, k) then
+        chunks[#chunks + 1] = { "* " .. k .. ":" .. string.rep(" ", 9 - #k) }
+        chunks[#chunks + 1] = { tostring(v) .. "\n", "Number" }
+      end
+    end
   end
   vim.api.nvim_echo(chunks, true, {})
 end
diff --git a/lua/lazy/core/loader.lua b/lua/lazy/core/loader.lua
index 0f1b977..833b432 100644
--- a/lua/lazy/core/loader.lua
+++ b/lua/lazy/core/loader.lua
@@ -4,6 +4,7 @@ local Handler = require("lazy.core.handler")
 local Cache = require("lazy.core.cache")
 local Plugin = require("lazy.core.plugin")
 
+---@class LazyCoreLoader
 local M = {}
 
 local DEFAULT_PRIORITY = 50
@@ -450,27 +451,35 @@ function M.colorscheme(name)
   end
 end
 
+function M.auto_load(modname, modpath)
+  local plugin = Plugin.find(modpath)
+  if plugin and modpath:find(plugin.dir, 1, true) == 1 then
+    -- don't load if we're loading specs or if the plugin is already loaded
+    if not (Plugin.loading or plugin._.loaded) then
+      if plugin.module == false then
+        error("Plugin " .. plugin.name .. " is not loaded and is configured with module=false")
+      end
+      M.load(plugin, { require = modname })
+    end
+    return true
+  end
+  return false
+end
+
 ---@param modname string
 function M.loader(modname)
-  local modpath = Cache.find(modname, { rtp = false, paths = Util.get_unloaded_rtp(modname) })
+  local paths = Util.get_unloaded_rtp(modname)
+  local modpath, hash = Cache._Cache.find(modname, { rtp = false, paths = paths })
+  -- print(modname .. " " .. paths[1])
   if modpath then
-    local plugin = Plugin.find(modpath)
-    if plugin and modpath:find(plugin.dir, 1, true) == 1 then
-      -- don't load if we're loading specs or if the plugin is already loaded
-      if not (Plugin.loading or plugin._.loaded) then
-        if plugin.module == false then
-          error("Plugin " .. plugin.name .. " is not loaded and is configured with module=false")
-        end
-        M.load(plugin, { require = modname })
+    M.auto_load(modname, modpath)
+    local mod = package.loaded[modname]
+    if type(mod) == "table" then
+      return function()
+        return mod
       end
-      local mod = package.loaded[modname]
-      if type(mod) == "table" then
-        return function()
-          return mod
-        end
-      end
-      return loadfile(modpath)
     end
+    return Cache.load(modpath, { hash = hash })
   end
 end
 
diff --git a/lua/lazy/init.lua b/lua/lazy/init.lua
index e81c4ec..cb6b732 100644
--- a/lua/lazy/init.lua
+++ b/lua/lazy/init.lua
@@ -34,7 +34,13 @@ function M.setup(spec, opts)
   local start = vim.loop.hrtime()
 
   -- load module cache before anything else
-  if not (opts and opts.performance and opts.performance.cache and opts.performance.cache.enabled == false) then
+  local enable_cache = not (
+    opts
+    and opts.performance
+    and opts.performance.cache
+    and opts.performance.cache.enabled == false
+  )
+  if enable_cache then
     require("lazy.core.cache").enable()
   end
 
@@ -43,8 +49,13 @@ function M.setup(spec, opts)
   local Util = require("lazy.core.util")
   local Config = require("lazy.core.config")
   local Loader = require("lazy.core.loader")
+
   table.insert(package.loaders, 3, Loader.loader)
 
+  if vim.g.profile_loaders then
+    require("lazy.core.cache").profile_loaders()
+  end
+
   Util.track({ plugin = "lazy.nvim" }) -- setup start
   Util.track("module", vim.loop.hrtime() - start)
 
diff --git a/lua/lazy/view/render.lua b/lua/lazy/view/render.lua
index ec5570e..56ac095 100644
--- a/lua/lazy/view/render.lua
+++ b/lua/lazy/view/render.lua
@@ -670,16 +670,21 @@ function M:debug()
   end)
   self:nl()
 
-  self:append("Cache.find()", "LazyH2"):nl()
-  self:props({
-    { "total", Cache.stats.find.total, "Number" },
-    { "time", self:ms(Cache.stats.find.time, 3), "Bold" },
-    { "avg time", self:ms(Cache.stats.find.time / Cache.stats.find.total, 3), "Bold" },
-    { "index", Cache.stats.find.index, "Number" },
-    { "fs_stat", Cache.stats.find.stat, "Number" },
-    { "not found", Cache.stats.find.not_found, "Number" },
-  }, { indent = 2 })
-  self:nl()
+  Util.foreach(Cache.stats, function(name, stats)
+    self:append(name, "LazyH2"):nl()
+    local props = {
+      { "total", stats.total or 0, "Number" },
+      { "time", self:ms(stats.time or 0, 3), "Bold" },
+      { "avg time", self:ms((stats.time or 0) / (stats.total or 0), 3), "Bold" },
+    }
+    for k, v in pairs(stats) do
+      if k ~= "total" and k ~= "time" then
+        props[#props + 1] = { k, v, "Number" }
+      end
+    end
+    self:props(props, { indent = 2 })
+    self:nl()
+  end)
 end
 
 return M