From 953c2791d8c391bf720ae68e734078bb558329f6 Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Tue, 3 Jan 2023 22:31:18 +0100
Subject: [PATCH] fix(util): made `Util.lsmod` more robust. See #298

---
 lua/lazy/core/cache.lua  |  8 +++---
 lua/lazy/core/util.lua   |  6 ++---
 tests/core/util_spec.lua | 53 ++++++++++++++++++++++++++++++++++------
 3 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/lua/lazy/core/cache.lua b/lua/lazy/core/cache.lua
index eee6b16..3efcb55 100644
--- a/lua/lazy/core/cache.lua
+++ b/lua/lazy/core/cache.lua
@@ -257,14 +257,14 @@ function M.get_topmods(path)
 end
 
 ---@param modname string
----@return string?, string?
-function M.find_dir(modname)
+---@return string?
+function M.find_root(modname)
   if M.cache[modname] then
     -- check if modname is in cache
     local modpath = M.cache[modname].modpath
     if M.check_path(modname, modpath) then
       local root = modpath:gsub("/init%.lua$", ""):gsub("%.lua$", "")
-      return root, modpath
+      return root
     end
   else
     -- in case modname is just a directory and not a real mod,
@@ -286,7 +286,7 @@ function M.find_dir(modname)
   local modpath = M.find(modname, { patterns = { "" } })
   if modpath then
     local root = modpath:gsub("/init%.lua$", ""):gsub("%.lua$", "")
-    return root, (modpath ~= root and modpath or nil)
+    return root
   end
 end
 
diff --git a/lua/lazy/core/util.lua b/lua/lazy/core/util.lua
index 31acdec..b1ba858 100644
--- a/lua/lazy/core/util.lua
+++ b/lua/lazy/core/util.lua
@@ -195,13 +195,13 @@ end
 ---@param fn fun(modname:string, modpath:string)
 function M.lsmod(modname, fn)
   local Cache = require("lazy.core.cache")
-  local root, modpath = Cache.find_dir(modname)
+  local root = Cache.find_root(modname)
   if not root then
     return
   end
 
-  if modpath and not modpath:find("/init%.lua$") and vim.loop.fs_stat(modpath) then
-    fn(modname, modpath)
+  if vim.loop.fs_stat(root .. ".lua") then
+    fn(modname, root .. ".lua")
   end
 
   M.ls(root, function(path, name, type)
diff --git a/tests/core/util_spec.lua b/tests/core/util_spec.lua
index 3dee0a8..655b71f 100644
--- a/tests/core/util_spec.lua
+++ b/tests/core/util_spec.lua
@@ -10,31 +10,68 @@ describe("util", function()
   it("lsmod lists all mods in dir", function()
     local tests = {
       {
+        root = "lua/foo",
+        mod = "foo",
         files = { "lua/foo/one.lua", "lua/foo/two.lua", "lua/foo/init.lua" },
-        mods = { "foo", "foo.one", "foo.two" },
+        mods = { "foo.one", "foo.two", "foo" },
       },
       {
+        root = "lua/foo",
+        mod = "foo",
         files = { "lua/foo/one.lua", "lua/foo/two.lua", "lua/foo.lua" },
-        mods = { "foo", "foo.one", "foo.two" },
+        mods = { "foo.one", "foo.two", "foo" },
       },
       {
+        root = "lua/foo",
+        mod = "foo",
         files = { "lua/foo/one.lua", "lua/foo/two.lua" },
         mods = { "foo.one", "foo.two" },
       },
+      {
+        root = "lua/load-plugins",
+        mod = "load-plugins",
+        files = { "lua/load-plugins.lua" },
+        mods = { "load-plugins" },
+      },
     }
 
     vim.opt.rtp:append(Helpers.path(""))
-    for _, test in ipairs(tests) do
-      Cache.cache = {}
-      table.sort(test.mods)
+    for t, test in ipairs(tests) do
+      local expected = vim.deepcopy(test.mods)
+      table.sort(expected)
       Helpers.fs_rm("")
-      Helpers.fs_create(test.files)
+      local files = Helpers.fs_create(test.files)
+
+      -- test with empty cache
+      Cache.cache = {}
+      Cache.indexed = {}
+      Cache.indexed_rtp = false
+      local root = Cache.find_root(test.mod)
+      assert(root, "no root found for " .. test.mod .. " (test " .. t .. ")")
+      assert.same(Helpers.path(test.root), root)
       local mods = {}
-      Util.lsmod("foo", function(modname, modpath)
+      Util.lsmod(test.mod, function(modname, modpath)
         mods[#mods + 1] = modname
       end)
       table.sort(mods)
-      assert.same(test.mods, mods)
+      assert.same(expected, mods)
+
+      -- fill the cache
+      Cache.cache = {}
+      for i, file in ipairs(files) do
+        Cache.cache[test.mods[i]] = { modpath = file }
+      end
+      Cache.indexed = {}
+      Cache.indexed_rtp = false
+      root = Cache.find_root(test.mod)
+      assert(root, "no root found for " .. test.mod .. " (test " .. t .. ")")
+      assert.same(Helpers.path(test.root), root)
+      mods = {}
+      Util.lsmod(test.mod, function(modname, modpath)
+        mods[#mods + 1] = modname
+      end)
+      table.sort(mods)
+      assert.same(expected, mods)
     end
   end)
 end)