Refactor (#19)

- Refactor and clean up old code.
- Fix the trie printing (it's perfect now :o)
- Add the ability to search from a starting point. This can help
avoid allocations.
- Avoid an allocation in longest_prefix.
- Refactor and allow configuring of color_name_parser setup

1 files changed, 52 insertions, 34 deletions

diff --git a/lua/colorizer.lua b/lua/colorizer.lua
index f127b19..00bd949 100644
--- a/lua/colorizer.lua
+++ b/lua/colorizer.lua
@@ -11,20 +11,34 @@ local nvim_buf_get_lines = vim.api.nvim_buf_get_lines
 local nvim_get_current_buf = vim.api.nvim_get_current_buf
 local band, lshift, bor, tohex = bit.band, bit.lshift, bit.bor, bit.tohex
 local rshift = bit.rshift
-local floor = math.floor
+local floor, min, max = math.floor, math.min, math.max
 
 local COLOR_MAP
 local COLOR_TRIE
+local COLOR_NAME_MINLEN, COLOR_NAME_MAXLEN
+local COLOR_NAME_SETTINGS = {
+	lowercase = false;
+	strip_digits = false;
+}
 
 --- Setup the COLOR_MAP and COLOR_TRIE
 local function initialize_trie()
 	if not COLOR_TRIE then
-		COLOR_MAP = nvim.get_color_map()
+		COLOR_MAP = {}
 		COLOR_TRIE = Trie()
-
-		for k, v in pairs(COLOR_MAP) do
-			COLOR_MAP[k] = tohex(v, 6)
-			COLOR_TRIE:insert(k)
+		for k, v in pairs(nvim.get_color_map()) do
+			if not (COLOR_NAME_SETTINGS.strip_digits and k:match("%d+$")) then
+				COLOR_NAME_MINLEN = COLOR_NAME_MINLEN and min(#k, COLOR_NAME_MINLEN) or #k
+				COLOR_NAME_MAXLEN = COLOR_NAME_MAXLEN and max(#k, COLOR_NAME_MAXLEN) or #k
+				local rgb_hex = tohex(v, 6)
+				COLOR_MAP[k] = rgb_hex
+				COLOR_TRIE:insert(k)
+				if COLOR_NAME_SETTINGS.lowercase then
+					local lowercase = k:lower()
+					COLOR_MAP[lowercase] = rgb_hex
+					COLOR_TRIE:insert(lowercase)
+				end
+			end
 		end
 	end
 end
@@ -53,6 +67,15 @@ local DEFAULT_OPTIONS = {
 	mode     = 'background'; -- Set the display mode.
 }
 
+-- -- TODO use rgb as the return value from the matcher functions
+-- -- instead of the rgb_hex. Can be the highlight key as well
+-- -- when you shift it left 8 bits. Use the lower 8 bits for
+-- -- indicating which highlight mode to use.
+-- ffi.cdef [[
+-- typedef struct { uint8_t r, g, b; } colorizer_rgb;
+-- ]]
+-- local rgb_t = ffi.typeof 'colorizer_rgb'
+
 -- Create a lookup table where the bottom 4 bits are used to indicate the
 -- category and the top 4 bits are the hex value of the ASCII byte.
 local BYTE_CATEGORY = ffi.new 'uint8_t[256]'
@@ -144,11 +167,12 @@ local function hsl_to_rgb(h, s, l)
 	return 255*hue_to_rgb(p, q, h + 1/3), 255*hue_to_rgb(p, q, h), 255*hue_to_rgb(p, q, h - 1/3)
 end
 
-local function name_parser(line, i)
+local function color_name_parser(line, i)
 	if i > 1 and byte_is_alphanumeric(line:byte(i-1)) then
 		return
 	end
-	local prefix = COLOR_TRIE:longest_prefix(line:sub(i))
+	if #line < i + COLOR_NAME_MINLEN - 1 then return end
+	local prefix = COLOR_TRIE:longest_prefix(line, i)
 	if prefix then
 		-- Check if there is a letter here so as to disallow matching here.
 		-- Take the Blue out of Blueberry
@@ -174,7 +198,7 @@ local function rgb_hex_parser(line, i, minlen, maxlen)
 	local n = j + maxlen
 	local alpha
 	local v = 0
-	while j <= math.min(n, #line) do
+	while j <= min(n, #line) do
 		local b = line:byte(j)
 		if not byte_is_hex(b) then break end
 		if j - i >= 7 then
@@ -205,23 +229,22 @@ end
 -- Things like pumblend might be useful here.
 local css_fn = {}
 do
-	local css_rgb_fn_minimum_length = #'rgb(0,0,0)' - 1
-	local css_rgba_fn_minimum_length = #'rgba(0,0,0,0)' - 1
-	local css_hsl_fn_minimum_length = #'hsl(0,0%,0%)' - 1
-	local css_hsla_fn_minimum_length = #'hsla(0,0%,0%,0)' - 1
+	local CSS_RGB_FN_MINIMUM_LENGTH = #'rgb(0,0,0)' - 1
+	local CSS_RGBA_FN_MINIMUM_LENGTH = #'rgba(0,0,0,0)' - 1
+	local CSS_HSL_FN_MINIMUM_LENGTH = #'hsl(0,0%,0%)' - 1
+	local CSS_HSLA_FN_MINIMUM_LENGTH = #'hsla(0,0%,0%,0)' - 1
 	function css_fn.rgb(line, i)
-		if #line < i + css_rgb_fn_minimum_length then return end
+		if #line < i + CSS_RGB_FN_MINIMUM_LENGTH then return end
 		local r, g, b, match_end = line:sub(i):match("^rgb%(%s*(%d+%%?)%s*,%s*(%d+%%?)%s*,%s*(%d+%%?)%s*%)()")
 		if not match_end then return end
 		r = percent_or_hex(r) if not r then return end
 		g = percent_or_hex(g) if not g then return end
 		b = percent_or_hex(b) if not b then return end
-		local rgb_hex = ("%02x%02x%02x"):format(r,g,b)
-		if #rgb_hex ~= 6 then return end
+		local rgb_hex = tohex(bor(lshift(r, 16), lshift(g, 8), b), 6)
 		return match_end - 1, rgb_hex
 	end
 	function css_fn.hsl(line, i)
-		if #line < i + css_hsl_fn_minimum_length then return end
+		if #line < i + CSS_HSL_FN_MINIMUM_LENGTH then return end
 		local h, s, l, match_end = line:sub(i):match("^hsl%(%s*(%d+)%s*,%s*(%d+)%%%s*,%s*(%d+)%%%s*%)()")
 		if not match_end then return end
 		h = tonumber(h) if h > 360 then return end
@@ -229,24 +252,22 @@ do
 		l = tonumber(l) if l > 100 then return end
 		local r, g, b = hsl_to_rgb(h/360, s/100, l/100)
 		if r == nil or g == nil or b == nil then return end
-		local rgb_hex = ("%02x%02x%02x"):format(floor(r), floor(g), floor(b))
-		if #rgb_hex ~= 6 then return end
+		local rgb_hex = tohex(bor(lshift(floor(r), 16), lshift(floor(g), 8), floor(b)), 6)
 		return match_end - 1, rgb_hex
 	end
 	function css_fn.rgba(line, i)
-		if #line < i + css_rgba_fn_minimum_length then return end
+		if #line < i + CSS_RGBA_FN_MINIMUM_LENGTH then return end
 		local r, g, b, a, match_end = line:sub(i):match("^rgba%(%s*(%d+%%?)%s*,%s*(%d+%%?)%s*,%s*(%d+%%?)%s*,%s*([.%d]+)%s*%)()")
 		if not match_end then return end
 		a = tonumber(a) if not a or a > 1 then return end
 		r = percent_or_hex(r) if not r then return end
 		g = percent_or_hex(g) if not g then return end
 		b = percent_or_hex(b) if not b then return end
-		local rgb_hex = ("%02x%02x%02x"):format(floor(r*a), floor(g*a), floor(b*a))
-		if #rgb_hex ~= 6 then return end
+		local rgb_hex = tohex(bor(lshift(floor(r*a), 16), lshift(floor(g*a), 8), floor(b*a)), 6)
 		return match_end - 1, rgb_hex
 	end
 	function css_fn.hsla(line, i)
-		if #line < i + css_hsla_fn_minimum_length then return end
+		if #line < i + CSS_HSLA_FN_MINIMUM_LENGTH then return end
 		local h, s, l, a, match_end = line:sub(i):match("^hsla%(%s*(%d+)%s*,%s*(%d+)%%%s*,%s*(%d+)%%%s*,%s*([.%d]+)%s*%)()")
 		if not match_end then return end
 		a = tonumber(a) if not a or a > 1 then return end
@@ -255,8 +276,7 @@ do
 		l = tonumber(l) if l > 100 then return end
 		local r, g, b = hsl_to_rgb(h/360, s/100, l/100)
 		if r == nil or g == nil or b == nil then return end
-		local rgb_hex = ("%02x%02x%02x"):format(floor(r*a), floor(g*a), floor(b*a))
-		if #rgb_hex ~= 6 then return end
+		local rgb_hex = tohex(bor(lshift(floor(r*a), 16), lshift(floor(g*a), 8), floor(b*a)), 6)
 		return match_end - 1, rgb_hex
 	end
 end
@@ -377,15 +397,15 @@ local function make_matcher(options)
 
 	local loop_matchers = {}
 	if enable_names then
-		table.insert(loop_matchers, name_parser)
+		table.insert(loop_matchers, color_name_parser)
 	end
 	do
 		local valid_lengths = {[3] = enable_RGB, [6] = enable_RRGGBB, [8] = enable_RRGGBBAA}
 		local minlen, maxlen
 		for k, v in pairs(valid_lengths) do
 			if v then
-				minlen = math.min(k, minlen or 99)
-				maxlen = math.max(k, maxlen or 0)
+				minlen = minlen and min(k, minlen) or k
+				maxlen = maxlen and max(k, maxlen) or k
 			end
 		end
 		if minlen then
@@ -535,19 +555,18 @@ end
 -- @param[opt={'*'}] filetypes A table/array of filetypes to selectively enable and/or customize. By default, enables all filetypes.
 -- @tparam[opt] {[string]=string} default_options Default options to apply for the filetypes enable.
 -- @usage require'colorizer'.setup()
-local function setup(filetypes, default_options)
+local function setup(filetypes, user_default_options)
 	if not nvim.o.termguicolors then
 		nvim.err_writeln("&termguicolors must be set")
 		return
 	end
-	initialize_trie()
 	FILETYPE_OPTIONS = {}
 	SETUP_SETTINGS = {
 		exclusions = {};
-		default_options = merge(DEFAULT_OPTIONS, default_options or {});
+		default_options = merge(DEFAULT_OPTIONS, user_default_options or {});
 	}
-	-- This is just in case I accidentally reference the wrong thing here.
-	default_options = SETUP_SETTINGS.default_options
+	-- Initialize this AFTER setting COLOR_NAME_SETTINGS
+	initialize_trie()
 	function COLORIZER_SETUP_HOOK()
 		local filetype = nvim.bo.filetype
 		if SETUP_SETTINGS.exclusions[filetype] then
@@ -558,7 +577,6 @@ local function setup(filetypes, default_options)
 	end
 	nvim.ex.augroup("ColorizerSetup")
 	nvim.ex.autocmd_()
-	-- nvim.ex.autocmd("VimEnter * lua COLORIZER_SETUP_HOOK()")
 	if not filetypes then
 		nvim.ex.autocmd("FileType * lua COLORIZER_SETUP_HOOK()")
 	else