From 6fe425c91acbf2b9b948b23673e22a0c61150249 Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Fri, 2 Dec 2022 12:43:34 +0100
Subject: [PATCH] perf: caching strategy is now configurable

---
 lua/lazy/core/{module.lua => cache.lua} | 95 ++++++++++++++++++-------
 lua/lazy/core/config.lua                |  7 +-
 lua/lazy/init.lua                       | 16 ++++-
 3 files changed, 89 insertions(+), 29 deletions(-)
 rename lua/lazy/core/{module.lua => cache.lua} (68%)

diff --git a/lua/lazy/core/module.lua b/lua/lazy/core/cache.lua
similarity index 68%
rename from lua/lazy/core/module.lua
rename to lua/lazy/core/cache.lua
index a1070db..a1ff412 100644
--- a/lua/lazy/core/module.lua
+++ b/lua/lazy/core/cache.lua
@@ -5,7 +5,19 @@ local uv = vim.loop
 local M = {}
 M.dirty = false
 
-local cache_path = vim.fn.stdpath("state") .. "/lazy.state"
+---@class LazyCacheConfig
+M.config = {
+  enabled = true,
+  path = vim.fn.stdpath("state") .. "/lazy.state",
+  -- choose what should be cached
+  --  * lazy: cache all lazy.nvim core modules and your config files
+  --  * init: all of the above and any module needed to init your plugins
+  --  * VimEnter: any module till VimEnter
+  --  * VeryLazy: any module till VeryLazy
+  --  * allthethings: all mdules. Not recommended
+  strategy = "VimEnter", ---@type "lazy"|"init"|"VimEnter"|"allthethings"
+}
+
 ---@type CacheHash
 local cache_hash
 
@@ -27,13 +39,25 @@ function M.check_load(modname, modpath)
   require("lazy.core.loader").autoload(modname, modpath)
 end
 
+---@param step? string
+function M.disable(step)
+  if not M.enabled then
+    return
+  end
+  if step and M.config.strategy ~= step then
+    return
+  end
+
+  local idx = M.idx()
+  if idx then
+    table.remove(package.loaders, idx)
+  end
+  M.enabled = false
+end
+
 ---@param modname string
 ---@return any
 function M.loader(modname)
-  if not M.enabled then
-    return "lazy loader is disabled"
-  end
-
   local entry = M.cache[modname]
 
   local chunk, err
@@ -68,11 +92,10 @@ function M.loader(modname)
   return chunk or error(err)
 end
 
----@param modname string
-function M.find(modname)
+function M.idx()
   -- update our loader position if needed
   if package.loaders[M.loader_idx] ~= M.loader then
-    M.loader_idx = 1
+    M.loader_idx = nil
     ---@diagnostic disable-next-line: no-unknown
     for i, loader in ipairs(package.loaders) do
       if loader == M.loader then
@@ -81,29 +104,49 @@ function M.find(modname)
       end
     end
   end
+  return M.loader_idx
+end
 
-  -- find the module and its modpath
-  for i = M.loader_idx + 1, #package.loaders do
-    ---@diagnostic disable-next-line: no-unknown
-    local chunk = package.loaders[i](modname)
-    if type(chunk) == "function" then
-      local info = debug.getinfo(chunk, "S")
-      return chunk, (info.what ~= "C" and info.source:sub(2))
+---@param modname string
+function M.find(modname)
+  if M.idx() then
+    -- find the module and its modpath
+    for i = M.loader_idx + 1, #package.loaders do
+      ---@diagnostic disable-next-line: no-unknown
+      local chunk = package.loaders[i](modname)
+      if type(chunk) == "function" then
+        local info = debug.getinfo(chunk, "S")
+        return chunk, (info.what ~= "C" and info.source:sub(2))
+      end
     end
   end
 end
 
-function M.setup()
+---@param opts? LazyConfig
+function M.setup(opts)
+  -- no fancy deep extend here. just set the options
+  if opts and opts.performance and opts.performance.cache then
+    for k, v in pairs(opts.performance.cache) do
+      M.config[k] = v
+    end
+  end
+
   M.load_cache()
   table.insert(package.loaders, M.loader_idx, M.loader)
 
-  vim.api.nvim_create_autocmd("VimEnter", {
-    once = true,
-    callback = function()
-      -- startup done, so stop caching
-      M.enabled = false
-    end,
-  })
+  if M.config.strategy == "VimEnter" then
+    vim.api.nvim_create_autocmd("VimEnter", {
+      once = true,
+      callback = function()
+        -- use schedule so all other VimEnter handlers will have run
+        vim.schedule(function()
+          -- startup done, so stop caching
+          M.disable()
+        end)
+      end,
+    })
+  end
+  return M
 end
 
 ---@return CacheHash?
@@ -118,7 +161,7 @@ function M.eq(h1, h2)
 end
 
 function M.save_cache()
-  local f = assert(uv.fs_open(cache_path, "w", 438))
+  local f = assert(uv.fs_open(M.config.path, "w", 438))
   for modname, entry in pairs(M.cache) do
     if entry.used > os.time() - M.ttl then
       entry.modname = modname
@@ -142,7 +185,7 @@ end
 
 function M.load_cache()
   M.cache = {}
-  local f = uv.fs_open(cache_path, "r", 438)
+  local f = uv.fs_open(M.config.path, "r", 438)
   if f then
     cache_hash = uv.fs_fstat(f) --[[@as CacheHash]]
     local data = uv.fs_read(f, cache_hash.size, 0) --[[@as string]]
@@ -172,7 +215,7 @@ function M.autosave()
   vim.api.nvim_create_autocmd("VimLeavePre", {
     callback = function()
       if M.dirty then
-        local hash = M.hash(cache_path)
+        local hash = M.hash(M.config.path)
         -- abort when the file was changed in the meantime
         if hash == nil or M.eq(cache_hash, hash) then
           M.save_cache()
diff --git a/lua/lazy/core/config.lua b/lua/lazy/core/config.lua
index de04c3c..b452625 100644
--- a/lua/lazy/core/config.lua
+++ b/lua/lazy/core/config.lua
@@ -48,6 +48,10 @@ M.defaults = {
     },
     throttle = 20, -- how frequently should the ui process render events
   },
+  performance = {
+    ---@type LazyCacheConfig
+    cache = nil,
+  },
 }
 
 M.ns = vim.api.nvim_create_namespace("lazy")
@@ -72,6 +76,7 @@ M.options = {}
 function M.setup(spec, opts)
   M.spec = spec
   M.options = vim.tbl_deep_extend("force", M.defaults, opts or {})
+  M.options.performance.cache = require("lazy.core.cache")
   M.root = M.options.package.path .. "/pack/" .. M.options.package.name .. "/opt"
 
   if M.options.package.reset then
@@ -84,7 +89,7 @@ function M.setup(spec, opts)
     pattern = "VeryLazy",
     once = true,
     callback = function()
-      require("lazy.core.module").autosave()
+      require("lazy.core.cache").autosave()
       require("lazy.view").setup()
     end,
   })
diff --git a/lua/lazy/init.lua b/lua/lazy/init.lua
index 2b760b9..1883d82 100644
--- a/lua/lazy/init.lua
+++ b/lua/lazy/init.lua
@@ -8,8 +8,12 @@ function M.setup(spec, opts)
   end
   local start = vim.loop.hrtime()
 
-  -- load module cache before anything else
-  require("lazy.core.module").setup()
+  local Cache
+  if not (opts and opts.performance and opts.performance.cache and opts.performance.cache.enabled == false) then
+    -- load module cache before anything else
+    Cache = require("lazy.core.cache").setup(opts)
+  end
+
   local Util = require("lazy.core.util")
   local Config = require("lazy.core.config")
   local Loader = require("lazy.core.loader")
@@ -36,9 +40,17 @@ function M.setup(spec, opts)
     Config.plugins["lazy.nvim"]._.loaded = { time = delta, source = "init.lua" }
   end
 
+  if Cache then
+    Cache.disable("lazy")
+  end
+
   -- load plugins with lazy=false or Plugin.init
   Loader.init_plugins()
 
+  if Cache then
+    Cache.disable("init")
+  end
+
   -- all done!
   vim.cmd("do User LazyDone")
 end