From dc7216abbaf03ca6bc16e4397231fd8bb4c081de Mon Sep 17 00:00:00 2001
From: Meow Honk <21010072+catgoose@users.noreply.github.com>
Date: Sun, 22 Dec 2024 10:46:22 -0600
Subject: [PATCH] Fix/110 add names custom nonalpha chars to trie (#113)

* fix: fixes incorrect reference to sass name parser

* feat(Trie): adds method to add additional valid characters to Trie index_lookup

* chore: moves TODO comment which causes luadoc to skip function
---
 doc/colorizer.txt               |  4 +--
 doc/modules/colorizer.sass.html |  6 ++--
 lua/colorizer/config.lua        |  2 +-
 lua/colorizer/matcher.lua       |  1 +
 lua/colorizer/parser/names.lua  | 63 ++++++++++++++++++++++++---------
 lua/colorizer/sass.lua          |  2 +-
 lua/colorizer/trie.lua          | 21 ++++++++---
 lua/colorizer/utils.lua         |  1 +
 test/expect.lua                 | 22 +++++-------
 9 files changed, 80 insertions(+), 42 deletions(-)

diff --git a/doc/colorizer.txt b/doc/colorizer.txt
index c54f4d3..1f3b8be 100644
--- a/doc/colorizer.txt
+++ b/doc/colorizer.txt
@@ -910,7 +910,7 @@ LUA API                                                 *colorizer.sass-lua-api*
 Functions: ~
     |cleanup| - Cleanup sass variables and watch handlers
 
-    |name_parser| - Parse the given line for sass color names
+    |parser| - Parse the given line for sass color names
  check for value in state[buf].definitions_all
 
     |update_variables| - Parse the given lines for sass variabled and add to
@@ -925,7 +925,7 @@ cleanup({bufnr})                                        *colorizer.sass.cleanup*
 
 
 
-name_parser({line}, {i}, {bufnr})                   *colorizer.sass.name_parser*
+parser({line}, {i}, {bufnr})                             *colorizer.sass.parser*
     Parse the given line for sass color names
      check for value in state[buf].definitions_all
 
diff --git a/doc/modules/colorizer.sass.html b/doc/modules/colorizer.sass.html
index 7e2d66d..70e4a88 100644
--- a/doc/modules/colorizer.sass.html
+++ b/doc/modules/colorizer.sass.html
@@ -75,7 +75,7 @@ <h2><a href="#Functions">Functions</a></h2>
 	<td class="summary">Cleanup sass variables and watch handlers</td>
 	</tr>
 	<tr>
-	<td class="name" nowrap><a href="#name_parser">name_parser (line, i, bufnr)</a></td>
+	<td class="name" nowrap><a href="#parser">parser (line, i, bufnr)</a></td>
 	<td class="summary">Parse the given line for sass color names
  check for value in state[buf].definitions_all</td>
 	</tr>
@@ -113,8 +113,8 @@ <h3>Parameters:</h3>
 
 </dd>
     <dt>
-    <a name = "name_parser"></a>
-    <strong>name_parser (line, i, bufnr)</strong>
+    <a name = "parser"></a>
+    <strong>parser (line, i, bufnr)</strong>
     </dt>
     <dd>
     Parse the given line for sass color names
diff --git a/lua/colorizer/config.lua b/lua/colorizer/config.lua
index 42a3403..9adff0a 100644
--- a/lua/colorizer/config.lua
+++ b/lua/colorizer/config.lua
@@ -233,7 +233,7 @@ function M.parse_buffer_options(options)
     end
   end
 
-  -- https://github.com/NvChad/nvim-colorizer.lua/issues/48
+  -- https://github.com/catgoose/nvim-colorizer.lua/issues/48
   handle_alias("css", options, default)
   handle_alias("css_fn", options, default)
 
diff --git a/lua/colorizer/matcher.lua b/lua/colorizer/matcher.lua
index 25b076c..551b6c2 100644
--- a/lua/colorizer/matcher.lua
+++ b/lua/colorizer/matcher.lua
@@ -15,6 +15,7 @@ local parsers = {
   hsl_function = require("colorizer.parser.hsl").parser,
   rgb_function = require("colorizer.parser.rgb").parser,
   rgba_hex = require("colorizer.parser.rgba_hex").parser,
+  --  TODO: 2024-12-21 - Should this be moved into parsers module?
   sass_name = require("colorizer.sass").parser,
 }
 
diff --git a/lua/colorizer/parser/names.lua b/lua/colorizer/parser/names.lua
index 0ee2920..8988ede 100644
--- a/lua/colorizer/parser/names.lua
+++ b/lua/colorizer/parser/names.lua
@@ -34,6 +34,22 @@ local function add_color(name, value)
   names_cache.color_trie:insert(name)
 end
 
+--- Extract non-alphanumeric characters to add as a valid index in the Trie
+-- @param tbl table The table to extract non-alphanumeric characters from.
+local function extract_non_alphanum_keys(tbl)
+  local non_alphanum_chars = {}
+  for key, _ in pairs(tbl) do
+    for char in key:gmatch("[^%w]") do
+      non_alphanum_chars[char] = true
+    end
+  end
+  local result = ""
+  for char in pairs(non_alphanum_chars) do
+    result = result .. char
+  end
+  return result
+end
+
 --- Handles additional color names provided as a table or function.
 -- @param names_custom table|function|nil Additional color names to add.
 local function handle_names_custom(names_custom)
@@ -42,7 +58,6 @@ local function handle_names_custom(names_custom)
   end
 
   local extra_data = {}
-
   if type(names_custom) == "table" then
     extra_data = names_custom
   elseif type(names_custom) == "function" then
@@ -57,6 +72,10 @@ local function handle_names_custom(names_custom)
     end
   end
 
+  -- Add additional characters found in names_custom keys
+  local additonal_chars = extract_non_alphanum_keys(names_custom)
+  names_cache.color_trie:additional_chars(additonal_chars)
+
   for name, hex in pairs(extra_data) do
     if type(hex) == "string" then
       local normalized_hex = hex:gsub("^#", ""):gsub("%s", "")
@@ -73,6 +92,30 @@ local function handle_names_custom(names_custom)
   end
 end
 
+--- Handles Tailwind CSS colors and adds them to the Trie and map.
+local function handle_tailwind()
+  names_cache.color_trie:additional_chars("-")
+  local tailwind = require("colorizer.tailwind_colors")
+  for name, hex in pairs(tailwind.colors) do
+    for _, prefix in ipairs(tailwind.prefixes) do
+      add_color(prefix .. "-" .. name, hex)
+    end
+  end
+end
+
+--- Handles Vim's color map and adds colors to the Trie and map.
+local function handle_names()
+  for name, value in pairs(vim.api.nvim_get_color_map()) do
+    if not (names_cache.color_name_settings.strip_digits and name:match("%d+$")) then
+      local rgb_hex = tohex(value, 6)
+      add_color(name, rgb_hex)
+      if names_cache.color_name_settings.lowercase then
+        add_color(name:lower(), rgb_hex)
+      end
+    end
+  end
+end
+
 --- Populates the Trie and map with colors based on options.
 -- @param opts table Configuration options for color names and Tailwind CSS.
 local function populate_colors(opts)
@@ -82,25 +125,12 @@ local function populate_colors(opts)
 
   -- Add Vim's color map
   if opts.color_names then
-    for name, value in pairs(vim.api.nvim_get_color_map()) do
-      if not (names_cache.color_name_settings.strip_digits and name:match("%d+$")) then
-        local rgb_hex = tohex(value, 6)
-        add_color(name, rgb_hex)
-        if names_cache.color_name_settings.lowercase then
-          add_color(name:lower(), rgb_hex)
-        end
-      end
-    end
+    handle_names()
   end
 
   -- Add Tailwind colors
   if opts.tailwind then
-    local tailwind = require("colorizer.tailwind_colors")
-    for name, hex in pairs(tailwind.colors) do
-      for _, prefix in ipairs(tailwind.prefixes) do
-        add_color(prefix .. "-" .. name, hex)
-      end
-    end
+    handle_tailwind()
   end
   names_cache.tailwind_enabled = opts.tailwind
 
@@ -117,6 +147,7 @@ end
 -- @return number|nil, string|nil Length of match and hex value if found.
 function M.parser(line, i, opts)
   if not names_cache.color_trie or opts.tailwind ~= names_cache.tailwind_enabled then
+    --  TODO: 2024-12-21 - Ensure that this is not being called too many times
     populate_colors(opts)
   end
 
diff --git a/lua/colorizer/sass.lua b/lua/colorizer/sass.lua
index eb9d55a..4726c47 100644
--- a/lua/colorizer/sass.lua
+++ b/lua/colorizer/sass.lua
@@ -37,7 +37,7 @@ end
 ---@param i number: Index of line from where to start parsing
 ---@param bufnr number: Buffer number
 ---@return number|nil, string|nil
-function M.name_parser(line, i, bufnr)
+function M.parser(line, i, bufnr)
   local variable_name = line:match("^%$([%w_-]+)", i)
   if variable_name then
     local rgb_hex = state[bufnr].definitions_all[variable_name]
diff --git a/lua/colorizer/trie.lua b/lua/colorizer/trie.lua
index fbb7263..14571a4 100644
--- a/lua/colorizer/trie.lua
+++ b/lua/colorizer/trie.lua
@@ -54,12 +54,9 @@ end
 
 local total_char = 255
 local index_lookup = ffi.new("uint8_t[?]", total_char)
-local char_lookup = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz-"
+local char_lookup = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
 do
   local b = string.byte
-  local extra_char = {
-    [b("-")] = true,
-  }
   local byte = {
     ["0"] = b("0"),
     ["9"] = b("9"),
@@ -75,7 +72,6 @@ do
       index_lookup[i] = i - byte["A"] + 10
     elseif i >= byte["a"] and i <= byte["z"] then
       index_lookup[i] = i - byte["a"] + 10 + 26
-    elseif extra_char[i] then
     else
       index_lookup[i] = total_char
     end
@@ -247,6 +243,20 @@ local function trie_to_string(trie)
   return table.concat(print_trie_table(as_table), "\n")
 end
 
+local function trie_additional_chars(trie, chars)
+  if trie == nil or type(chars) ~= "string" then
+    return
+  end
+  for i = 1, #chars do
+    local char = chars:sub(i, i)
+    local char_byte = string.byte(char)
+    if index_lookup[char_byte] == total_char then
+      char_lookup = char_lookup .. char
+      index_lookup[char_byte] = total_char + 1
+    end
+  end
+end
+
 local Trie_mt = {
   __new = function(_, init)
     local trie = trie_create()
@@ -261,6 +271,7 @@ local Trie_mt = {
     longest_prefix = trie_longest_prefix,
     extend = trie_extend,
     destroy = trie_destroy,
+    additional_chars = trie_additional_chars,
   },
   __tostring = trie_to_string,
   __gc = trie_destroy,
diff --git a/lua/colorizer/utils.lua b/lua/colorizer/utils.lua
index 00e39c4..a9f73a5 100644
--- a/lua/colorizer/utils.lua
+++ b/lua/colorizer/utils.lua
@@ -68,6 +68,7 @@ end
 ---@param byte number The byte to check.
 ---@return boolean `true` if the byte is valid, otherwise `false`.
 function M.byte_is_valid_colorchar(byte)
+  --  TODO: 2024-12-21 - Is this check required?
   return M.byte_is_alphanumeric(byte) or byte == ("-"):byte()
 end
 
diff --git a/test/expect.lua b/test/expect.lua
index a02ab26..621fe5d 100644
--- a/test/expect.lua
+++ b/test/expect.lua
@@ -6,15 +6,11 @@ local opts = {
     lua = {
       names = true,
       names_custom = {
-        -- names = "#1F4770",
-        names = "#791497",
-        cool = "#3F3347",
-        lua = "#107d3c",
-        ["notcool"] = "#ee9240",
-        redgreen = "#970000",
-        asdf = "#234311",
-        eeee = "#112238",
-        -- lua
+        red_purple = "#017dac",
+        ["red=green"] = "#3700c2",
+        ["green@blue"] = "#e9e240",
+        ["green!blue"] = "#a9e042",
+        ["green!!blue"] = "#09e392",
       },
       -- names_custom = function()
       --   local colors = require("kanagawa.colors").setup()
@@ -25,11 +21,6 @@ local opts = {
   buftypes = { "*", "!prompt", "!popup" },
   user_commands = true,
   user_default_options = {
-    names_custom = {
-      -- names = "#1F4770",
-      names = "#1740F7",
-      lua = "#7407F1",
-    },
     names = false,
     RGB = true,
     RRGGBB = true,
@@ -69,6 +60,9 @@ Extra names:
 oniViolet oniViolet2 crystalBlue springViolet1 springViolet2 springBlue
 lightBlue waveAqua2
 
+Additional names with non-alphanumeric characters
+red_purple red=green green@blue green!blue green!!blue
+
 Hexadecimal:
 #RGB:
   #F0F