perf(chat): optimize message storage and access

lua/CopilotChat/config/mappings.lua

@@ -279,9 +279,10 @@ return {
     normal = 'gqa',
     callback = function()
       local items = {}
-      for i, message in ipairs(copilot.chat.messages) do
+      local messages = copilot.chat:get_messages()
+      for i, message in ipairs(messages) do
         if message.section and message.role == constants.ROLE.ASSISTANT then
-          local prev_message = copilot.chat.messages[i - 1]
+          local prev_message = messages[i - 1]
           local text = ''
           if prev_message then
             text = prev_message.content
@@ -305,8 +306,8 @@ return {
     normal = 'gqd',
     callback = function(source)
       local items = {}
-
-      for _, message in ipairs(copilot.chat.messages) do
+      local messages = copilot.chat:get_messages()
+      for _, message in ipairs(messages) do
         if message.section then
           for _, block in ipairs(message.section.blocks) do
             local diff = get_diff(source.bufnr, block)

lua/CopilotChat/init.lua

@@ -356,7 +356,7 @@ function M.resolve_functions(prompt, config)
   local resolved_resources = {}
   local resolved_tools = {}
   local tool_calls = {}
-  for _, message in ipairs(M.chat.messages) do
+  for _, message in ipairs(M.chat:get_messages()) do
     if message.tool_calls then
       for _, tool_call in ipairs(message.tool_calls) do
         table.insert(tool_calls, tool_call)
@@ -868,7 +868,7 @@ function M.ask(prompt, config)
 
       local ask_response = client.ask(client, prompt, {
         headless = config.headless,
-        history = M.chat.messages,
+        history = M.chat:get_messages(),
         resources = resolved_resources,
         tools = selected_tools,
         system_prompt = system_prompt,
@@ -948,7 +948,7 @@ function M.save(name, history_path)
     return
   end
 
-  local history = vim.deepcopy(M.chat.messages)
+  local history = vim.deepcopy(M.chat:get_messages())
   for _, message in ipairs(history) do
     message.section = nil
   end

lua/CopilotChat/ui/chat.lua

@@ -4,6 +4,7 @@ local constants = require('CopilotChat.constants')
 local notify = require('CopilotChat.notify')
 local utils = require('CopilotChat.utils')
 local class = require('CopilotChat.utils.class')
+local orderedmap = require('CopilotChat.utils.orderedmap')
 
 function CopilotChatFoldExpr(lnum, separator)
   local to_match = separator .. '$'
@@ -93,7 +94,7 @@ end
 ---@field config CopilotChat.config.Shared
 ---@field token_count number?
 ---@field token_max_count number?
----@field messages table<CopilotChat.client.Message>
+---@field private messages OrderedMap<string, CopilotChat.client.Message>
 ---@field private layout CopilotChat.config.Layout?
 ---@field private headers table<string, string>
 ---@field private separator string
@@ -106,7 +107,7 @@ local Chat = class(function(self, config, on_buf_create)
   self.config = config
   self.token_count = nil
   self.token_max_count = nil
-  self.messages = {}
+  self.messages = orderedmap()
 
   self.layout = nil
   self.headers = {}
@@ -168,6 +169,9 @@ end
 ---@param cursor boolean? If true, returns the block closest to the cursor position
 ---@return CopilotChat.ui.chat.Block?
 function Chat:get_block(role, cursor)
+  self:parse()
+  local messages = self:get_messages()
+
   if cursor then
     if not self:visible() then
       return nil
@@ -178,7 +182,7 @@ function Chat:get_block(role, cursor)
     local closest_block = nil
     local max_line_below_cursor = -1
 
-    for _, message in ipairs(self.messages) do
+    for _, message in ipairs(messages) do
       local section = message.section
       local matches_role = not role or message.role == role
       if matches_role and section and section.blocks then
@@ -194,21 +198,28 @@ function Chat:get_block(role, cursor)
     return closest_block
   end
 
-  for i = #self.messages, 1, -1 do
-    local message = self.messages[i]
+  for i = #messages, 1, -1 do
+    local message = messages[i]
     local matches_role = not role or message.role == role
     if matches_role and message.section and message.section.blocks and #message.section.blocks > 0 then
       return message.section.blocks[#message.section.blocks]
     end
   end
 end
 
+--- Get list of all chat messages
+---@return table<CopilotChat.ui.chat.Message>
+function Chat:get_messages()
+  return self.messages:values()
+end
+
 --- Get last message by role in the chat window.
 ---@param role string? If specified, only considers sections of the given role
 ---@param cursor boolean? If true, returns the message closest to the cursor position
 ---@return CopilotChat.ui.chat.Message?
 function Chat:get_message(role, cursor)
   self:parse()
+  local messages = self:get_messages()
 
   if cursor then
     if not self:visible() then
@@ -220,7 +231,7 @@ function Chat:get_message(role, cursor)
     local closest_message = nil
     local max_line_below_cursor = -1
 
-    for _, message in ipairs(self.messages) do
+    for _, message in ipairs(messages) do
       local section = message.section
       local matches_role = not role or message.role == role
       if matches_role and section.start_line <= cursor_line and section.start_line > max_line_below_cursor then
@@ -232,8 +243,8 @@ function Chat:get_message(role, cursor)
     return closest_message
   end
 
-  for i = #self.messages, 1, -1 do
-    local message = self.messages[i]
+  for i = #messages, 1, -1 do
+    local message = messages[i]
     local matches_role = not role or message.role == role
     if matches_role then
       return message
@@ -479,7 +490,8 @@ end
 function Chat:add_message(message, replace)
   self:parse()
 
-  local current_message = self.messages[#self.messages]
+  local messages = self:get_messages()
+  local current_message = messages[#messages]
   local is_new = not current_message
     or current_message.role ~= message.role
     or (message.id and current_message.id ~= message.id)
@@ -488,7 +500,7 @@ function Chat:add_message(message, replace)
     -- Add appropriate header based on role and generate a new ID if not provided
     message.id = message.id or utils.uuid()
     local header = self.headers[message.role]
-    table.insert(self.messages, message)
+    self.messages:set(message.id, message)
 
     if current_message then
       self:append('\n')
@@ -546,12 +558,7 @@ function Chat:remove_message(role, cursor)
   vim.bo[self.bufnr].modifiable = modifiable
 
   -- Remove the message from the messages list
-  for i, msg in ipairs(self.messages) do
-    if msg.id == message.id then
-      table.remove(self.messages, i)
-      break
-    end
-  end
+  self.messages:remove(message.id)
 end
 
 --- Append text to the chat window.
@@ -585,7 +592,7 @@ function Chat:clear()
   self:validate()
   self.token_count = nil
   self.token_max_count = nil
-  self.messages = {}
+  self.messages = orderedmap()
 
   local modifiable = vim.bo[self.bufnr].modifiable
   vim.bo[self.bufnr].modifiable = true
@@ -718,15 +725,8 @@ function Chat:parse()
   -- Finish last message
   current_message.section.end_line = vim.api.nvim_buf_line_count(self.bufnr)
 
-  -- Build lookup table for previous messages by id
-  local old_messages_by_id = {}
-  for _, msg in ipairs(self.messages or {}) do
-    if msg.id then
-      old_messages_by_id[msg.id] = msg
-    end
-  end
-
   -- Format new messages and preserve extra fields from old messages
+  local messages = orderedmap()
   for _, message in ipairs(new_messages) do
     message.content = vim.trim(table.concat(message.content, '\n'))
     if message.section then
@@ -735,17 +735,20 @@ function Chat:parse()
       end
     end
 
-    local old = old_messages_by_id[message.id]
+    local old = self.messages:get(message.id)
     if old then
       for k, v in pairs(old) do
         if message[k] == nil then
           message[k] = v
         end
       end
     end
+
+    messages:set(message.id, message)
   end
 
-  self.messages = new_messages
+  -- Update messages
+  self.messages = messages
 end
 
 --- Render the chat window.
@@ -757,7 +760,9 @@ function Chat:render()
   vim.api.nvim_buf_clear_namespace(self.bufnr, highlight_ns, 0, -1) -- Clear previous highlights
   self:show_help() -- Clear previous help
 
-  for i, message in ipairs(self.messages) do
+  local messages = self:get_messages()
+
+  for i, message in ipairs(messages) do
     if self.config.highlight_headers then
       -- Overlay section header with nice display
       local header_value = self.headers[message.role]
@@ -847,7 +852,7 @@ function Chat:render()
       end
     end
 
-    if i == #self.messages and message.role == constants.ROLE.USER then
+    if i == #messages and message.role == constants.ROLE.USER then
       -- Highlight tools in the last user message
       local assistant_msg = self:get_message(constants.ROLE.ASSISTANT)
       if assistant_msg and assistant_msg.tool_calls and #assistant_msg.tool_calls > 0 then
@@ -883,7 +888,7 @@ function Chat:render()
 
     -- Auto fold non-assistant messages if enabled
     if self.config.auto_fold and self:visible() then
-      if message.role ~= constants.ROLE.ASSISTANT and message.section and i < #self.messages then
+      if message.role ~= constants.ROLE.ASSISTANT and message.section and i < #messages then
         vim.api.nvim_win_call(self.winnr, function()
           local fold_level = vim.fn.foldlevel(message.section.start_line)
           if fold_level > 0 and vim.fn.foldclosed(message.section.start_line) == -1 then

lua/CopilotChat/utils/orderedmap.lua

@@ -1,6 +1,7 @@
 ---@class OrderedMap<K, V>
 ---@field set fun(self:OrderedMap, key:any, value:any)
 ---@field get fun(self:OrderedMap, key:any):any
+---@field remove fun(self:OrderedMap, key:any)
 ---@field keys fun(self:OrderedMap):table
 ---@field values fun(self:OrderedMap):table
 
@@ -22,6 +23,18 @@ local function orderedmap()
       return self._data[key]
     end,
 
+    remove = function(self, key)
+      if self._data[key] then
+        self._data[key] = nil
+        for i, k in ipairs(self._keys) do
+          if k == key then
+            table.remove(self._keys, i)
+            break
+          end
+        end
+      end
+    end,
+
     keys = function(self)
       return self._keys
     end,

tests/orderedmap_spec.lua

@@ -25,4 +25,13 @@ describe('CopilotChat.utils.orderedmap', function()
     assert.are.same({ 'a' }, map:keys())
     assert.are.same({ 2 }, map:values())
   end)
+
+  it('removes values and updates order', function()
+    local map = orderedmap()
+    map:set('a', 1)
+    map:set('b', 2)
+    map:remove('a')
+    assert.are.same({ 'b' }, map:keys())
+    assert.are.same({ 2 }, map:values())
+  end)
 end)