feat(perf): pre-compute conceal information a single time per parse cycle

## Details

Our current logic for figuring out what is concealed relies on calling the
treesitter get_captures_at_pos API repeatedly. Under the hood this call
is doing lots of repeated work iterating over the same queries.

Instead store the concealed nodes for a buffer in some table. Then use
this much smaller subset of nodes when computing concealed widths.

Depending on the concents of a buffer this improves performance
somewhere between 3x & 10x, nice!

Author: MeanderingProgrammer

benches/medium_spec.lua

@@ -39,8 +39,8 @@ describe('medium.md', function()
         end)
         eq(3000, #util.get_actual_marks())
 
-        truthy(start_time > 400 and start_time < 500, 'expected start time (400, 500)')
-        truthy(refresh_time > 400 and refresh_time < 500, 'expected refresh time (400, 500)')
+        truthy(start_time > 25 and start_time < 75, 'expected start time (25, 75)')
+        truthy(refresh_time > 25 and refresh_time < 50, 'expected refresh time (25, 50)')
         truthy(move_time > 1 and move_time < 5, 'expected move time (1, 5)')
     end)
 end)

doc/todo.md

@@ -2,7 +2,26 @@
 
 - See if there is a stable way to align table cells according to delimiter
   alignment info.
-- Add a performance test suite to measure improvements, relative value on
-  same hardware should still be useful.
 - Figure out how to display the many configuration options & impact
 - Potentially change LuaRocks icon dependency to [mini.icons](https://luarocks.org/modules/neorocks/mini.icons)
+- Force non-markdown handlers to run first in `ui.lua`. Cache information
+  at buffer level from inline parser to avoid computing it:
+
+```lua
+-- Parse marks
+local marks = {}
+-- Parse markdown after all other nodes to take advantage of state
+local markdown_roots = {}
+parser:for_each_tree(function(tree, language_tree)
+    local language = language_tree:lang()
+    if language == 'markdown' then
+        table.insert(markdown_roots, tree:root())
+    else
+        vim.list_extend(marks, M.parse_tree(buf, language, tree:root()))
+    end
+end)
+for _, root in ipairs(markdown_roots) do
+    vim.list_extend(marks, M.parse_tree(buf, 'markdown', root))
+end
+return marks
+```

lua/render-markdown/colors.lua

@@ -1,8 +1,5 @@
----@class render.md.ColorCache
-local cache = {
-    ---@type string[]
-    highlights = {},
-}
+---@type string[]
+local cache = {}
 
 ---@class render.md.Colors
 local M = {}
@@ -65,7 +62,7 @@ end
 ---@return string
 function M.combine(foreground, background)
     local name = string.format('%s_%s_%s', M.prefix, foreground, background)
-    if not vim.tbl_contains(cache.highlights, name) then
+    if not vim.tbl_contains(cache, name) then
         local fg = M.get_hl(foreground)
         local bg = M.get_hl(background)
         vim.api.nvim_set_hl(0, name, {
@@ -76,7 +73,7 @@ function M.combine(foreground, background)
             ---@diagnostic disable-next-line: undefined-field
             ctermbg = bg.ctermbg,
         })
-        table.insert(cache.highlights, name)
+        table.insert(cache, name)
     end
     return name
 end
@@ -85,7 +82,7 @@ end
 ---@return string
 function M.inverse(highlight)
     local name = string.format('%s_Inverse_%s', M.prefix, highlight)
-    if not vim.tbl_contains(cache.highlights, name) then
+    if not vim.tbl_contains(cache, name) then
         local hl = M.get_hl(highlight)
         vim.api.nvim_set_hl(0, name, {
             fg = hl.bg,
@@ -95,7 +92,7 @@ function M.inverse(highlight)
             ---@diagnostic disable-next-line: undefined-field
             ctermfg = hl.ctermbg,
         })
-        table.insert(cache.highlights, name)
+        table.insert(cache, name)
     end
     return name
 end

lua/render-markdown/handler/latex.lua

@@ -2,11 +2,8 @@ local logger = require('render-markdown.logger')
 local state = require('render-markdown.state')
 local ts = require('render-markdown.ts')
 
----@class render.md.LatexCache
-local cache = {
-    ---@type table<string, string[]>
-    expressions = {},
-}
+---@type table<string, string[]>
+local cache = {}
 
 ---@class render.md.handler.Latex: render.md.Handler
 local M = {}
@@ -27,7 +24,7 @@ function M.parse(root, buf)
     local info = ts.info(root, buf)
     logger.debug_node_info('latex', info)
 
-    local expressions = cache.expressions[info.text]
+    local expressions = cache[info.text]
     if expressions == nil then
         local raw_expression = vim.fn.system(latex.converter, info.text)
         expressions = vim.split(raw_expression, '\n', { plain = true, trimempty = true })
@@ -37,7 +34,7 @@ function M.parse(root, buf)
         for _ = 1, latex.bottom_pad do
             table.insert(expressions, '')
         end
-        cache.expressions[info.text] = expressions
+        cache[info.text] = expressions
     end
 
     local latex_lines = vim.tbl_map(function(expression)

lua/render-markdown/manager.lua

@@ -2,11 +2,8 @@ local state = require('render-markdown.state')
 local ui = require('render-markdown.ui')
 local util = require('render-markdown.util')
 
----@class render.md.manager.Data
-local data = {
-    ---@type integer[]
-    buffers = {},
-}
+---@type integer[]
+local buffers = {}
 
 ---@class render.md.Manager
 local M = {}
@@ -30,7 +27,7 @@ function M.setup()
         callback = function()
             for _, win in ipairs(vim.v.event.windows) do
                 local buf = util.win_to_buf(win)
-                if vim.tbl_contains(data.buffers, buf) then
+                if vim.tbl_contains(buffers, buf) then
                     ui.schedule_render(buf, true)
                 end
             end
@@ -43,7 +40,7 @@ function M.set_all(enabled)
     -- Attempt to attach current buffer in case this is from a lazy load
     M.attach(vim.api.nvim_get_current_buf())
     state.enabled = enabled
-    for _, buf in ipairs(data.buffers) do
+    for _, buf in ipairs(buffers) do
         ui.schedule_render(buf, true)
     end
 end
@@ -61,10 +58,10 @@ function M.attach(buf)
     if util.file_size_mb(buf) > config.max_file_size then
         return
     end
-    if vim.tbl_contains(data.buffers, buf) then
+    if vim.tbl_contains(buffers, buf) then
         return
     end
-    table.insert(data.buffers, buf)
+    table.insert(buffers, buf)
     -- Events that do not imply modifications to buffer so can avoid re-parsing
     -- This relies on the ui parsing the buffer anyway if it is the first time it is seen
     vim.api.nvim_create_autocmd({ 'BufWinEnter', 'BufLeave' }, {

lua/render-markdown/request.lua

@@ -0,0 +1,55 @@
+local util = require('render-markdown.util')
+
+---@class render.md.RequestCache
+local cache = {
+    ---@type table<integer, TSNode[]>
+    conceal = {},
+}
+
+---@class render.md.Request
+local M = {}
+
+---@param buf integer
+function M.reset_buf(buf)
+    cache.conceal[buf] = {}
+end
+
+---@param buf integer
+---@param parser vim.treesitter.LanguageTree
+function M.compute_conceal(buf, parser)
+    if util.get_win(util.buf_to_win(buf), 'conceallevel') == 0 then
+        cache.conceal[buf] = {}
+        return
+    end
+
+    local nodes = {}
+    parser:for_each_tree(function(tree, language_tree)
+        local language = language_tree:lang()
+        if vim.tbl_contains({ 'markdown', 'markdown_inline' }, language) then
+            local query = vim.treesitter.query.get(language, 'highlights')
+            if query ~= nil then
+                for _, node, metadata in query:iter_captures(tree:root(), buf, 0, -1) do
+                    if metadata.conceal ~= nil then
+                        table.insert(nodes, node)
+                    end
+                end
+            end
+        end
+    end)
+    cache.conceal[buf] = nodes
+end
+
+---@param buf integer
+---@param row integer
+---@param col integer
+---@return boolean
+function M.concealed(buf, row, col)
+    for _, node in ipairs(cache.conceal[buf]) do
+        if vim.treesitter.is_in_node_range(node, row, col) then
+            return true
+        end
+    end
+    return false
+end
+
+return M

lua/render-markdown/state.lua

@@ -1,10 +1,7 @@
 local util = require('render-markdown.util')
 
----@class render.md.StateCache
-local cache = {
-    ---@type table<integer, render.md.BufferConfig>
-    configs = {},
-}
+---@type table<integer, render.md.BufferConfig>
+local configs = {}
 
 ---@class render.md.State
 ---@field private config render.md.Config
@@ -24,7 +21,7 @@ local M = {}
 ---@param user_config? render.md.UserConfig
 function M.setup(default_config, user_config)
     -- Reset cache to pickup config changes
-    cache.configs = {}
+    configs = {}
 
     -- Create top level config from default & user
     local config = vim.deepcopy(vim.tbl_deep_extend('force', default_config, user_config or {}), true)
@@ -46,7 +43,7 @@ end
 ---@param buf integer
 ---@return render.md.BufferConfig
 M.get_config = function(buf)
-    local config = cache.configs[buf]
+    local config = configs[buf]
     if config == nil then
         local buftype = util.get_buf(buf, 'buftype')
         local buftype_config = M.config.overrides.buftype[buftype]
@@ -55,7 +52,7 @@ M.get_config = function(buf)
         else
             config = M.default_buffer_config()
         end
-        cache.configs[buf] = config
+        configs[buf] = config
     end
     return config
 end

lua/render-markdown/ts.lua

@@ -1,5 +1,5 @@
+local request = require('render-markdown.request')
 local str = require('render-markdown.str')
-local util = require('render-markdown.util')
 
 ---@param node TSNode
 ---@return boolean
@@ -110,18 +110,12 @@ end
 ---@param info render.md.NodeInfo
 ---@return integer
 function M.concealed(buf, info)
-    if util.get_win(util.buf_to_win(buf), 'conceallevel') == 0 then
-        return 0
-    end
     local result = 0
     local col = info.start_col
     for _, index in ipairs(vim.fn.str2list(info.text)) do
         local ch = vim.fn.nr2char(index)
-        local captures = vim.treesitter.get_captures_at_pos(buf, info.start_row, col)
-        for _, capture in ipairs(captures) do
-            if capture.metadata.conceal ~= nil then
-                result = result + str.width(ch)
-            end
+        if request.concealed(buf, info.start_row, col) then
+            result = result + str.width(ch)
         end
         col = col + #ch
     end

lua/render-markdown/ui.lua

@@ -1,4 +1,5 @@
 local logger = require('render-markdown.logger')
+local request = require('render-markdown.request')
 local state = require('render-markdown.state')
 local util = require('render-markdown.util')
 
@@ -9,11 +10,8 @@ local builtin_handlers = {
     latex = require('render-markdown.handler.latex'),
 }
 
----@class render.md.UiCache
-local cache = {
-    ---@type table<integer, render.md.Mark[]>
-    marks = {},
-}
+---@type table<integer, render.md.Mark[]>
+local cache = {}
 
 ---@class render.md.Ui
 local M = {}
@@ -22,7 +20,7 @@ local M = {}
 M.namespace = vim.api.nvim_create_namespace('render-markdown.nvim')
 
 function M.invalidate_cache()
-    cache.marks = {}
+    cache = {}
 end
 
 ---@param buf integer
@@ -59,13 +57,13 @@ function M.render(buf, mode, parse)
         end
 
         -- Re-compute marks, needed if missing or between text changes
-        local marks = cache.marks[buf]
+        local marks = cache[buf]
         local parsed = marks == nil or parse
         if parsed then
             logger.start()
             marks = M.parse_buffer(buf)
             logger.flush()
-            cache.marks[buf] = marks
+            cache[buf] = marks
         end
 
         local row = util.cursor_row(buf)
@@ -110,6 +108,9 @@ function M.parse_buffer(buf)
     if not parser:is_valid() then
         parser:parse(true)
     end
+    -- Pre-compute conceal information after reseting buffer cache
+    request.reset_buf(buf)
+    request.compute_conceal(buf, parser)
     -- Parse marks
     local marks = {}
     parser:for_each_tree(function(tree, language_tree)