Merge pull request #6 from al-ce/main

Fix issue where remapped keys are not accurately picked up by DelayTrain

lua/delaytrain/init.lua

@@ -15,20 +15,22 @@ local keymaps = {
 
 local is_enabled = false
 
-function M.try_delay_keypress(key)
+local custom_mappings = {}
+
+function M.try_delay_keypress(key, keypress)
     current_interval = current_grace_period_intervals[key]
 
     -- ignore user defined patterns
     for _,ign_ft in ipairs(ignore_filetypes) do
         if vim.o.filetype:match(ign_ft) then
-            M.send_keypress(key)
+            M.send_keypress(keypress)
             return
         end
     end
 
     -- Ingore on macro execution
     if vim.fn.reg_executing() ~= "" then
-        M.send_keypress(key)
+        M.send_keypress(keypress)
         return
     end
 
@@ -42,13 +44,13 @@ function M.try_delay_keypress(key)
     -- Pass the key through only if we haven't reached the grace period
     if current_interval < vim.g.delaytrain_grace_period then
         current_grace_period_intervals[key] = current_interval + 1
-        M.send_keypress(key)
+        M.send_keypress(keypress)
     end
 end
 
-function M.send_keypress(key)
+function M.send_keypress(keypress)
     vim.api.nvim_feedkeys(
-        vim.api.nvim_replace_termcodes(key, true, false, true),
+        vim.api.nvim_replace_termcodes(keypress, true, false, true),
         'n', 
         false
     )
@@ -77,18 +79,56 @@ end
 function M.enable()
     is_enabled = true
 
+    -- Get an array of all the keys we want to delay regardless of remap status
+    local function get_key_mappings (key, mode_array)
+        local key_mappings = {}
+        for _, mode in ipairs(mode_array) do
+            local keypress = key
+            local remapped = vim.fn.maparg(key, mode, false, true).rhs
+            if remapped then
+                keypress = remapped
+            end
+
+            -- If keypress values differ across modes, add the new value here
+            if key_mappings[keypress] then
+                table.insert(key_mappings[keypress].modes, mode)
+            else
+                key_mappings[keypress] = {modes = {mode}, isremap = remapped ~= nil}
+            end
+        end
+
+      return key_mappings
+    end
+
+    -- Preserve old keymap so it can be restored after calling M.disable()
+    local function preserve_custom_mappings(lhs, rhs, modes)
+        if custom_mappings[lhs] then
+            table.insert(custom_mappings[lhs], {rhs = rhs, modes = modes})
+        else
+            custom_mappings[lhs] = {{rhs = rhs, modes = modes}}
+        end
+    end
+
     for modes, keys in pairs(keymaps) do
         mode_array = {}
         for mode in modes:gmatch"."  do
             table.insert(mode_array, mode)
         end
-        for _, key in ipairs(keys) do
-            -- Set the current grace period for the given key
-            current_grace_period_intervals[key] = 0
 
-            vim.keymap.set(mode_array, key, function() M.try_delay_keypress(key) end, {expr = true})
+        for _, key in ipairs(keys) do
+            local key_mappings = get_key_mappings(key, mode_array)
+            for keypress, key_data in pairs(key_mappings) do
+                if key_data.isremap then
+                  preserve_custom_mappings(key, keypress, key_data.modes)
+                end
+                -- Set the current grace period for the given key
+                current_grace_period_intervals[key] = 0
+                vim.keymap.set(key_data.modes, key, function() M.try_delay_keypress(key, keypress) end, {expr = true})
+            end
         end
     end
+
+
 end
 
 function M.disable()
@@ -101,6 +141,12 @@ function M.disable()
         end
         for _, key in ipairs(keys) do
             vim.keymap.del(mode_array, key)
+            -- Re-enable custom mappings
+            if custom_mappings[key] then
+                for _, map_values in pairs(custom_mappings[key]) do
+                  vim.keymap.set(map_values.modes, key, map_values.rhs)
+                end
+            end
         end
     end
 end