Merge pull request #16124 from mjlbach/feat/bjorn-bait

feat: add vim.str_utf_pos

2 files changed, 32 insertions, 0 deletions

diff --git a/src/nvim/lua/executor.c b/src/nvim/lua/executor.c
index 59ebafd9c8..6e095ce7fe 100644
--- a/src/nvim/lua/executor.c
+++ b/src/nvim/lua/executor.c
@@ -176,6 +176,29 @@ static int nlua_str_utfindex(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   return 2;
 }
 
+/// return byte indices of codepoints in a string (only supports utf-8 currently).
+///
+/// Expects a string.
+///
+/// Returns a list of codepoints.
+static int nlua_str_utf_pos(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
+{
+  size_t s1_len;
+  const char *s1 = luaL_checklstring(lstate, 1, &s1_len);
+  lua_newtable(lstate);
+
+  size_t idx = 1;
+  size_t clen;
+  for (size_t i = 0; i < s1_len && s1[i] != NUL; i += clen) {
+    clen = (size_t)utf_ptr2len_len((const char_u *)(s1)+i, (int)(s1_len-i));
+    lua_pushinteger(lstate, (long)i + 1);
+    lua_rawseti(lstate, -2, (int)idx);
+    idx++;
+  }
+
+  return 1;
+}
+
 /// convert UTF-32 or UTF-16 indices to byte index.
 ///
 /// Expects up to three args: string, index and use_utf16.
@@ -435,6 +458,9 @@ static int nlua_state_init(lua_State *const lstate) FUNC_ATTR_NONNULL_ALL
   // str_byteindex
   lua_pushcfunction(lstate, &nlua_str_byteindex);
   lua_setfield(lstate, -2, "str_byteindex");
+  // str_utf_pos
+  lua_pushcfunction(lstate, &nlua_str_utf_pos);
+  lua_setfield(lstate, -2, "str_utf_pos");
   // neovim version
   lua_pushcfunction(lstate, &nlua_nvim_version);
   lua_setfield(lstate, -2, "version");
diff --git a/test/functional/lua/vim_spec.lua b/test/functional/lua/vim_spec.lua
index 3123063b8c..8c691b1121 100644
--- a/test/functional/lua/vim_spec.lua
+++ b/test/functional/lua/vim_spec.lua
@@ -179,6 +179,12 @@ describe('lua stdlib', function()
     end
   end)
 
+  it("vim.str_utf_pos", function()
+    exec_lua([[_G.test_text = "xy åäö ɧ 汉语 ↥ 🤦x🦄 å بِيَّ"]])
+    local expected_positions = { 1,2,3,4,6,8,10,11,13,14,17,20,21,24,25,29,30,34,35,36,38,39,41,43,45,47 }
+    eq(expected_positions, exec_lua("return vim.str_utf_pos(_G.test_text)"))
+  end)
+
   it("vim.schedule", function()
     exec_lua([[
       test_table = {}