Optimize rpc msg parse loop

Author: mfussenegger

lua/dap/rpc.lua

@@ -1,65 +1,79 @@
 local utils = require('dap.utils')
 local M = {}
 
--- Copied from neovim rpc.lua
 ---@param header string
-local function parse_headers(header)
-  local headers = {}
-  for line in vim.gsplit(header, '\r\n', {plain = true}) do
-    if line == '' then
-      break
-    end
-    local key, value = line:match('^%s*(%S+)%s*:%s*(.+)%s*$')
-    if key then
-      key = key:lower():gsub('%-', '_')
-      headers[key] = value
-    else
-      error(string.format("Invalid header line %q", line))
+---@return integer?
+local function get_content_length(header)
+  for line in header:gmatch("(.-)\r\n") do
+    local key, value = line:match('^%s*(%S+)%s*:%s*(%d+)%s*$')
+    if key and key:lower() == "content-length" then
+      return tonumber(value)
     end
   end
-  headers.content_length = tonumber(headers.content_length)
-    or error(string.format("Content-Length not found in headers. %q", header))
-  return headers
 end
 
 
--- Mostly copied from neovim rpc.lua
-local header_start_pattern = ("content"):gsub("%w", function(c) return "["..c..c:upper().."]" end)
-local function parse_chunk_loop()
-  local buffer = ''
-  while true do
-    local start, finish = buffer:find('\r\n\r\n', 1, true)
-    if start then
-      local buffer_start = buffer:find(header_start_pattern)
-      if not buffer_start then
-        error(string.format(
-          "Headers were expected but debug adapter sent: %s",
-          buffer
-        ))
-      end
-      local headers = parse_headers(buffer:sub(buffer_start, start - 1))
-      local content_length = headers.content_length
-      local body_chunks = {buffer:sub(finish + 1)}
-      local body_length = #body_chunks[1]
-      while body_length < content_length do
+local parse_chunk_loop
+local has_strbuffer, strbuffer = pcall(require, "string.buffer")
+
+if has_strbuffer then
+  parse_chunk_loop = function()
+    local buf = strbuffer.new()
+    while true do
+      local msg = buf:tostring()
+      local header_end = msg:find('\r\n\r\n', 1, true)
+      if header_end then
+        local header = buf:get(header_end + 1)
+        buf:skip(2) -- skip past header boundary
+        local content_length = get_content_length(header)
+        if not content_length then
+          error("Content-Length not found in headers: " .. header)
+        end
+        while #buf < content_length do
+          local chunk = coroutine.yield()
+          buf:put(chunk)
+        end
+        local body = buf:get(content_length)
+        coroutine.yield(body)
+      else
         local chunk = coroutine.yield()
-          or error("Expected more data for the body. The server may have died.")
-        table.insert(body_chunks, chunk)
-        body_length = body_length + #chunk
+        buf:put(chunk)
       end
-      local last_chunk = body_chunks[#body_chunks]
+    end
+  end
+else
+  parse_chunk_loop = function()
+    local buffer = ''
+    while true do
+      local header_end, body_start = buffer:find('\r\n\r\n', 1, true)
+      if header_end then
+        local header = buffer:sub(1, header_end + 1)
+        local content_length = get_content_length(header)
+        if not content_length then
+          error("Content-Length not found in headers: " .. header)
+        end
+        local body_chunks = {buffer:sub(body_start + 1)}
+        local body_length = #body_chunks[1]
+        while body_length < content_length do
+          local chunk = coroutine.yield()
+            or error("Expected more data for the body. The server may have died.")
+          table.insert(body_chunks, chunk)
+          body_length = body_length + #chunk
+        end
+        local last_chunk = body_chunks[#body_chunks]
 
-      body_chunks[#body_chunks] = last_chunk:sub(1, content_length - body_length - 1)
-      local rest = ''
-      if body_length > content_length then
-        rest = last_chunk:sub(content_length - body_length)
+        body_chunks[#body_chunks] = last_chunk:sub(1, content_length - body_length - 1)
+        local rest = ''
+        if body_length > content_length then
+          rest = last_chunk:sub(content_length - body_length)
+        end
+        local body = table.concat(body_chunks)
+        buffer = rest .. (coroutine.yield(body)
+          or error("Expected more data for the body. The server may have died."))
+      else
+        buffer = buffer .. (coroutine.yield()
+          or error("Expected more data for the header. The server may have died."))
       end
-      local body = table.concat(body_chunks)
-      buffer = rest .. (coroutine.yield(headers, body)
-        or error("Expected more data for the body. The server may have died."))
-    else
-      buffer = buffer .. (coroutine.yield()
-        or error("Expected more data for the header. The server may have died."))
     end
   end
 end
@@ -80,10 +94,10 @@ function M.create_read_loop(handle_body, on_no_chunk)
       return
     end
     while true do
-      local headers, body = parse_chunk(chunk)
-      if headers then
+      local body = parse_chunk(chunk)
+      if body then
         handle_body(body)
-        chunk = ''
+        chunk = ""
       else
         break
       end