From 72ba4ad52e2ba58435d3297d30394a20afe8a177 Mon Sep 17 00:00:00 2001
From: Aaron Pham <contact@aarnphm.xyz>
Date: Mon, 26 Aug 2024 17:18:37 -0400
Subject: [PATCH] refactor: cleanup diff highlights (#247)

make sure to run on scheme change

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
---
 lua/avante/diff.lua                      |  33 +-------
 lua/avante/highlights.lua                | 101 +++++++++++++++++++++--
 lua/avante/init.lua                      |  11 +++
 lua/avante/{utils/init.lua => utils.lua} |   3 +-
 lua/avante/utils/colors.lua              |  40 ---------
 5 files changed, 106 insertions(+), 82 deletions(-)
 rename lua/avante/{utils/init.lua => utils.lua} (99%)
 delete mode 100644 lua/avante/utils/colors.lua

diff --git a/lua/avante/diff.lua b/lua/avante/diff.lua
index ca3d8b9..f7e800e 100644
--- a/lua/avante/diff.lua
+++ b/lua/avante/diff.lua
@@ -4,6 +4,7 @@ local M = {}
 
 local Config = require("avante.config")
 local Utils = require("avante.utils")
+local Highlights = require("avante.highlights")
 
 local fn = vim.fn
 local api = vim.api
@@ -492,44 +493,14 @@ local function clear_buffer_mappings(bufnr)
   vim.b[bufnr].conflict_mappings_set = false
 end
 
------------------------------------------------------------------------------//
--- Highlights
------------------------------------------------------------------------------//
-
----Derive the colour of the section label highlights based on each sections highlights
----@param highlights AvanteConflictHighlights
-local function set_highlights(highlights)
-  local current_color = Utils.get_hl(highlights.current)
-  local incoming_color = Utils.get_hl(highlights.incoming)
-  local ancestor_color = Utils.get_hl(highlights.ancestor)
-  local current_bg = current_color.bg or DEFAULT_CURRENT_BG_COLOR
-  local incoming_bg = incoming_color.bg or DEFAULT_INCOMING_BG_COLOR
-  local ancestor_bg = ancestor_color.bg or DEFAULT_ANCESTOR_BG_COLOR
-  local current_label_bg = Utils.colors.shade_color(current_bg, 60)
-  local incoming_label_bg = Utils.colors.shade_color(incoming_bg, 60)
-  local ancestor_label_bg = Utils.colors.shade_color(ancestor_bg, 60)
-  api.nvim_set_hl(0, CURRENT_HL, { bg = current_bg, bold = true, default = true })
-  api.nvim_set_hl(0, INCOMING_HL, { bg = incoming_bg, bold = true, default = true })
-  api.nvim_set_hl(0, ANCESTOR_HL, { bg = ancestor_bg, bold = true, default = true })
-  api.nvim_set_hl(0, CURRENT_LABEL_HL, { bg = current_label_bg, default = true })
-  api.nvim_set_hl(0, INCOMING_LABEL_HL, { bg = incoming_label_bg, default = true })
-  api.nvim_set_hl(0, ANCESTOR_LABEL_HL, { bg = ancestor_label_bg, default = true })
-end
-
 function M.setup()
-  set_highlights(Config.diff.highlights)
+  Highlights.conflict_highlights()
 
   set_commands()
 
   set_plug_mappings()
 
   local augroup = api.nvim_create_augroup(AUGROUP_NAME, { clear = true })
-  api.nvim_create_autocmd("ColorScheme", {
-    group = augroup,
-    callback = function()
-      set_highlights(Config.diff.highlights)
-    end,
-  })
 
   api.nvim_create_autocmd("User", {
     group = augroup,
diff --git a/lua/avante/highlights.lua b/lua/avante/highlights.lua
index 9b4477a..9389b3c 100644
--- a/lua/avante/highlights.lua
+++ b/lua/avante/highlights.lua
@@ -1,8 +1,10 @@
 local api = vim.api
 
 local Config = require("avante.config")
+local bit = require("bit")
+local rshift, band = bit.rshift, bit.band
 
-local H = {
+local Highlights = {
   TITLE = { name = "AvanteTitle", fg = "#1e222a", bg = "#98c379" },
   REVERSED_TITLE = { name = "AvanteReversedTitle", fg = "#98c379" },
   SUBTITLE = { name = "AvanteSubtitle", fg = "#1e222a", bg = "#56b6c2" },
@@ -11,6 +13,18 @@ local H = {
   REVERSED_THIRD_TITLE = { name = "AvanteReversedThirdTitle", fg = "#353B45" },
 }
 
+Highlights.conflict = {
+  CURRENT = { name = "AvanteConflictCurrent", bg = 4218238 }, -- #405d7e
+  CURRENT_LABEL = { name = "AvanteConflictCurrentLabel", link = "CURRENT", shade = 60 },
+  INCOMING = { name = "AvanteConflictIncoming", bg = 3229523 }, -- #314753
+  INCOMING_LABEL = { name = "AvanteConflictIncomingLabel", link = "INCOMING", shade = 60 },
+  ANCESTOR = { name = "AvanteConflictAncestor", bg = 6824314 }, -- #68217A
+  ANCESTOR_LABEL = { name = "AvanteConflictAncestorLabel", link = "ANCESTOR", shade = 60 },
+}
+
+--- helper
+local H = {}
+
 local M = {}
 
 M.input_ns = api.nvim_create_namespace("avante_input")
@@ -26,26 +40,95 @@ M.setup = function()
   local normal_float = api.nvim_get_hl(0, { name = "NormalFloat" })
 
   if Config.behaviour.auto_set_highlight_group then
-    vim.iter(H):each(function(_, hl)
-      if not has_set_colors(hl.name) then
-        api.nvim_set_hl(0, hl.name, { fg = hl.fg, bg = hl.bg or nil })
-      end
-    end)
+    vim
+      .iter(Highlights)
+      :filter(function(k, _)
+        -- return all uppercase key with underscore or fully uppercase key
+        return k:match("^%u+_") or k:match("^%u+$")
+      end)
+      :each(function(_, hl)
+        if not has_set_colors(hl.name) then
+          api.nvim_set_hl(0, hl.name, { fg = hl.fg or nil, bg = hl.bg or nil })
+        end
+      end)
+
+    M.conflict_highlights()
   end
 
   api.nvim_set_hl(M.hint_ns, "NormalFloat", { fg = normal_float.fg, bg = normal_float.bg })
-
   api.nvim_set_hl(M.input_ns, "NormalFloat", { fg = normal_float.fg, bg = normal_float.bg })
   api.nvim_set_hl(M.input_ns, "FloatBorder", { fg = normal.fg, bg = normal.bg })
 end
 
+---@param opts? AvanteConflictHighlights
+M.conflict_highlights = function(opts)
+  opts = opts or Config.diff.highlights
+
+  local get_default_colors = function(key, hl)
+    local getter = api.nvim_get_hl(0, { name = opts[key:lower()] })
+    return getter.bg or hl.bg
+  end
+
+  local get_shade = function(hl)
+    local color = get_default_colors(hl.link, Highlights.conflict[hl.link])
+    return H.shade_color(color, hl.shade)
+  end
+
+  vim.iter(Highlights.conflict):each(function(key, hl)
+    if not has_set_colors(hl.name) then
+      if hl.link ~= nil then
+        api.nvim_set_hl(0, hl.name, { bg = get_shade(hl), default = true })
+      else
+        api.nvim_set_hl(0, hl.name, { bg = get_default_colors(key, hl), default = true, bold = true })
+      end
+    end
+  end)
+end
+
 setmetatable(M, {
   __index = function(t, k)
-    if H[k] ~= nil then
-      return H[k].name
+    if Highlights[k] ~= nil then
+      return Highlights[k].name
+    elseif Highlights.conflict[k] ~= nil then
+      return Highlights.conflict[k].name
     end
     return t[k]
   end,
 })
 
+--- Returns a table containing the RGB values encoded inside 24 least
+--- significant bits of the number @rgb_24bit
+---
+---@param rgb_24bit number 24-bit RGB value
+---@return {r: integer, g: integer, b: integer} with keys 'r', 'g', 'b' in [0,255]
+H.decode_24bit_rgb = function(rgb_24bit)
+  vim.validate({ rgb_24bit = { rgb_24bit, "n", true } })
+  local r = band(rshift(rgb_24bit, 16), 255)
+  local g = band(rshift(rgb_24bit, 8), 255)
+  local b = band(rgb_24bit, 255)
+  return { r = r, g = g, b = b }
+end
+
+---@param attr integer
+---@param percent integer
+H.alter = function(attr, percent)
+  return math.floor(attr * (100 + percent) / 100)
+end
+
+---@source https://stackoverflow.com/q/5560248
+---@see https://stackoverflow.com/a/37797380
+---Darken a specified hex color
+---@param color number
+---@param percent number
+---@return string
+H.shade_color = function(color, percent)
+  local rgb = H.decode_24bit_rgb(color)
+  if not rgb.r or not rgb.g or not rgb.b then
+    return "NONE"
+  end
+  local r, g, b = H.alter(rgb.r, percent), H.alter(rgb.g, percent), H.alter(rgb.b, percent)
+  r, g, b = math.min(r, 255), math.min(g, 255), math.min(b, 255)
+  return string.format("#%02x%02x%02x", r, g, b)
+end
+
 return M
diff --git a/lua/avante/init.lua b/lua/avante/init.lua
index fbfaaf7..b236d0b 100644
--- a/lua/avante/init.lua
+++ b/lua/avante/init.lua
@@ -64,6 +64,8 @@ H.keymaps = function()
   })
 end
 
+H.augroup = api.nvim_create_augroup("avante-autocmds", { clear = true })
+
 H.autocmds = function()
   local ok, LazyConfig = pcall(require, "lazy.core.config")
 
@@ -95,6 +97,7 @@ H.autocmds = function()
   end
 
   api.nvim_create_autocmd("TabEnter", {
+    group = H.augroup,
     pattern = "*",
     once = true,
     callback = function(ev)
@@ -107,6 +110,7 @@ H.autocmds = function()
   })
 
   api.nvim_create_autocmd("TabClosed", {
+    group = H.augroup,
     pattern = "*",
     callback = function(ev)
       local tab = tonumber(ev.file)
@@ -131,6 +135,13 @@ H.autocmds = function()
     end
   end)
 
+  api.nvim_create_autocmd("ColorSchemePre", {
+    group = H.augroup,
+    callback = function()
+      require("avante.highlights").setup()
+    end,
+  })
+
   -- automatically setup Avante filetype to markdown
   vim.treesitter.language.register("markdown", "Avante")
 end
diff --git a/lua/avante/utils/init.lua b/lua/avante/utils.lua
similarity index 99%
rename from lua/avante/utils/init.lua
rename to lua/avante/utils.lua
index 0c7a7d1..54fb978 100644
--- a/lua/avante/utils/init.lua
+++ b/lua/avante/utils.lua
@@ -1,7 +1,6 @@
 local api = vim.api
 
----@class avante.Utils: LazyUtilCore
----@field colors avante.util.colors
+---@class avante.utils: LazyUtilCore
 local M = {}
 
 setmetatable(M, {
diff --git a/lua/avante/utils/colors.lua b/lua/avante/utils/colors.lua
deleted file mode 100644
index e5baf40..0000000
--- a/lua/avante/utils/colors.lua
+++ /dev/null
@@ -1,40 +0,0 @@
----@class avante.util.colors
-local M = {}
-
-local bit = require("bit")
-local rshift, band = bit.rshift, bit.band
-
---- Returns a table containing the RGB values encoded inside 24 least
---- significant bits of the number @rgb_24bit
----
---@param rgb_24bit (number) 24-bit RGB value
---@returns (table) with keys 'r', 'g', 'b' in [0,255]
-local function decode_24bit_rgb(rgb_24bit)
-  vim.validate({ rgb_24bit = { rgb_24bit, "n", true } })
-  local r = band(rshift(rgb_24bit, 16), 255)
-  local g = band(rshift(rgb_24bit, 8), 255)
-  local b = band(rgb_24bit, 255)
-  return { r = r, g = g, b = b }
-end
-
-local function alter(attr, percent)
-  return math.floor(attr * (100 + percent) / 100)
-end
-
----@source https://stackoverflow.com/q/5560248
----@see: https://stackoverflow.com/a/37797380
----Darken a specified hex color
----@param color string
----@param percent number
----@return string
-function M.shade_color(color, percent)
-  local rgb = decode_24bit_rgb(color)
-  if not rgb.r or not rgb.g or not rgb.b then
-    return "NONE"
-  end
-  local r, g, b = alter(rgb.r, percent), alter(rgb.g, percent), alter(rgb.b, percent)
-  r, g, b = math.min(r, 255), math.min(g, 255), math.min(b, 255)
-  return string.format("#%02x%02x%02x", r, g, b)
-end
-
-return M