perf(renderer): cache absolute folds in ctx.part_folds for incremental updates

Streaming flushes touch at most a few parts at a time, but
set_all_folds walks ctx.formatted_parts and recomputes every part's
folds every call. For an append-heavy tool output that's
quadratic in part count.

Add ctx.part_folds as a self-checked cache of absolute fold ranges,
keyed by part_id. Cache entries store line_start so rebuild can
detect shift_all invalidation without external notification; the
rebuild also drops entries whose source part is gone from
ctx.formatted_parts, covering removal paths. update_part_folds(id)
becomes the streaming hot path: it computes or reuses the target
part's entry, then rebuild scans the cache (cheap) and dispatches
the sorted fold list. set_all_folds remains the cold path for
bulk mode and explicit resets.

Microbench on N=20 parts, single-part update:
  set_all_folds:      4.05 us
  update_part_folds:  0.30 us   (13.5x)
On N=100 hot-part streaming, 67x.

The bug the previous incarnation had (stale fold ranges surviving
in the cache after part removal, eventually fed to foldopen! and
raising E16) is structurally closed: rebuild reads ctx.formatted_parts
as the source of truth and drops any cache entry without one. Tested
explicitly in renderer_buffer_spec.

Author: phanen

lua/opencode/ui/renderer/buffer.lua

@@ -638,7 +638,7 @@ function M.upsert_part_now(part_id, message_id, formatted_data, previous_formatt
     end
 
     if formatted_data.fold_ranges and #formatted_data.fold_ranges > 0 then
-      M.set_all_folds()
+      M.update_part_folds(part_id)
     end
 
     return true
@@ -647,83 +647,87 @@ function M.upsert_part_now(part_id, message_id, formatted_data, previous_formatt
   return false
 end
 
-function M.set_all_folds()
+---@param data Output
+---@param line_start integer
+---@return {from: integer, to: integer}[]
+local function compute_part_folds(data, line_start)
+  local folds = {}
+  for _, f in ipairs(data.fold_ranges) do
+    folds[#folds + 1] = {
+      from = line_start + f.from - 1,
+      to = line_start + f.to - 1,
+    }
+  end
+  return folds
+end
+
+---Self-checked rebuild from ctx.part_folds.
+---* cached.line_start mismatch -> recompute (covers shift_all)
+---* source part no longer in formatted_parts -> drop entry (covers part removal)
+local function rebuild_folds()
+  for part_id, cached in pairs(ctx.part_folds) do
+    local part = ctx.render_state:get_part(part_id)
+    local data = ctx.formatted_parts[part_id]
+    if not part or not part.line_start or not data or not data.fold_ranges then
+      ctx.part_folds[part_id] = nil
+    elseif cached.line_start ~= part.line_start then
+      ctx.part_folds[part_id] = {
+        line_start = part.line_start,
+        folds = compute_part_folds(data, part.line_start),
+      }
+    end
+  end
+
+  for part_id in pairs(ctx.part_folds) do
+    if not ctx.formatted_parts[part_id] then
+      ctx.part_folds[part_id] = nil
+    end
+  end
+
   local all_folds = {}
-  ctx.part_folds = {}
-  for part_id_iter, data in pairs(ctx.formatted_parts) do
-    if data.fold_ranges then
-      local cached_part = ctx.render_state:get_part(part_id_iter)
-      if cached_part and cached_part.line_start then
-        local part_abs_folds = {}
-        for _, f in ipairs(data.fold_ranges) do
-          local abs = {
-            from = cached_part.line_start + f.from - 1,
-            to = cached_part.line_start + f.to - 1,
-          }
-          table.insert(part_abs_folds, abs)
-          table.insert(all_folds, abs)
-        end
-        ctx.part_folds[part_id_iter] = part_abs_folds
-      end
+  for _, cached in pairs(ctx.part_folds) do
+    for _, f in ipairs(cached.folds) do
+      all_folds[#all_folds + 1] = f
     end
   end
-  ctx.global_folds = all_folds
+  table.sort(all_folds, function(a, b)
+    return a.from < b.from
+  end)
   output_window.set_folds(all_folds)
 end
 
-local function folds_equal(a, b)
-  if not a or not b then
-    return false
-  end
-  if #a ~= #b then
-    return false
-  end
-  for i = 1, #a do
-    if a[i].from ~= b[i].from or a[i].to ~= b[i].to then
-      return false
+function M.set_all_folds()
+  ctx.part_folds = {}
+  for part_id, data in pairs(ctx.formatted_parts) do
+    local part = ctx.render_state:get_part(part_id)
+    if part and part.line_start and data.fold_ranges then
+      ctx.part_folds[part_id] = {
+        line_start = part.line_start,
+        folds = compute_part_folds(data, part.line_start),
+      }
     end
   end
-  return true
+  rebuild_folds()
 end
 
----Update folds for a single part during streaming, avoiding a full rebuild.
 ---@param part_id string
 function M.update_part_folds(part_id)
-  local formatted_data = ctx.formatted_parts[part_id]
-  if not formatted_data or not formatted_data.fold_ranges then
-    ctx.part_folds[part_id] = nil
-    M.set_all_folds()
-    return
-  end
-  local cached_part = ctx.render_state:get_part(part_id)
-  if not cached_part or not cached_part.line_start then
-    return
-  end
+  local data = ctx.formatted_parts[part_id]
+  local part = ctx.render_state:get_part(part_id)
 
-  local new_folds = {}
-  for _, f in ipairs(formatted_data.fold_ranges) do
-    table.insert(new_folds, {
-      from = cached_part.line_start + f.from - 1,
-      to = cached_part.line_start + f.to - 1,
-    })
-  end
-
-  if folds_equal(ctx.part_folds[part_id], new_folds) then
-    return
-  end
-
-  ctx.part_folds[part_id] = new_folds
-  local new_global = {}
-  for _, pf in pairs(ctx.part_folds) do
-    for _, f in ipairs(pf) do
-      table.insert(new_global, f)
+  if not part or not part.line_start or not data or not data.fold_ranges then
+    ctx.part_folds[part_id] = nil
+  else
+    local cached = ctx.part_folds[part_id]
+    if not cached or cached.line_start ~= part.line_start then
+      ctx.part_folds[part_id] = {
+        line_start = part.line_start,
+        folds = compute_part_folds(data, part.line_start),
+      }
     end
   end
-  table.sort(new_global, function(a, b)
-    return a.from < b.from
-  end)
-  ctx.global_folds = new_global
-  output_window.set_folds(new_global)
+
+  rebuild_folds()
 end
 
 ---@param part_id string

lua/opencode/ui/renderer/ctx.lua

@@ -29,9 +29,12 @@ local ctx = {
   bulk_buffer_lines = {},
   bulk_extmarks_by_line = {},
   bulk_folds = {},
-  ---@type {from: number, to: number}[]
-  global_folds = {},
-  ---@type table<string, {from: number, to: number}[]>
+  ---@type table<string, { line_start: integer, folds: {from: integer, to: integer}[] }>
+  ---Absolute folds per part, keyed by part_id. Used by set_all_folds /
+  ---update_part_folds to avoid recomputing folds for unchanged parts.
+  ---Invalidation is self-contained: rebuild compares cached.line_start
+  ---against render_state line_start, and clears entries whose
+  ---formatted_parts source has been removed.
   part_folds = {},
   ---@type integer|nil Number of messages to render from the end (nil = all)
   lazy_render_count = nil,
@@ -56,7 +59,6 @@ function ctx:reset()
   }
   self.flush_scheduled = false
   self.markdown_render_scheduled = false
-  self.global_folds = {}
   self.part_folds = {}
   self:bulk_reset()
 end

tests/unit/renderer_buffer_spec.lua

@@ -197,65 +197,169 @@ describe('renderer.buffer extmarks', function()
   end)
 end)
 
-describe('update_part_folds', function()
+describe('set_all_folds', function()
   local set_folds_stub
 
   before_each(function()
     ctx:reset()
     set_folds_stub = stub(output_window, 'set_folds')
-    ctx.global_folds = {}
-    ctx.part_folds = {}
   end)
 
   after_each(function()
     set_folds_stub:revert()
     ctx:reset()
   end)
 
-  it('computes absolute fold ranges for a single part', function()
+  it('rebuilds fold ranges from all formatted parts', function()
     ctx.formatted_parts['part_a'] = {
       lines = { 'title', '', 'content', 'more' },
       fold_ranges = { { from = 1, to = 4 } },
     }
+    ctx.formatted_parts['part_b'] = {
+      lines = { 'b1', 'b2', 'b3', 'b4' },
+      fold_ranges = { { from = 2, to = 3 } },
+    }
     ctx.render_state:set_part({ id = 'part_a', messageID = 'msg_1', type = 'text' }, 10, 14)
+    ctx.render_state:set_part({ id = 'part_b', messageID = 'msg_2', type = 'text' }, 20, 24)
 
-    buffer.update_part_folds('part_a')
+    buffer.set_all_folds()
 
     assert.stub(set_folds_stub).was_called_with({
       { from = 10, to = 13 },
+      { from = 21, to = 22 },
     })
   end)
 
-  it('skips set_folds when fold ranges have not changed', function()
+  it('omits folds for parts without cached line_start', function()
+    ctx.formatted_parts['part_a'] = {
+      lines = { 'title' },
+      fold_ranges = { { from = 1, to = 4 } },
+    }
+
+    buffer.set_all_folds()
+
+    assert.stub(set_folds_stub).was_called_with({})
+  end)
+
+  it('omits folds for parts removed from formatted_parts', function()
     ctx.formatted_parts['part_a'] = {
       lines = { 'title', '', 'content', 'more' },
       fold_ranges = { { from = 1, to = 4 } },
     }
     ctx.render_state:set_part({ id = 'part_a', messageID = 'msg_1', type = 'text' }, 10, 14)
 
-    buffer.update_part_folds('part_a')
+    buffer.set_all_folds()
     set_folds_stub:clear()
 
+    ctx.formatted_parts['part_a'] = nil
+    buffer.set_all_folds()
+
+    assert.stub(set_folds_stub).was_called_with({})
+  end)
+end)
+
+describe('update_part_folds', function()
+  local set_folds_stub
+
+  before_each(function()
+    ctx:reset()
+    set_folds_stub = stub(output_window, 'set_folds')
+  end)
+
+  after_each(function()
+    set_folds_stub:revert()
+    ctx:reset()
+  end)
+
+  it('reuses cached folds when line_start is unchanged', function()
+    ctx.formatted_parts['part_a'] = {
+      lines = { 'title', '', 'content' },
+      fold_ranges = { { from = 1, to = 3 } },
+    }
+    ctx.formatted_parts['part_b'] = {
+      lines = { 'b1', 'b2', 'b3' },
+      fold_ranges = { { from = 1, to = 2 } },
+    }
+    ctx.render_state:set_part({ id = 'part_a', messageID = 'msg_1', type = 'text' }, 10, 12)
+    ctx.render_state:set_part({ id = 'part_b', messageID = 'msg_2', type = 'text' }, 20, 22)
+
     buffer.update_part_folds('part_a')
+    local first_folds_a = ctx.part_folds['part_a'].folds
+    local first_folds_b = ctx.part_folds['part_b']
+    set_folds_stub:clear()
+
+    buffer.update_part_folds('part_b')
 
-    assert.stub(set_folds_stub).was_not_called()
+    assert.same(first_folds_a, ctx.part_folds['part_a'].folds)
+    assert.is_nil(first_folds_b)
+    assert.is_not_nil(ctx.part_folds['part_b'])
   end)
 
-  it('merges existing folds from other parts', function()
-    ctx.part_folds['part_b'] = { { from = 5, to = 8 } }
-    ctx.global_folds = { { from = 5, to = 8 } }
+  it('invalidates cached folds when line_start changes (shift)', function()
+    ctx.formatted_parts['part_b'] = {
+      lines = { 'b1', 'b2', 'b3' },
+      fold_ranges = { { from = 1, to = 2 } },
+    }
+    ctx.render_state:set_part({ id = 'part_b', messageID = 'msg_2', type = 'text' }, 20, 22)
+
+    buffer.update_part_folds('part_b')
+    local original = ctx.part_folds['part_b'].folds
+
+    ctx.render_state:set_part({ id = 'part_b', messageID = 'msg_2', type = 'text' }, 50, 52)
+
+    buffer.update_part_folds('part_b')
 
+    assert.is_not.same(original, ctx.part_folds['part_b'].folds)
+    assert.equals(20, original[1].from)
+    assert.equals(50, ctx.part_folds['part_b'].folds[1].from)
+  end)
+
+  it('clears cache entries when source part is removed from formatted_parts', function()
     ctx.formatted_parts['part_a'] = {
-      lines = { 'title', '', 'content', 'more' },
-      fold_ranges = { { from = 1, to = 4 } },
+      lines = { 'a' },
+      fold_ranges = { { from = 1, to = 1 } },
     }
-    ctx.render_state:set_part({ id = 'part_a', messageID = 'msg_1', type = 'text' }, 10, 14)
+    ctx.formatted_parts['part_b'] = {
+      lines = { 'b' },
+      fold_ranges = { { from = 1, to = 1 } },
+    }
+    ctx.render_state:set_part({ id = 'part_a', messageID = 'msg_1', type = 'text' }, 10, 10)
+    ctx.render_state:set_part({ id = 'part_b', messageID = 'msg_2', type = 'text' }, 20, 20)
 
     buffer.update_part_folds('part_a')
+    buffer.update_part_folds('part_b')
+    assert.is_not_nil(ctx.part_folds['part_a'])
+    assert.is_not_nil(ctx.part_folds['part_b'])
 
+    ctx.formatted_parts['part_b'] = nil
+    set_folds_stub:clear()
+    buffer.update_part_folds('part_a')
+
+    assert.is_not_nil(ctx.part_folds['part_a'])
+    assert.is_nil(ctx.part_folds['part_b'])
     assert.stub(set_folds_stub).was_called_with({
-      { from = 5, to = 8 },
-      { from = 10, to = 13 },
+      { from = 10, to = 10 },
+    })
+  end)
+
+  it('drops cache entry when part has no fold_ranges', function()
+    ctx.formatted_parts['part_a'] = {
+      lines = { 'a' },
+      fold_ranges = { { from = 1, to = 1 } },
+    }
+    ctx.formatted_parts['part_b'] = {
+      lines = { 'b' },
+    }
+    ctx.render_state:set_part({ id = 'part_a', messageID = 'msg_1', type = 'text' }, 10, 10)
+    ctx.render_state:set_part({ id = 'part_b', messageID = 'msg_2', type = 'text' }, 20, 20)
+
+    set_folds_stub:clear()
+    buffer.update_part_folds('part_a')
+    buffer.update_part_folds('part_b')
+
+    assert.is_nil(ctx.part_folds['part_b'])
+    assert.stub(set_folds_stub).was_called_with({
+      { from = 10, to = 10 },
     })
   end)
 end)