From a543134b8c1b17c2396a757b08951b6d91b14402 Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Mon, 21 Nov 2022 21:50:16 +0100
Subject: [PATCH] perf: way better compilation and caching

---
 lua/lazy/cache.lua  | 392 +++++++++++++++++++++++++++++++++++++-------
 lua/lazy/config.lua |  33 +---
 lua/lazy/init.lua   |  36 ++--
 lua/lazy/loader.lua |  16 +-
 lua/lazy/plugin.lua |  42 +++--
 lua/lazy/util.lua   |   4 +-
 6 files changed, 390 insertions(+), 133 deletions(-)

diff --git a/lua/lazy/cache.lua b/lua/lazy/cache.lua
index 20bf4bc..3dd167e 100644
--- a/lua/lazy/cache.lua
+++ b/lua/lazy/cache.lua
@@ -1,87 +1,363 @@
-local cache_file = vim.fn.stdpath("cache") .. "/lazy/cache.mpack"
-vim.fn.mkdir(vim.fn.fnamemodify(cache_file, ":p:h"), "p")
-
 local M = {}
----@alias CacheEntry {hash:string, chunk:string, used:boolean}
 
----@type table<string, CacheEntry>
-M.cache = {}
+---@class CacheOptions
+M.options = {
+  module = "config.plugins",
+  cache = vim.fn.stdpath("state") .. "/lazy/plugins.state",
+  trim = false,
+}
+
 M.dirty = false
-M.did_setup = false
+M.cache_hash = ""
+
+---@alias ModEntry {file: string, hash?:string, chunk?: string|fun(), used?: boolean}
+---@type table<string,ModEntry>
+M.modules = {}
+
+---@type LazyState?
+M.state = nil
+
+---@alias DirEntry {name: string, type: "file"|"directory"|"link"}
+
+function M.walk(dir, modname, fn)
+  local d = vim.loop.fs_opendir(dir, nil, 100)
+  if d then
+    ---@type DirEntry[]
+    local entries = vim.loop.fs_readdir(d)
+    while entries do
+      for _, entry in ipairs(entries) do
+        local path = dir .. "/" .. entry.name
+        if entry.type == "directory" then
+          M.walk(path, modname .. "." .. entry.name, fn)
+        else
+          local child = entry.name == "init.lua" and modname or (modname .. "." .. entry.name:match("^(.*)%.lua$"))
+          if child then
+            fn(child, path)
+          end
+        end
+      end
+      entries = vim.loop.fs_readdir(d)
+    end
+    vim.loop.fs_closedir(d)
+  end
+end
 
 function M.hash(modpath)
   local stat = vim.loop.fs_stat(modpath)
   if stat then
     return stat.mtime.sec .. stat.mtime.nsec .. stat.size
   end
-  error("Could not hash " .. modpath)
 end
 
-function M.load_cache()
-  local f = io.open(cache_file, "rb")
-  if f then
-    M.cache = vim.mpack.decode(f:read("*a")) or {}
-    f:close()
-  end
-end
-
-function M.save_cache()
-  if M.dirty then
-    for key, entry in pairs(M.cache) do
-      if not entry.used then
-        M.cache[key] = nil
-      end
-      entry.used = nil
+---@param opts? CacheOptions
+function M.boot(opts)
+  if opts then
+    for k, _ in pairs(M.options) do
+      M.options[k] = opts[k] or M.options[k]
+    end
+  end
+  M.load_state()
+
+  -- preload core modules
+  local root = debug.getinfo(1, "S").source:sub(2)
+  root = vim.fn.fnamemodify(root, ":p:h")
+  for _, modname in ipairs({ "util", "config", "plugin", "loader" }) do
+    local file = root .. "/" .. modname .. ".lua"
+    modname = "lazy." .. modname
+    if not M.modules[modname] then
+      M.modules[modname] = { file = file }
+    end
+    package.preload[modname] = function()
+      return M.load(modname)
     end
-    local f = assert(io.open(cache_file, "wb"))
-    f:write(vim.mpack.encode(M.cache))
-    f:close()
   end
 end
 
 function M.setup()
-  M.load_cache()
-  vim.api.nvim_create_autocmd("VimLeave", {
+  vim.api.nvim_create_autocmd("User", {
+    pattern = "LazyDone",
+    once = true,
     callback = function()
-      M.save_cache()
+      vim.api.nvim_create_autocmd("VimLeavePre", {
+        callback = function()
+          if M.dirty then
+            local hash = M.hash(M.options.cache)
+            -- abort when the file was changed in the meantime
+            if M.hash == nil or M.cache_hash == hash then
+              M.compile()
+            end
+          end
+        end,
+      })
     end,
   })
-end
 
-function M.load(modpath, modname)
-  if not M.did_setup then
-    M.setup()
-    M.did_setup = true
-  end
-  if type(package.loaded[modname]) ~= "table" then
-    ---@type fun()?, string?
-    local chunk, err
-    local entry = M.cache[modname]
-
-    if entry and M.hash(modpath) == entry.hash then
-      entry.used = true
-      chunk, err = loadstring(entry.chunk, "@" .. modpath)
-    end
-
-    -- not cached, or failed to load chunk
-    if not chunk then
-      vim.schedule(function()
-        vim.notify("not cached")
-      end)
-      chunk, err = loadfile(modpath)
-      if chunk then
-        M.cache[modname] = { hash = M.hash(modpath), chunk = string.dump(chunk, true), used = true }
-        M.dirty = true
+  if M.state and M.load_plugins() then
+    return true
+  else
+    M.dirty = true
+    -- FIXME: what if module is a file
+    local root = vim.fn.stdpath("config") .. "/lua/" .. M.options.module:gsub("%.", "/")
+    if vim.loop.fs_stat(root .. ".lua") then
+      if not M.modules[M.options.module] then
+        M.modules[M.options.module] = { file = root .. ".lua" }
       end
     end
+    M.walk(root, M.options.module, function(modname, modpath)
+      if not M.modules[modname] then
+        M.modules[modname] = { file = modpath }
+      end
+    end)
+  end
+end
 
-    if not chunk then
+---@param modname string
+function M.load(modname)
+  local info = M.modules[modname]
+
+  if type(package.loaded[modname]) == "table" then
+    if info then
+      info.used = true
+    end
+    return package.loaded[modname]
+  end
+
+  if info then
+    local hash = M.hash(info.file)
+    if hash ~= info.hash then
+      info.chunk = nil
+    end
+    local err
+    if not info.chunk then
+      vim.schedule(function()
+        vim.notify("loading " .. modname)
+      end)
+      info.chunk, err = loadfile(info.file)
+      info.hash = hash
+      M.dirty = true
+    end
+    if type(info.chunk) == "string" then
+      info.chunk, err = loadstring(info.chunk --[[@as string]], "@" .. info.file)
+    end
+    if not info.chunk then
       error(err)
     end
-    ---@diagnostic disable-next-line: no-unknown
-    package.loaded[modname] = chunk()
+    info.used = true
+    ---@type table
+    local mod = info.chunk()
+    package.loaded[modname] = mod
+    return mod
   end
-  return package.loaded[modname]
+end
+
+---@param state LazyState
+function M.write(state)
+  local chunks = state.chunks
+  state.chunks = nil
+
+  local header = loadstring("return " .. M.dump(state))
+  assert(header)
+  table.insert(chunks, string.dump(header, true))
+
+  vim.fn.mkdir(vim.fn.fnamemodify(M.options.cache, ":p:h"), "p")
+  local f = assert(io.open(M.options.cache, "wb"))
+  for _, chunk in ipairs(chunks) do
+    f:write(tostring(#chunk), "\0", chunk)
+  end
+  f:close()
+end
+
+---@return LazyState?
+function M.read()
+  M.cache_hash = M.hash(M.options.cache)
+  local f = io.open(M.options.cache, "rb")
+  if f then
+    ---@type string
+    local data = f:read("*a")
+    f:close()
+
+    local from = 1
+    local to = data:find("\0", from, true)
+    ---@type string[]
+    local chunks = {}
+    while to do
+      local len = tonumber(data:sub(from, to - 1))
+      from = to + 1
+      local chunk = data:sub(from, from + len - 1)
+      table.insert(chunks, chunk)
+      from = from + len
+      to = data:find("\0", from, true)
+    end
+
+    local state = loadstring(table.remove(chunks))
+    assert(state)
+    ---@type LazyState
+    local ret = state()
+    ret.chunks = chunks
+    return ret
+  end
+end
+
+function M.compile()
+  local Config = require("lazy.config")
+
+  ---@class LazyState
+  local state = {
+    ---@type LazyPlugin[]
+    plugins = {},
+    ---@type table<string, {file:string, hash:string, chunk:number}>
+    modules = {},
+    loaders = require("lazy.loader").loaders,
+    -- config = Config.options,
+    ---@type string[]
+    chunks = {},
+  }
+
+  local skip = { installed = true, loaded = true }
+
+  -- plugins
+  for _, plugin in pairs(Config.plugins) do
+    -- mark module as used
+    if M.modules[plugin.modname] then
+      ---@diagnostic disable-next-line: no-unknown
+      M.modules[plugin.modname].used = true
+    end
+
+    ---@type LazyPlugin | {_chunks: string[] | table<string, number>}
+    local save = {}
+    table.insert(state.plugins, save)
+    for k, v in pairs(plugin) do
+      if type(v) == "function" then
+        save._chunks = save._chunks or {}
+        if plugin.modname then
+          table.insert(save._chunks, k)
+        else
+          table.insert(state.chunks, string.dump(v, M.options.trim))
+          save._chunks[k] = #state.chunks
+        end
+      elseif not skip[k] then
+        save[k] = v
+      end
+    end
+  end
+
+  -- modules
+  for modname, entry in pairs(M.modules) do
+    if entry.used and entry.chunk then
+      table.insert(
+        state.chunks,
+        type(entry.chunk) == "string" and entry.chunk or string.dump(entry.chunk --[[@as fun()]], M.options.trim)
+      )
+      state.modules[modname] = { file = entry.file, hash = entry.hash, chunk = #state.chunks }
+    end
+  end
+  M.write(state)
+end
+
+function M._dump(value, result)
+  local t = type(value)
+  if t == "number" or t == "boolean" then
+    table.insert(result, tostring(value))
+  elseif t == "string" then
+    table.insert(result, ("%q"):format(value))
+  elseif t == "table" then
+    table.insert(result, "{")
+    local i = 1
+    ---@diagnostic disable-next-line: no-unknown
+    for k, v in pairs(value) do
+      if k == i then
+      elseif type(k) == "string" then
+        table.insert(result, ("[%q]="):format(k))
+      else
+        table.insert(result, k .. "=")
+      end
+      M._dump(v, result)
+      table.insert(result, ",")
+      i = i + 1
+    end
+    table.insert(result, "}")
+  else
+    error("Unsupported type " .. t)
+  end
+end
+
+function M.dump(value)
+  local result = {}
+  M._dump(value, result)
+  return table.concat(result, "")
+end
+
+function M.load_state()
+  M.state = M.read()
+
+  if not M.state then
+    return
+  end
+  local reload = false
+  for modname, entry in pairs(M.state.modules) do
+    entry.chunk = M.state.chunks[entry.chunk]
+    ---@cast entry ModEntry
+    if M.hash(entry.file) ~= entry.hash then
+      -- keep loading modules, but reset state (reload plugins)
+      reload = true
+    end
+    M.modules[modname] = entry
+  end
+  if reload then
+    M.state = nil
+  end
+end
+
+local function load_plugin(plugin, fun, ...)
+  local mod = M.load(plugin.modname)
+  for k, v in pairs(mod) do
+    if type(v) == "function" then
+      plugin[k] = v
+    end
+  end
+  return mod[fun](...)
+end
+
+function M.load_plugins()
+  local Config = require("lazy.config")
+
+  if not vim.deepcopy(Config.options, M.state.config) then
+    return false
+  end
+
+  -- plugins
+  for _, plugin in ipairs(M.state.plugins) do
+    plugin.loaded = false
+    plugin.installed = vim.loop.fs_stat(plugin.dir) and true
+    if plugin._chunks then
+      if plugin.modname then
+        for _, fun in ipairs(plugin._chunks) do
+          plugin[fun] = function(...)
+            return load_plugin(plugin, fun, ...)
+          end
+        end
+      else
+        for fun, value in pairs(plugin._chunks) do
+          plugin[fun] = function(...)
+            plugin[fun] = loadstring(M.state.chunks[value])
+            return plugin[fun](...)
+          end
+        end
+      end
+      plugin._chunks = nil
+    end
+  end
+
+  -- loaders
+  local Loader = require("lazy.loader")
+  Loader.loaders = M.state.loaders
+
+  -- save plugins
+  Config.plugins = {}
+  for _, plugin in ipairs(M.state.plugins) do
+    Config.plugins[plugin.name] = plugin
+  end
+  return true
 end
 
 return M
diff --git a/lua/lazy/config.lua b/lua/lazy/config.lua
index f7f5599..4936151 100644
--- a/lua/lazy/config.lua
+++ b/lua/lazy/config.lua
@@ -5,16 +5,10 @@ local M = {}
 ---@class LazyConfig
 M.defaults = {
   opt = true,
-  plugins = {},
+  plugins = "config.plugins",
   plugins_local = {
     path = vim.fn.expand("~/projects"),
-    patterns = {
-      "folke",
-    },
-  },
-  plugins_config = {
-    module = "plugins",
-    path = vim.fn.stdpath("config") .. "/lua/plugins",
+    patterns = {},
   },
   package_path = vim.fn.stdpath("data") .. "/site/pack/lazy",
 }
@@ -27,32 +21,11 @@ M.plugins = {}
 ---@type LazyConfig
 M.options = {}
 
----@type table<string, string>
-M.has_config = {}
-
 ---@param opts? LazyConfig
 function M.setup(opts)
   M.options = vim.tbl_deep_extend("force", M.defaults, opts or {})
 
-  vim.fn.mkdir(M.options.package_path, "p")
-
-  for _, entry in ipairs(Util.scandir(M.options.plugins_config.path)) do
-    local name, modpath
-
-    if entry.type == "file" then
-      modpath = entry.path
-      name = entry.name:match("(.*)%.lua")
-    elseif entry.type == "directory" then
-      modpath = M.options.plugins_config.path .. "/" .. entry.name .. "/init.lua"
-      if vim.loop.fs_stat(modpath) then
-        name = entry.name
-      end
-    end
-
-    if name then
-      M.has_config[M.options.plugins_config.module .. "." .. name] = modpath
-    end
-  end
+  -- vim.fn.mkdir(M.options.package_path, "p")
 
   vim.api.nvim_create_autocmd("User", {
     pattern = "VeryLazy",
diff --git a/lua/lazy/init.lua b/lua/lazy/init.lua
index 717e701..3d36501 100644
--- a/lua/lazy/init.lua
+++ b/lua/lazy/init.lua
@@ -4,28 +4,42 @@ local M = {}
 function M.setup(opts)
   --FIXME:  preload()
 
+  local Cache = require("lazy.cache")
+
+  local start = vim.loop.hrtime()
+  Cache.boot()
+
   local Util = require("lazy.util")
   local Config = require("lazy.config")
   local Plugin = require("lazy.plugin")
 
+  Util.track("lazy_boot", vim.loop.hrtime() - start)
+
   Util.track("lazy_setup")
 
   Util.track("lazy_config")
   Config.setup(opts)
   Util.track()
 
-  Util.track("plugin_normalize")
-  Plugin.normalize(Config.options.plugins)
-  if not Config.plugins.lazy then
-    Plugin.plugin({
-      "folke/lazy.nvim",
-      opt = false,
-    })
-  end
-  Util.track()
+  Util.track("lazy_plugins")
+  if not Cache.setup() then
+    vim.schedule(function()
+      vim.notify("Reloading")
+    end)
+    Util.track("plugin_normalize")
+    Plugin.normalize(require(Config.options.plugins))
+    if not Config.plugins.lazy then
+      Plugin.plugin({
+        "folke/lazy.nvim",
+        opt = false,
+      })
+    end
+    Util.track()
 
-  Util.track("plugin_process")
-  Plugin.process()
+    Util.track("plugin_process")
+    Plugin.process()
+    Util.track()
+  end
   Util.track()
 
   Util.track("lazy_install")
diff --git a/lua/lazy/loader.lua b/lua/lazy/loader.lua
index aff1a2a..b7aacf5 100644
--- a/lua/lazy/loader.lua
+++ b/lua/lazy/loader.lua
@@ -3,7 +3,7 @@ local Config = require("lazy.config")
 
 local M = {}
 
----@alias LoaderType "event"|"ft"|"module"|"keys"|"cmd"
+---@alias LoaderType "event"|"ft"|"module"|"keys"|"cmd"|"init"
 ---@type LoaderType[]
 M.types = {
   "event",
@@ -13,20 +13,17 @@ M.types = {
   "cmd",
 }
 
----@type table<LoaderType, table<string, string[]>>
-M.loaders = {}
+---@class LazyLoaders: table<LoaderType, table<string, string[]>>|{init: string[]}
+M.loaders = { init = {} }
 
 for _, type in ipairs(M.types) do
   M.loaders[type] = {}
 end
 
----@type LazyPlugin[]
-M.need_setup = {}
-
 ---@param plugin LazyPlugin
 function M.add(plugin)
   if plugin.init or plugin.opt == false and plugin.config then
-    table.insert(M.need_setup, plugin)
+    table.insert(M.loaders.init, plugin.name)
   end
 
   for _, loader_type in ipairs(M.types) do
@@ -130,7 +127,8 @@ end
 
 function M.init_plugins()
   Util.track("loader_plugin_init")
-  for _, plugin in ipairs(M.need_setup) do
+  for _, name in ipairs(M.loaders.init) do
+    local plugin = Config.plugins[name]
     if plugin.init then
       Util.track(plugin.name)
       plugin.init()
@@ -152,7 +150,7 @@ function M.module(modname)
     local plugins = M.loaders.module[name]
     if plugins then
       M.load(plugins)
-      M.loaders.module[name] = nil
+      -- M.loaders.module[name] = nil
     end
     idx = modname:find(".", idx + 1, true)
   end
diff --git a/lua/lazy/plugin.lua b/lua/lazy/plugin.lua
index fe79161..01d86e9 100644
--- a/lua/lazy/plugin.lua
+++ b/lua/lazy/plugin.lua
@@ -5,14 +5,15 @@ local Cache = require("lazy.cache")
 
 local M = {}
 
----@class LazyPlugin: {[1]: string}
+---@class LazyPlugin
+---@field [1] string
+---@field name string display name and name used for plugin config files
+---@field pack string package name
 ---@field uri string
----@field as? string
+---@field modname? string
 ---@field branch? string
 ---@field dir string
 ---@field opt? boolean
----@field name string display name and name used for plugin config files
----@field pack string package name
 ---@field init? fun(LazyPlugin) Will always be run
 ---@field config? fun(LazyPlugin) Will be executed when loading the plugin
 ---@field event? string|string[]
@@ -74,27 +75,22 @@ end
 ---@param plugin LazyPlugin
 function M.process_config(plugin)
   local name = plugin.name
-  local modname = Config.options.plugins_config.module .. "." .. name
+  local modname = Config.options.plugins .. "." .. name
 
-  local file = Config.has_config[modname]
-  if file then
-    -- use dofile and then add to modules. Since we know where to look, this is faster
-    local ok, spec = pcall(Cache.load, file, modname)
-    if ok then
-      -- add to loaded modules
-      if spec.requires then
-        spec.requires = M.normalize(spec.requires)
-      end
-
-      ---@diagnostic disable-next-line: no-unknown
-      for k, v in pairs(spec) do
-        ---@diagnostic disable-next-line: no-unknown
-        plugin[k] = v
-      end
-      M.plugin(plugin)
-    else
-      Util.error("Failed to load plugin config for " .. name .. "\n" .. spec)
+  local spec = Cache.load(modname)
+  if spec then
+    -- add to loaded modules
+    if spec.requires then
+      spec.requires = M.normalize(spec.requires)
     end
+
+    ---@diagnostic disable-next-line: no-unknown
+    for k, v in pairs(spec) do
+      ---@diagnostic disable-next-line: no-unknown
+      plugin[k] = v
+    end
+    plugin.modname = modname
+    M.plugin(plugin)
   end
 end
 
diff --git a/lua/lazy/util.lua b/lua/lazy/util.lua
index 7e09599..2062335 100644
--- a/lua/lazy/util.lua
+++ b/lua/lazy/util.lua
@@ -25,7 +25,7 @@ function M.track(name, time)
 end
 
 function M.time()
-  return vim.loop.hrtime() / 1000000
+  return vim.loop.hrtime()
 end
 
 function M.file_exists(file)
@@ -115,7 +115,7 @@ function M.profile()
 
     table.insert(
       lines,
-      ("  "):rep(depth) .. "- " .. entry.name .. ": **" .. math.floor((entry.time or 0) * 100) / 100 .. "ms**"
+      ("  "):rep(depth) .. "- " .. entry.name .. ": **" .. math.floor((entry.time or 0) / 1e6 * 100) / 100 .. "ms**"
     )
 
     for _, child in ipairs(entry) do