From 19d970623b752448abe023bb6684a76141ca200a Mon Sep 17 00:00:00 2001
From: Tymoteusz Makowski <tymek.makowski@gmail.com>
Date: Tue, 29 Oct 2024 22:54:07 +0100
Subject: [PATCH] fix: preserve PrecognitionHighlight through colorscheme
 changes (#90)

---
 lua/precognition/init.lua       | 12 ++++++++++--
 tests/precognition/e2e_spec.lua | 10 ++++++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/lua/precognition/init.lua b/lua/precognition/init.lua
index fdf804c..f652432 100644
--- a/lua/precognition/init.lua
+++ b/lua/precognition/init.lua
@@ -439,6 +439,10 @@ function M.toggle()
     return visible
 end
 
+local function setup_highlights()
+    vim.api.nvim_set_hl(0, "PrecognitionHighlight", config.highlightColor)
+end
+
 ---@param opts Precognition.PartialConfig
 function M.setup(opts)
     opts = opts or {}
@@ -450,8 +454,12 @@ function M.setup(opts)
     ns = vim.api.nvim_create_namespace("precognition")
     au = vim.api.nvim_create_augroup("precognition", { clear = true })
 
-    local hl_name = "PrecognitionHighlight"
-    vim.api.nvim_set_hl(0, hl_name, config.highlightColor)
+    setup_highlights()
+    vim.api.nvim_create_autocmd("ColorScheme", {
+        desc = "Set precognition.nvim's highlights up",
+        group = au,
+        callback = setup_highlights,
+    })
 
     create_command()
 
diff --git a/tests/precognition/e2e_spec.lua b/tests/precognition/e2e_spec.lua
index 214d468..6917d7d 100644
--- a/tests/precognition/e2e_spec.lua
+++ b/tests/precognition/e2e_spec.lua
@@ -10,10 +10,10 @@ describe("e2e tests", function()
 
     it("auto commands are set", function()
         local autocmds = vim.api.nvim_get_autocmds({ group = "precognition" })
-        eq(4, vim.tbl_count(autocmds))
+        eq(5, vim.tbl_count(autocmds))
         precognition.peek()
         autocmds = vim.api.nvim_get_autocmds({ group = "precognition" })
-        eq(7, vim.tbl_count(autocmds))
+        eq(8, vim.tbl_count(autocmds))
     end)
 
     -- it("namespace is created", function()
@@ -223,6 +223,12 @@ describe("e2e tests", function()
         eq("PrecognitionHighlight", extmarks[3].virt_lines[1][1][2])
         eq(customMark, vim.api.nvim_get_hl(0, { name = extmarks[3].virt_lines[1][1][2] }))
     end)
+
+    it("preserves highlight groups through a colorscheme change", function()
+        vim.cmd.colorscheme("default")
+        local hl = vim.api.nvim_get_hl(0, { name = "PrecognitionHighlight" })
+        eq(false, vim.tbl_isempty(hl))
+    end)
 end)
 
 describe("Gutter Priority", function()