perf: fix performance regression

Author: itepechi

lua/hlchunk/mods/chunk/init.lua

@@ -114,23 +114,23 @@ function ChunkMod:get_chunk_data(range, virt_text_list, row_list, virt_text_win_
         local virt_text, virt_text_win_col =
             chunkHelper.calc(beg_virt_text, start_col, self.meta.leftcol, self.meta.shiftwidth)
         local char_list = utf8Split(virt_text)
-        vim.list_extend(virt_text_list, vim.fn.reverse(char_list))
-        vim.list_extend(row_list, vim.fn["repeat"]({ range.start }, #char_list))
-        vim.list_extend(
+        chunkHelper.list_extend(virt_text_list, chunkHelper.listReverse(char_list))
+        chunkHelper.list_extend(row_list, chunkHelper.repeated(range.start, #char_list))
+        chunkHelper.list_extend(
             virt_text_win_col_list,
-            vim.fn.reverse(chunkHelper.getColList(char_list, virt_text_win_col, self.meta.shiftwidth))
+            chunkHelper.listReverse(chunkHelper.getColList(char_list, virt_text_win_col, self.meta.shiftwidth))
         )
     end
 
     local mid_row_nums = range.finish - range.start - 1
-    vim.list_extend(row_list, rangeFromTo((range.start + 1), (range.finish - 1)))
-    vim.list_extend(virt_text_win_col_list, vim.fn["repeat"]({ start_col - self.meta.leftcol }, mid_row_nums))
+    chunkHelper.list_extend(row_list, rangeFromTo((range.start + 1), (range.finish - 1)))
+    chunkHelper.list_extend(virt_text_win_col_list, chunkHelper.repeated(start_col - self.meta.leftcol, mid_row_nums))
     ---@type string[]
     local chars
     if start_col - self.meta.leftcol < 0 then
-        chars = vim.fn["repeat"]({ "" }, mid_row_nums)
+        chars = chunkHelper.repeated("", mid_row_nums)
     else
-        chars = vim.fn["repeat"]({ self.conf.chars.vertical_line }, mid_row_nums)
+        chars = chunkHelper.repeated(self.conf.chars.vertical_line, mid_row_nums)
         -- when use click `<<` or `>>` to indent, we should make sure the line would not encounter the indent char
         for i = 1, mid_row_nums do
             local line = cFunc.get_line(range.bufnr, range.start + i)
@@ -144,7 +144,7 @@ function ChunkMod:get_chunk_data(range, virt_text_list, row_list, virt_text_win_
             end
         end
     end
-    vim.list_extend(virt_text_list, chars)
+    chunkHelper.list_extend(virt_text_list, chars)
 
     if end_blank_len > 0 then
         local virt_text_width = end_blank_len - start_col
@@ -183,9 +183,9 @@ function ChunkMod:get_chunk_data(range, virt_text_list, row_list, virt_text_win_
         local virt_text, virt_text_win_col =
             chunkHelper.calc(end_virt_text, start_col, self.meta.leftcol, self.meta.shiftwidth)
         local char_list = utf8Split(virt_text)
-        vim.list_extend(virt_text_list, char_list)
-        vim.list_extend(row_list, vim.fn["repeat"]({ range.finish }, #char_list))
-        vim.list_extend(
+        chunkHelper.list_extend(virt_text_list, char_list)
+        chunkHelper.list_extend(row_list, chunkHelper.repeated(range.finish, #char_list))
+        chunkHelper.list_extend(
             virt_text_win_col_list,
             chunkHelper.getColList(char_list, virt_text_win_col, self.meta.shiftwidth)
         )

lua/hlchunk/utils/chunkHelper.lua

@@ -91,7 +91,68 @@ local function get_chunk_range_by_treesitter(pos)
     return chunkHelper.CHUNK_RANGE_RET.NO_CHUNK, Scope(pos.bufnr, -1, -1)
 end
 
----@param opts? {pos: HlChunk.Pos, use_treesitter: boolean}
+---@param char string
+---@param shiftwidth integer
+---@return integer
+local function virt_text_char_width(char, shiftwidth)
+    local b1 = char:byte(1)
+    if b1 == 0x00 then
+        -- NULL is a terminator when used in virtual texts
+        return 0
+    elseif b1 == 0x09 then
+        return shiftwidth
+    elseif b1 <= 0x1F or b1 == 0x7F then
+        -- control chars other than NULL and TAB are two cells wide
+        return 2
+    elseif b1 <= 0x7F then
+        -- other ASCII chars are single cell wide
+        return 1
+    else
+        return vim.api.nvim_strwidth(char)
+    end
+end
+
+---faster alternative to `vim.fn.reverse()`
+---unlike the original, this only supports lists
+---@generic T
+---@param list T[]
+---@return T[]
+function chunkHelper.listReverse(list)
+    local dst = {}
+    for i, v in ipairs(list) do
+        dst[#list + 1 - i] = v
+    end
+    return dst
+end
+
+---faster alternative to `vim.fn.repeat()`
+---unlike the original, the input will be repeated as-is and the output will always be a list
+---@generic T
+---@param input T
+---@param count integer
+---@return T[]
+function chunkHelper.repeated(input, count)
+    local dst = {}
+    for i = 1, count do
+        dst[i] = input
+    end
+    return dst
+end
+
+---faster alternative to `vim.list_extend()` (mutates dst!)
+---unlike the original, this function lacks validation and range support
+---@generic T
+---@param dst T[]
+---@param src T[]
+---@return T[] dst
+function chunkHelper.list_extend(dst, src)
+    for i = 1, #src do
+        dst[#dst + 1] = src[i]
+    end
+    return dst
+end
+
+---@param opts? {pos: Pos, use_treesitter: boolean}
 ---@return CHUNK_RANGE_RETCODE enum
 ---@return HlChunk.Scope
 function chunkHelper.get_chunk_range(opts)
@@ -127,7 +188,7 @@ end
 function chunkHelper.utf8Split(inputstr)
     local list = {}
     for uchar in string.gmatch(inputstr, "[^\128-\191][\128-\191]*") do
-        table.insert(list, uchar)
+        list[#list + 1] = uchar
     end
     return list
 end
@@ -139,7 +200,7 @@ function chunkHelper.rangeFromTo(i, j, step)
     local t = {}
     step = step or 1
     for x = i, j, step do
-        table.insert(t, x)
+        t[#t + 1] = x
     end
     return t
 end
@@ -152,8 +213,8 @@ function chunkHelper.getColList(char_list, leftcol, shiftwidth)
     local t = {}
     local next_col = leftcol
     for i = 1, #char_list do
-        table.insert(t, next_col)
-        next_col = next_col + chunkHelper.virtTextStrWidth(char_list[i], shiftwidth)
+        t[#t + 1] = next_col
+        next_col = next_col + virt_text_char_width(char_list[i], shiftwidth)
     end
     return t
 end
@@ -183,7 +244,7 @@ function chunkHelper.repeatToWidth(str, width, shiftwidth)
     local current_width = str_width * repeatable_len
     local i = 1
     while i <= #chars do
-        local char_width = chunkHelper.virtTextStrWidth(chars[i], shiftwidth)
+        local char_width = virt_text_char_width(chars[i], shiftwidth)
         ---assumed to be an out-of-bounds char (like in nerd fonts) followed by a whitespace if true
         local likely_oob_char =
             -- single-cell
@@ -235,18 +296,13 @@ end
 ---@return boolean
 function chunkHelper.checkCellsBlank(line, start_col, end_col, shiftwidth)
     local current_col = 1
-    local current_byte = 1
     local current_char = 1
-    while current_byte <= #line and current_col <= end_col do
-        local final_byte = vim.str_byteindex(line, current_char)
-        local char = line:sub(current_byte, final_byte)
+    local chars = chunkHelper.utf8Split(line)
+    while current_char <= #chars and current_col <= end_col do
+        local char = chars[current_char]
         local b1, b2, b3 = char:byte(1, 3)
-        if char == "" then
-            break
-        end
         ---@type integer
         local next_col
-        local next_byte = final_byte + 1
         local next_char = current_char + 1
         if char == " " then
             next_col = current_col + 1
@@ -260,13 +316,11 @@ function chunkHelper.checkCellsBlank(line, start_col, end_col, shiftwidth)
             next_col = current_col + 1
         else
             local char_width = vim.api.nvim_strwidth(char)
-            local next_byte_peek = line:byte(final_byte + 1)
-            if char_width == 1 and next_byte_peek == 0x20 then
+            if char_width == 1 and chars[current_char + 1] == " " then
                 -- the char is assumed to be an out-of-bounds char (like in nerd fonts)
                 -- followed by a whitespace
                 next_col = current_col + 2
                 -- skip the whitespace part of out-of-bounds char + " "
-                next_byte = next_byte + 1
                 next_char = next_char + 1
             else
                 next_col = current_col + char_width
@@ -338,7 +392,6 @@ function chunkHelper.checkCellsBlank(line, start_col, end_col, shiftwidth)
             return false
         end
         current_col = next_col
-        current_byte = next_byte
         current_char = next_char
     end
     return true
@@ -351,25 +404,11 @@ end
 function chunkHelper.virtTextStrWidth(str, shiftwidth, stop_on_null)
     local current_width = 0
     for _, char in ipairs(chunkHelper.utf8Split(str)) do
-        if char == "\0" then
-            if stop_on_null then
-                return current_width
-            end
-            -- just ignore otherwise
-        elseif char == "\t" then
-            current_width = current_width + shiftwidth
-        else
-            local b1 = char:byte(1)
-            if b1 <= 0x1F or b1 == 0x7F then
-                -- control chars other than NULL and TAB are two cells wide
-                current_width = current_width + 2
-            elseif b1 <= 0x7F then
-                -- other ASCII chars are single cell wide
-                current_width = current_width + 1
-            else
-                current_width = current_width + vim.api.nvim_strwidth(char)
-            end
+        if stop_on_null and char == "\0" then
+            -- NULL is a terminator when used in virtual texts
+            return current_width
         end
+        current_width = current_width + virt_text_char_width(char, shiftwidth)
     end
     return current_width
 end

test/features/chunkHelper_spec.lua

@@ -134,6 +134,32 @@ describe("indentHelper", function()
         end
     end)
 
+    it("listReverse happy path", function()
+        local inputList = {
+            { t = {}, res = {} },
+            { t = { 1 }, res = { 1 } },
+            { t = { 1, 2, 3 }, res = { 3, 2, 1 } },
+            { t = { 1, 2, 3, 4 }, res = { 4, 3, 2, 1 } },
+        }
+
+        for _, testCase in ipairs(inputList) do
+            local res = chunkHelper.listReverse(testCase.t)
+            assert.same(res, testCase.res)
+        end
+    end)
+
+    it("repeated happy path", function()
+        local inputList = {
+            { input = 1, repeat_to = 1, res = { 1 } },
+            { input = 1, repeat_to = 3, res = { 1, 1, 1 } },
+        }
+
+        for _, testCase in ipairs(inputList) do
+            local res = chunkHelper.repeated(testCase.input, testCase.repeat_to)
+            assert.same(res, testCase.res)
+        end
+    end)
+
     it("checkCellsBlank happy path", function()
         -- bunch of edge cases
         local inputList = {
@@ -218,4 +244,17 @@ describe("indentHelper", function()
             assert.same(res, testCase.res)
         end
     end)
+
+    it("list_extend happy path", function()
+        local inputList = {
+            { dst = { 1, 2, 3 }, src = {}, res = { 1, 2, 3 } },
+            { dst = {}, src = { 4, 5, 6 }, res = { 4, 5, 6 } },
+            { dst = { 1, 2, 3 }, src = { 4, 5, 6 }, res = { 1, 2, 3, 4, 5, 6 } },
+        }
+
+        for _, testCase in ipairs(inputList) do
+            chunkHelper.list_extend(testCase.dst, testCase.src)
+            assert.same(testCase.dst, testCase.res)
+        end
+    end)
 end)