From bc6ab47840b487ff16c4d2c51e95dd5eaed49087 Mon Sep 17 00:00:00 2001
From: Steven Xu <stevenxxiu@gmail.com>
Date: Thu, 12 Dec 2024 19:15:23 +1100
Subject: [PATCH] refactor: move `Buffer:set_line_highlights()` ANSI highlights
 to `Buffer:set_ansi_highlights()`, to avoid finding hunk line numbers
 manually with the hard-coded `NeogitDiffContext`

- Add `ansi_hl` to `ComponentOptions`, to indicate whether the component needs ANSI highlighting.
- Add `ansi_highlight` to `RendererBuffer`, to pass the line numbers.
---
 lua/neogit/buffers/common.lua   | 17 ++++++++++++----
 lua/neogit/lib/buffer.lua       | 36 +++++++--------------------------
 lua/neogit/lib/ui/component.lua |  1 +
 lua/neogit/lib/ui/init.lua      |  1 +
 lua/neogit/lib/ui/renderer.lua  |  9 +++++++++
 5 files changed, 31 insertions(+), 33 deletions(-)

diff --git a/lua/neogit/buffers/common.lua b/lua/neogit/buffers/common.lua
index 80a948196..dc565177f 100644
--- a/lua/neogit/buffers/common.lua
+++ b/lua/neogit/buffers/common.lua
@@ -80,10 +80,19 @@ local HunkLine = Component.new(function(line)
 end)
 
 M.Hunk = Component.new(function(props)
-  return col.tag("Hunk")({
-    text.line_hl("NeogitHunkHeader")(props.header),
-    col.tag("HunkContent")(map(props.content, HunkLine)),
-  }, { foldable = true, folded = props.folded or false, context = true, hunk = props.hunk })
+  return col.tag("Hunk")(
+    {
+      text.line_hl("NeogitHunkHeader")(props.header),
+      col.tag("HunkContent")(map(props.content, HunkLine)),
+    },
+    {
+      foldable = true,
+      folded = props.folded or false,
+      context = true,
+      hunk = props.hunk,
+      ansi_hl = config.values.log_pager ~= nil,
+    }
+  )
 end)
 
 M.List = Component.new(function(props)
diff --git a/lua/neogit/lib/buffer.lua b/lua/neogit/lib/buffer.lua
index 47c20a077..4667b7433 100644
--- a/lua/neogit/lib/buffer.lua
+++ b/lua/neogit/lib/buffer.lua
@@ -139,38 +139,16 @@ function Buffer:set_extmarks(extmarks)
 end
 
 function Buffer:set_line_highlights(highlights)
-  local line_ansi_colorized_map = {}
-
   for _, hl in ipairs(highlights) do
-    local line_nr, hl_group = unpack(hl)
-    if hl_group == "NeogitDiffContext" then
-      line_ansi_colorized_map[line_nr] = true
-    else
-      self:add_line_highlight(unpack(hl))
-    end
-  end
-
-  local line_ansi_colorized = {}
-  for k in pairs(line_ansi_colorized_map) do
-    table.insert(line_ansi_colorized, k)
+    self:add_line_highlight(unpack(hl))
   end
-  table.sort(line_ansi_colorized)
+end
 
-  local start_line_nr, prev_line_nr
-  for _, line_nr in ipairs(line_ansi_colorized) do
-    if start_line_nr == nil then
-      start_line_nr = line_nr
-    end
-    if prev_line_nr ~= nil and line_nr ~= prev_line_nr + 1 then
-      local text = self:get_lines(start_line_nr, prev_line_nr + 1, false)
-      vim.g.baleia.buf_set_lines(self.handle, start_line_nr, prev_line_nr + 1, false, text)
-      start_line_nr = line_nr
-    end
-    prev_line_nr = line_nr
-  end
-  if start_line_nr ~= nil then
-    local text = self:get_lines(start_line_nr, prev_line_nr + 1, false)
-    vim.g.baleia.buf_set_lines(self.handle, start_line_nr, prev_line_nr + 1, false, text)
+function Buffer:set_ansi_highlights(highlights)
+  for _, hl in ipairs(highlights) do
+    local first_line, last_line = unpack(hl)
+    local text = self:get_lines(first_line, last_line, false)
+    vim.g.baleia.buf_set_lines(self.handle, first_line, last_line, false, text)
   end
 end
 
diff --git a/lua/neogit/lib/ui/component.lua b/lua/neogit/lib/ui/component.lua
index ae737977c..15c8f2375 100644
--- a/lua/neogit/lib/ui/component.lua
+++ b/lua/neogit/lib/ui/component.lua
@@ -13,6 +13,7 @@ local default_component_options = {
 
 ---@class ComponentOptions
 ---@field line_hl string
+---@field ansi_hl boolean
 ---@field highlight string
 ---@field align_right integer|nil
 ---@field padding_left integer
diff --git a/lua/neogit/lib/ui/init.lua b/lua/neogit/lib/ui/init.lua
index 95b2d5e0e..9d6eee222 100644
--- a/lua/neogit/lib/ui/init.lua
+++ b/lua/neogit/lib/ui/init.lua
@@ -686,6 +686,7 @@ function Ui:update()
     self.buf:set_highlights(renderer.buffer.highlight)
     self.buf:set_extmarks(renderer.buffer.extmark)
     self.buf:set_line_highlights(renderer.buffer.line_highlight)
+    self.buf:set_ansi_highlights(renderer.buffer.ansi_highlight)
     self.buf:set_folds(renderer.buffer.fold)
 
     self.statuscolumn = {}
diff --git a/lua/neogit/lib/ui/renderer.lua b/lua/neogit/lib/ui/renderer.lua
index 00fce647c..641e14f3f 100644
--- a/lua/neogit/lib/ui/renderer.lua
+++ b/lua/neogit/lib/ui/renderer.lua
@@ -65,6 +65,7 @@ end
 ---@field line string[]
 ---@field highlight table[]
 ---@field line_highlight table[]
+---@field ansi_highlight table[]
 ---@field extmark table[]
 ---@field fold table[]
 
@@ -93,6 +94,7 @@ function Renderer:new(layout, buffer)
       line = {},
       highlight = {},
       line_highlight = {},
+      ansi_highlight = {},
       extmark = {},
       fold = {},
     },
@@ -194,6 +196,13 @@ function Renderer:_render_child(child)
     table.insert(self.buffer.line_highlight, { #self.buffer.line - 1, line_hl })
   end
 
+  if child.options.ansi_hl then
+    table.insert(self.buffer.ansi_highlight, {
+      #self.buffer.line - (child.position.row_end - child.position.row_start),
+      #self.buffer.line,
+    })
+  end
+
   if child.options.virtual_text then
     table.insert(self.buffer.extmark, {
       self.namespace,