refactor(registry): refactor registry initialization

4 files changed, 144 insertions, 100 deletions

diff --git a/lua/mason-registry/sources/file.lua b/lua/mason-registry/sources/file.lua
index 62e7d7a4..84566479 100644
--- a/lua/mason-registry/sources/file.lua
+++ b/lua/mason-registry/sources/file.lua
@@ -14,6 +14,7 @@ local util = require "mason-registry.sources.util"
 local Channel = async_control.Channel
 
 ---@class FileRegistrySourceSpec
+---@field id string
 ---@field path string
 
 ---@class FileRegistrySource : RegistrySource
@@ -28,6 +29,7 @@ function FileRegistrySource:new(spec)
     ---@type FileRegistrySource
     local instance = {}
     setmetatable(instance, self)
+    instance.id = spec.id
     instance.spec = spec
     return instance
 end
@@ -68,10 +70,6 @@ function FileRegistrySource:get_all_package_names()
     return _.map(_.prop "name", self:get_all_package_specs())
 end
 
-function FileRegistrySource:get_installer()
-    return Optional.of(_.partial(self.install, self))
-end
-
 ---@async
 function FileRegistrySource:install()
     return Result.try(function(try)
diff --git a/lua/mason-registry/sources/github.lua b/lua/mason-registry/sources/github.lua
index b314d690..ceb503dc 100644
--- a/lua/mason-registry/sources/github.lua
+++ b/lua/mason-registry/sources/github.lua
@@ -15,7 +15,6 @@ local parse_checksums = _.compose(_.from_pairs, _.map(_.compose(_.reverse, _.spl
 
 ---@class GitHubRegistrySourceSpec
 ---@field id string
----@field repo string
 ---@field namespace string
 ---@field name string
 ---@field version string?
@@ -38,6 +37,7 @@ function GitHubRegistrySource:new(spec)
     local root_dir = InstallLocation.global():registry(path.concat { "github", spec.namespace, spec.name })
     instance.id = spec.id
     instance.spec = spec
+    instance.repo = ("%s/%s"):format(spec.namespace, spec.name)
     instance.root_dir = root_dir
     instance.data_file = path.concat { root_dir, "registry.json" }
     instance.info_file = path.concat { root_dir, "info.json" }
@@ -81,10 +81,6 @@ function GitHubRegistrySource:get_all_package_names()
     return _.map(_.prop "name", self:get_all_package_specs())
 end
 
-function GitHubRegistrySource:get_installer()
-    return Optional.of(_.partial(self.install, self))
-end
-
 ---@async
 function GitHubRegistrySource:install()
     local zzlib = require "mason-vendor.zzlib"
@@ -106,7 +102,7 @@ function GitHubRegistrySource:install()
             ---@type GitHubRelease
             local release = try(
                 providers.github
-                    .get_latest_release(self.spec.repo)
+                    .get_latest_release(self.repo)
                     :map_err(_.always "Failed to fetch latest registry version from GitHub API.")
             )
             version = release.tag_name
@@ -114,7 +110,7 @@ function GitHubRegistrySource:install()
         end
 
         local zip_file = path.concat { self.root_dir, "registry.json.zip" }
-        try(fetch(settings.current.github.download_url_template:format(self.spec.repo, version, "registry.json.zip"), {
+        try(fetch(settings.current.github.download_url_template:format(self.repo, version, "registry.json.zip"), {
             out_file = zip_file,
         }):map_err(_.always "Failed to download registry archive."))
         local zip_buffer = fs.async.read_file(zip_file)
@@ -126,7 +122,7 @@ function GitHubRegistrySource:install()
         pcall(fs.async.unlink, zip_file)
 
         local checksums = try(
-            fetch(settings.current.github.download_url_template:format(self.spec.repo, version, "checksums.txt")):map_err(
+            fetch(settings.current.github.download_url_template:format(self.repo, version, "checksums.txt")):map_err(
                 _.always "Failed to download checksums.txt."
             )
         )
@@ -158,17 +154,17 @@ end
 function GitHubRegistrySource:get_display_name()
     if self:is_installed() then
         local info = self:get_info()
-        return ("github.com/%s version: %s"):format(self.spec.repo, info.version)
+        return ("github.com/%s version: %s"):format(self.repo, info.version)
     else
-        return ("github.com/%s [uninstalled]"):format(self.spec.repo)
+        return ("github.com/%s [uninstalled]"):format(self.repo)
     end
 end
 
 function GitHubRegistrySource:__tostring()
     if self.spec.version then
-        return ("GitHubRegistrySource(repo=%s, version=%s)"):format(self.spec.repo, self.spec.version)
+        return ("GitHubRegistrySource(repo=%s, version=%s)"):format(self.repo, self.spec.version)
     else
-        return ("GitHubRegistrySource(repo=%s)"):format(self.spec.repo)
+        return ("GitHubRegistrySource(repo=%s)"):format(self.repo)
     end
 end
 
diff --git a/lua/mason-registry/sources/init.lua b/lua/mason-registry/sources/init.lua
index af1b9f6a..fd9cc573 100644
--- a/lua/mason-registry/sources/init.lua
+++ b/lua/mason-registry/sources/init.lua
@@ -1,6 +1,76 @@
-local _ = require "mason-core.functional"
+---@class RegistrySource
+---@field id string
+---@field get_package fun(self: RegistrySource, pkg_name: string): Package?
+---@field get_all_package_names fun(self: RegistrySource): string[]
+---@field get_all_package_specs fun(self: RegistrySource): RegistryPackageSpec[]
+---@field get_display_name fun(self: RegistrySource): string
+---@field is_installed fun(self: RegistrySource): boolean
+---@field install fun(self: RegistrySource): Result
 
-local M = {}
+---@alias RegistrySourceType '"github"' | '"lua"' | '"file"'
+
+---@class LazySource
+---@field type RegistrySourceType
+---@field id string
+---@field init fun(id: string): RegistrySource
+local LazySource = {}
+LazySource.__index = LazySource
+
+---@param id string
+function LazySource.GitHub(id)
+    local namespace, name = id:match "^(.+)/(.+)$"
+    if not namespace or not name then
+        error(("Failed to parse repository from GitHub registry: %q"):format(id), 0)
+    end
+    local name, version = unpack(vim.split(name, "@"))
+    local GitHubRegistrySource = require "mason-registry.sources.github"
+    return GitHubRegistrySource:new {
+        id = id,
+        namespace = namespace,
+        name = name,
+        version = version,
+    }
+end
+
+---@param id string
+function LazySource.Lua(id)
+    local LuaRegistrySource = require "mason-registry.sources.lua"
+    return LuaRegistrySource:new {
+        id = id,
+        mod = id,
+    }
+end
+
+---@param id string
+function LazySource.File(id)
+    local FileRegistrySource = require "mason-registry.sources.file"
+    return FileRegistrySource:new {
+        id = id,
+        path = id,
+    }
+end
+
+---@param type RegistrySourceType
+---@param id string
+---@param init fun(id: string): RegistrySource
+function LazySource:new(type, id, init)
+    local instance = setmetatable({}, self)
+    instance.type = type
+    instance.id = id
+    instance.init = init
+    return instance
+end
+
+function LazySource:get()
+    if not self.instance then
+        self.instance = self.init(self.id)
+    end
+    return self.instance
+end
+
+function LazySource:__tostring()
+    return ("LazySource(type=%s, id=%s)"):format(self.type, self.id)
+end
 
 ---@param str string
 local function split_once_left(str, char)
@@ -14,88 +84,46 @@ local function split_once_left(str, char)
 end
 
 ---@param registry_id string
----@return fun(): RegistrySource # Thunk to instantiate provider.
 local function parse(registry_id)
     local type, id = split_once_left(registry_id, ":")
+    assert(id, ("Malformed registry %q"):format(registry_id))
     if type == "github" then
-        local namespace, name = id:match "^(.+)/(.+)$"
-        if not namespace or not name then
-            error(("Failed to parse repository from GitHub registry: %q."):format(registry_id), 0)
-        end
-        local name, version = unpack(vim.split(name, "@"))
-        return function()
-            local GitHubRegistrySource = require "mason-registry.sources.github"
-            return GitHubRegistrySource:new {
-                id = registry_id,
-                repo = ("%s/%s"):format(namespace, name),
-                namespace = namespace,
-                name = name,
-                version = version,
-            }
-        end
+        return LazySource:new(type, id, LazySource.GitHub)
     elseif type == "lua" then
-        return function()
-            local LuaRegistrySource = require "mason-registry.sources.lua"
-            return LuaRegistrySource:new {
-                id = registry_id,
-                mod = id,
-            }
-        end
+        return LazySource:new(type, id, LazySource.Lua)
     elseif type == "file" then
-        return function()
-            local FileRegistrySource = require "mason-registry.sources.file"
-            return FileRegistrySource:new {
-                path = id,
-            }
-        end
-    elseif type ~= nil then
-        error(("Unknown registry type %q: %q."):format(type, registry_id), 0)
+        return LazySource:new(type, id, LazySource.File)
     end
-    error(("Malformed registry id: %q."):format(registry_id), 0)
+    error(("Unknown registry type: %s"):format(type))
 end
 
----@type ((fun(): RegistrySource) | RegistrySource)[]
-local registries = {}
+---@class LazySourceCollection
+---@field list LazySource[]
+local LazySourceCollection = {}
+LazySourceCollection.__index = LazySourceCollection
 
----@param registry_ids string[]
-function M.set_registries(registry_ids)
-    registries = {}
-    for _, registry in ipairs(registry_ids) do
-        local ok, err = pcall(function()
-            table.insert(registries, parse(registry))
-        end)
-        if not ok then
-            local log = require "mason-core.log"
-            local notify = require "mason-core.notify"
-            log.fmt_error("Failed to parse registry %q: %s", registry, err)
-            notify(err, vim.log.levels.ERROR)
-        end
-    end
+---@return LazySourceCollection
+function LazySourceCollection:new()
+    local instance = {}
+    setmetatable(instance, self)
+    instance.list = {}
+    return instance
 end
 
----@param opts? { include_uninstalled?: boolean }
-function M.iter(opts)
-    opts = opts or {}
-    local i = 1
-    return function()
-        while i <= #registries do
-            local registry = registries[i]
-            if type(registry) == "function" then
-                -- unwrap thunk
-                registry = registry()
-                registries[i] = registry
-            end
-            i = i + 1
-            if opts.include_uninstalled or registry:is_installed() then
-                return registry
-            end
-        end
-    end
+---@param registry string
+function LazySourceCollection:append(registry)
+    table.insert(self.list, parse(registry))
+    return self
+end
+
+---@param registry string
+function LazySourceCollection:prepend(registry)
+    table.insert(self.list, 1, parse(registry))
+    return self
 end
 
----@return boolean #Returns true if all sources are installed.
-function M.is_installed()
-    for source in M.iter { include_uninstalled = true } do
+function LazySourceCollection:is_all_installed()
+    for source in self:iterate { include_uninstalled = true } do
         if not source:is_installed() then
             return false
         end
@@ -103,15 +131,37 @@ function M.is_installed()
     return true
 end
 
----@return string # The sha256 checksum of the currently registered sources.
-function M.checksum()
+function LazySourceCollection:checksum()
     ---@type string[]
-    local registry_ids = {}
-    for source in M.iter { include_uninstalled = true } do
-        table.insert(registry_ids, source.id)
+    local registry_ids = vim.tbl_map(
+        ---@param source LazySource
+        function(source)
+            return source.id
+        end,
+        self.list
+    )
+    table.sort(registry_ids)
+    return vim.fn.sha256(table.concat(registry_ids, ""))
+end
+
+---@param opts? { include_uninstalled?: boolean }
+function LazySourceCollection:iterate(opts)
+    opts = opts or {}
+
+    local idx = 1
+    return function()
+        while idx <= #self.list do
+            local source = self.list[idx]:get()
+            idx = idx + 1
+            if opts.include_uninstalled or source:is_installed() then
+                return source
+            end
+        end
     end
-    local checksum = _.compose(vim.fn.sha256, _.join "", _.sort_by(_.identity))
-    return checksum(registry_ids)
 end
 
-return M
+function LazySourceCollection:__tostring()
+    return ("LazySourceCollection(list={%s})"):format(table.concat(vim.tbl_map(tostring, self.list), ", "))
+end
+
+return LazySourceCollection
diff --git a/lua/mason-registry/sources/lua.lua b/lua/mason-registry/sources/lua.lua
index fd88a49d..7fae54fe 100644
--- a/lua/mason-registry/sources/lua.lua
+++ b/lua/mason-registry/sources/lua.lua
@@ -1,4 +1,5 @@
 local Optional = require "mason-core.optional"
+local Result = require "mason-core.result"
 local _ = require "mason-core.functional"
 local log = require "mason-core.log"
 
@@ -35,6 +36,10 @@ function LuaRegistrySource:get_package(pkg_name)
     end
 end
 
+function LuaRegistrySource:install()
+    return Result.pcall(require, self.spec.mod)
+end
+
 ---@return string[]
 function LuaRegistrySource:get_all_package_names()
     local index = require(self.spec.mod)
@@ -53,11 +58,6 @@ function LuaRegistrySource:is_installed()
     return ok
 end
 
-function LuaRegistrySource:get_installer()
-    local Optional = require "mason-core.optional"
-    return Optional.empty()
-end
-
 function LuaRegistrySource:get_display_name()
     if self:is_installed() then
         return ("require(%q)"):format(self.spec.mod)