From b79099cc9d768241162bb45d284d6a243736b9fb Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Sun, 8 Oct 2023 10:11:25 +0200
Subject: [PATCH] feat(keys): refactor code and allow disabling keymaps per
 mode. mode no longer needs to be exactly the same in order to disable.

---
 lua/lazy/core/handler/init.lua |  3 ++
 lua/lazy/core/handler/keys.lua | 96 ++++++++++++++++++++++------------
 lua/lazy/types.lua             |  2 +-
 lua/lazy/view/render.lua       |  2 +-
 4 files changed, 68 insertions(+), 35 deletions(-)

diff --git a/lua/lazy/core/handler/init.lua b/lua/lazy/core/handler/init.lua
index 9c6b92d0..9a3fd5a5 100644
--- a/lua/lazy/core/handler/init.lua
+++ b/lua/lazy/core/handler/init.lua
@@ -5,6 +5,7 @@ local Config = require("lazy.core.config")
 ---@field type LazyHandlerTypes
 ---@field extends? LazyHandler
 ---@field active table<string,table<string,string>>
+---@field managed table<string,string>
 ---@field super LazyHandler
 local M = {}
 
@@ -64,6 +65,7 @@ function M.new(type)
   local self = setmetatable({}, { __index = setmetatable(handler, { __index = super }) })
   self.super = super
   self.active = {}
+  self.managed = {}
   self.type = type
   return self
 end
@@ -94,6 +96,7 @@ function M:add(plugin)
       self.active[key] = {}
       self:_add(value)
     end
+    self.managed[key] = self.managed[key]
     self.active[key][plugin.name] = plugin.name
   end
 end
diff --git a/lua/lazy/core/handler/keys.lua b/lua/lazy/core/handler/keys.lua
index 35d7fb06..92531027 100644
--- a/lua/lazy/core/handler/keys.lua
+++ b/lua/lazy/core/handler/keys.lua
@@ -1,60 +1,89 @@
 local Util = require("lazy.core.util")
 local Loader = require("lazy.core.loader")
 
----@class LazyKeys
----@field [1] string lhs
----@field [2]? string|fun()|false rhs
+---@class LazyKeysBase
 ---@field desc? string
----@field mode? string|string[]
 ---@field noremap? boolean
 ---@field remap? boolean
 ---@field expr? boolean
+---@field nowait? boolean
 ---@field ft? string|string[]
+
+---@class LazyKeysSpec: LazyKeysBase
+---@field [1] string lhs
+---@field [2]? string|fun()|false rhs
+---@field mode? string|string[]
+
+---@class LazyKeys: LazyKeysBase
+---@field lhs string lhs
+---@field rhs? string|fun() rhs
+---@field mode? string
 ---@field id string
 
 ---@class LazyKeysHandler:LazyHandler
 local M = {}
 
----@param value string|LazyKeys
-function M.parse(value)
-  local ret = vim.deepcopy(value)
-  ret = type(ret) == "string" and { ret } or ret --[[@as LazyKeys]]
-  ret.mode = ret.mode or "n"
-  ret.id = vim.api.nvim_replace_termcodes(ret[1] or "", true, true, true)
-  if ret.mode then
-    local mode = ret.mode
-    if type(mode) == "table" then
-      ---@cast mode string[]
-      table.sort(mode)
-      mode = table.concat(mode, ", ")
-    end
-    if mode ~= "n" then
-      ret.id = ret.id .. " (" .. mode .. ")"
-    end
+---@param value string|LazyKeysSpec
+---@param mode? string
+---@return LazyKeys
+function M.parse(value, mode)
+  value = type(value) == "string" and { value } or value --[[@as LazyKeysSpec]]
+  local ret = vim.deepcopy(value) --[[@as LazyKeys]]
+  ret.lhs = ret[1] or ""
+  ret.rhs = ret[2]
+  ---@diagnostic disable-next-line: no-unknown
+  ret[1] = nil
+  ---@diagnostic disable-next-line: no-unknown
+  ret[2] = nil
+  ret.mode = mode or "n"
+  ret.id = vim.api.nvim_replace_termcodes(ret.lhs, true, true, true)
+  if ret.mode ~= "n" then
+    ret.id = ret.id .. " (" .. ret.mode .. ")"
   end
   return ret
 end
 
+---@param lhs string
+---@param mode? string
+function M:have(lhs, mode)
+  local keys = M.parse(lhs, mode)
+  return self.managed[keys.id]
+end
+
 ---@param plugin LazyPlugin
 function M:values(plugin)
-  ---@type table<string,any>
+  return M.resolve(plugin.keys)
+end
+
+---@param spec? (string|LazyKeysSpec)[]
+function M.resolve(spec)
+  ---@type LazyKeys[]
   local values = {}
   ---@diagnostic disable-next-line: no-unknown
-  for _, value in ipairs(plugin[self.type] or {}) do
-    local keys = M.parse(value)
-    if keys[2] == vim.NIL or keys[2] == false then
-      values[keys.id] = nil
-    else
-      values[keys.id] = keys
+  for _, value in ipairs(spec or {}) do
+    value = type(value) == "string" and { value } or value --[[@as LazyKeysSpec]]
+    value.mode = value.mode or "n"
+    local modes = (type(value.mode) == "table" and value.mode or { value.mode }) --[=[@as string[]]=]
+    for _, mode in ipairs(modes) do
+      local keys = M.parse(value, mode)
+      if keys.rhs == vim.NIL or keys.rhs == false then
+        values[keys.id] = nil
+      else
+        values[keys.id] = keys
+      end
     end
   end
   return values
 end
 
+---@param keys LazyKeys
 function M.opts(keys)
-  local opts = {}
+  local skip = { mode = true, id = true, ft = true, rhs = true, lhs = true }
+  local opts = {} ---@type LazyKeysBase
+  ---@diagnostic disable-next-line: no-unknown
   for k, v in pairs(keys) do
-    if type(k) ~= "number" and k ~= "mode" and k ~= "id" and k ~= "ft" then
+    if type(k) ~= "number" and not skip[k] then
+      ---@diagnostic disable-next-line: no-unknown
       opts[k] = v
     end
   end
@@ -63,7 +92,7 @@ end
 
 ---@param keys LazyKeys
 function M:_add(keys)
-  local lhs = keys[1]
+  local lhs = keys.lhs
   local opts = M.opts(keys)
 
   ---@param buf? number
@@ -121,7 +150,7 @@ end
 -- mapping when needed
 ---@param keys LazyKeys
 function M:_del(keys)
-  pcall(vim.keymap.del, keys.mode, keys[1])
+  pcall(vim.keymap.del, keys.mode, keys.lhs)
   -- make sure to create global mappings when needed
   -- buffer-local mappings are managed by lazy
   if not keys.ft then
@@ -133,10 +162,11 @@ end
 ---@param keys LazyKeys
 ---@param buf number?
 function M:_set(keys, buf)
-  if keys[2] then
+  if keys.rhs then
     local opts = M.opts(keys)
+    ---@diagnostic disable-next-line: inject-field
     opts.buffer = buf
-    vim.keymap.set(keys.mode, keys[1], keys[2], opts)
+    vim.keymap.set(keys.mode, keys.lhs, keys.rhs, opts)
   end
 end
 
diff --git a/lua/lazy/types.lua b/lua/lazy/types.lua
index 30363b60..3e7f74db 100644
--- a/lua/lazy/types.lua
+++ b/lua/lazy/types.lua
@@ -32,7 +32,7 @@
 ---@field event? string[]
 ---@field cmd? string[]
 ---@field ft? string[]
----@field keys? (string|LazyKeys)[]
+---@field keys? (string|LazyKeysSpec)[]
 ---@field module? false
 
 ---@class LazyPluginRef
diff --git a/lua/lazy/view/render.lua b/lua/lazy/view/render.lua
index 5b1f9afa..319d627b 100644
--- a/lua/lazy/view/render.lua
+++ b/lua/lazy/view/render.lua
@@ -331,7 +331,7 @@ function M:reason(reason, opts)
         value = value:match("User (.*)") or value
       end
       if key == "keys" then
-        value = type(value) == "string" and value or value[1]
+        value = type(value) == "string" and value or value.lhs
       end
       local hl = "LazyReason" .. key:sub(1, 1):upper() .. key:sub(2)
       local icon = Config.options.ui.icons[key]