From 2feea7e6b2e27afd9b5571dd60baccc75ea9b160 Mon Sep 17 00:00:00 2001
From: Tristan Knight <admin@snappeh.com>
Date: Fri, 17 May 2024 15:27:22 +0100
Subject: [PATCH] fix: e motion (#21)

---
 lua/precognition/horizontal_motions.lua       | 21 ++++++++++++-------
 .../precognition/horizontal_motions_spec.lua  | 20 ++++++++++++------
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/lua/precognition/horizontal_motions.lua b/lua/precognition/horizontal_motions.lua
index 4782af3..0abaf64 100644
--- a/lua/precognition/horizontal_motions.lua
+++ b/lua/precognition/horizontal_motions.lua
@@ -48,11 +48,10 @@ end
 
 ---@param str string
 ---@param cursorcol integer
----@param _linelen integer
+---@param linelen integer
 ---@return Precognition.PlaceLoc
-function M.end_of_word(str, cursorcol, _linelen)
-    local len = vim.fn.strcharlen(str)
-    if cursorcol >= len then
+function M.end_of_word(str, cursorcol, linelen)
+    if cursorcol >= linelen then
         return 0
     end
     local offset = cursorcol
@@ -69,20 +68,26 @@ function M.end_of_word(str, cursorcol, _linelen)
     end
 
     if c_class ~= 0 and next_char_class ~= 0 then
-        while utils.char_class(char) == c_class and offset <= len do
+        while utils.char_class(char) == c_class and offset <= linelen do
             offset = offset + 1
             char = vim.fn.strcharpart(str, offset - 1, 1)
         end
     end
 
     if c_class == 0 or next_char_class == 0 then
-        local next_word_start = M.next_word_boundary(str, offset, 0)
+        local next_word_start = M.next_word_boundary(str, cursorcol, linelen)
         if next_word_start then
-            rev_offset = M.end_of_word(str, next_word_start + 1, 0)
+            next_char_class = utils.char_class(vim.fn.strcharpart(str, (next_word_start - 1) + 1, 1))
+            --next word is single char
+            if next_char_class == 0 then
+                rev_offset = next_word_start
+            else
+                rev_offset = M.end_of_word(str, next_word_start, linelen)
+            end
         end
     end
 
-    if rev_offset ~= nil and rev_offset <= 0 then
+    if rev_offset and rev_offset <= 0 then
         return 0
     end
 
diff --git a/tests/precognition/horizontal_motions_spec.lua b/tests/precognition/horizontal_motions_spec.lua
index 2ac76eb..ef12043 100644
--- a/tests/precognition/horizontal_motions_spec.lua
+++ b/tests/precognition/horizontal_motions_spec.lua
@@ -121,12 +121,12 @@ describe("boundaries", function()
             eq(3, hm.end_of_word("abc efg", 2, 7))
             eq(7, hm.end_of_word("abc efg", 3, 7))
 
-            eq(7, hm.end_of_word("slighly more complex test", 1, 22))
-            eq(7, hm.end_of_word("slighly more complex test", 2, 22))
-            eq(12, hm.end_of_word("slighly more complex test", 10, 22))
-            eq(20, hm.end_of_word("slighly more complex test", 13, 22))
-            eq(20, hm.end_of_word("slighly more complex test", 15, 22))
-            eq(25, hm.end_of_word("slighly more complex test", 21, 22))
+            eq(7, hm.end_of_word("slighly more complex test", 1, 25))
+            eq(7, hm.end_of_word("slighly more complex test", 2, 25))
+            eq(12, hm.end_of_word("slighly more complex test", 10, 25))
+            eq(20, hm.end_of_word("slighly more complex test", 13, 25))
+            eq(20, hm.end_of_word("slighly more complex test", 15, 25))
+            eq(25, hm.end_of_word("slighly more complex test", 21, 25))
 
             eq(14, hm.end_of_word("    myFunction(example, stuff)", 1, 30))
             eq(14, hm.end_of_word("    myFunction(example, stuff)", 2, 30))
@@ -175,4 +175,12 @@ describe("edge case", function()
         local str = "there is a space "
         eq(0, hm.end_of_word(str, 16, #str))
     end)
+
+    it("single character next word ends", function()
+        local str = "show_something = true,"
+        eq(14, hm.end_of_word(str, 1, #str))
+        eq(16, hm.end_of_word(str, 14, #str))
+        eq(16, hm.end_of_word(str, 15, #str))
+        eq(22, hm.end_of_word(str, 21, #str))
+    end)
 end)