Allow finding modules with multiple roots

lua/lazy/core/util.lua

@@ -240,26 +240,40 @@ function M.get_unloaded_rtp(modname)
   return rtp
 end
 
-function M.find_root(modname)
+function M.find_roots(modname)
   local ret = require("lazy.core.cache").find(modname, {
     rtp = true,
     paths = M.get_unloaded_rtp(modname),
     patterns = { "", ".lua" },
-  })[1]
-  if ret then
-    local root = ret.modpath:gsub("/init%.lua$", ""):gsub("%.lua$", "")
-    return root
+    all = true,
+  })
+
+  local roots = {}
+  local hash = {}
+  for _, root in ipairs(ret) do
+    local path = root.modpath:gsub("/init%.lua$", ""):gsub("%.lua$", "")
+    if not hash[path] then
+      roots[#roots + 1] = path
+      hash[path] = true
     end
   end
 
+  return roots
+end
+
+function M.find_root(modname)
+  return M.find_roots(modname)[1]
+end
+
 ---@param modname string
 ---@param fn fun(modname:string, modpath:string)
 function M.lsmod(modname, fn)
-  local root = M.find_root(modname)
-  if not root then
+  local roots = M.find_roots(modname)
+  if not roots then
     return
   end
 
+  for _, root in ipairs(roots) do
     if vim.loop.fs_stat(root .. ".lua") then
       fn(modname, root .. ".lua")
     end
@@ -274,6 +288,7 @@ function M.lsmod(modname, fn)
       end
     end)
   end
+end
 
 ---@generic T
 ---@param list T[]

tests/core/util_spec.lua

@@ -15,6 +15,82 @@ describe("util", function()
     assert(not vim.loop.fs_stat(Helpers.path("")), "fs root should be deleted")
   end)
 
+  it("find_roots lists the expected single root", function()
+    vim.opt.rtp:append(Helpers.path(""))
+
+    local tests = {
+      {
+        root = "lua/foo",
+        mod = "foo",
+        files = { "lua/foo/one.lua", "lua/foo/two.lua", "lua/foo/init.lua" },
+      },
+      {
+        root = "lua/foo",
+        mod = "foo",
+        files = { "lua/foo/one.lua", "lua/foo/two.lua", "lua/foo.lua" },
+      },
+      {
+        root = "lua/foo",
+        mod = "foo",
+        files = { "lua/foo/one.lua", "lua/foo/two.lua" },
+      },
+      {
+        root = "lua/load-plugins",
+        mod = "load-plugins",
+        files = { "lua/load-plugins.lua" },
+      },
+    }
+
+    for t, test in ipairs(tests) do
+      Helpers.fs_rm("")
+      assert(not vim.loop.fs_stat(Helpers.path("")), "fs root should be deleted")
+
+      local files = Helpers.fs_create(test.files)
+
+      Cache.reset()
+      local roots = Util.find_roots(test.mod)
+      assert.equal(#roots, 1, "wrong number of roots found for " .. test.mod .. " (test " .. t .. ")")
+
+      local expected_root = Helpers.path(test.root)
+      assert.same({expected_root}, roots, "wrong roots found (test " .. t .. ")")
+    end
+  end)
+
+  it("find_roots lists the correct multiple roots", function()
+    vim.opt.rtp:append(Helpers.path("first"))
+    vim.opt.rtp:append(Helpers.path("second"))
+
+    local tests = {
+      {
+        roots = { "first/lua/foo", "second/lua/foo" },
+        mod = "foo",
+        files = { "first/lua/foo/init.lua", "second/lua/foo/local.lua" },
+      },
+      {
+        roots = { "first/lua/foo", "second/lua/foo" },
+        mod = "foo",
+        files = { "first/lua/foo.lua", "second/lua/foo/baz.lua" },
+      },
+    }
+
+    for t, test in ipairs(tests) do
+      Helpers.fs_rm("")
+      assert(not vim.loop.fs_stat(Helpers.path("")), "fs root should be deleted")
+
+      local files = Helpers.fs_create(test.files)
+
+      Cache.reset()
+      local roots = Util.find_roots(test.mod)
+      assert(#roots > 0, "no roots found for " .. test.mod .. " (test " .. t .. ")")
+
+      local expected_roots = {}
+      for _, root in ipairs(test.roots) do
+        expected_roots[#expected_roots + 1] = Helpers.path(root)
+      end
+      assert.same(expected_roots, roots)
+    end
+  end)
+
   it("lsmod lists all mods in dir", function()
     vim.opt.rtp:append(Helpers.path(""))
     local tests = {
@@ -76,6 +152,50 @@ describe("util", function()
     end
   end)
 
+  it("lsmod lists modules in multiple roots", function()
+    vim.opt.rtp:append(Helpers.path("first"))
+    vim.opt.rtp:append(Helpers.path("second"))
+
+    local tests = {
+      {
+        roots = { "first/lua/foo", "second/lua/foo" },
+        mod = "foo",
+        files = { "first/lua/foo/init.lua", "second/lua/foo/local.lua" },
+        mods = { "foo", "foo.local" },
+      },
+      {
+        roots = { "first/lua/foo", "second/lua/foo" },
+        mod = "foo",
+        files = { "first/lua/foo.lua", "second/lua/foo/baz.lua" },
+        mods = { "foo", "foo.baz" },
+      },
+    }
+
+    for t, test in ipairs(tests) do
+      Helpers.fs_rm("")
+      assert(not vim.loop.fs_stat(Helpers.path("")), "fs root should be deleted")
+
+      local files = Helpers.fs_create(test.files)
+
+      Cache.reset()
+      local roots = Util.find_roots(test.mod)
+      assert(#roots > 0, "no roots found for " .. test.mod .. " (test " .. t .. ")")
+
+      local expected_roots = {}
+      for _, root in ipairs(test.roots) do
+        expected_roots[#expected_roots + 1] = Helpers.path(root)
+      end
+      assert.same(expected_roots, roots)
+
+      mods = {}
+      Util.lsmod(test.mod, function(modname, modpath)
+        mods[#mods + 1] = modname
+      end)
+      table.sort(mods)
+      assert.same(test.mods, mods)
+    end
+  end)
+
   it("find the correct root with dels", function()
     Cache.reset()
     vim.opt.rtp:append(Helpers.path("old"))