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

2 files changed, 82 insertions, 97 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
diff --git a/lua/trie.lua b/lua/trie.lua
index b665f76..21794ef 100644
--- a/lua/trie.lua
+++ b/lua/trie.lua
@@ -34,6 +34,19 @@ local function trie_create()
 	return ffi.cast(Trie_ptr_t, ptr)
 end
 
+local function trie_destroy(trie)
+	if trie == nil then
+		return
+	end
+	for i = 0, 61 do
+		local child = trie.character[i]
+		if child ~= nil then
+			trie_destroy(child)
+		end
+	end
+	ffi.C.free(trie)
+end
+
 local INDEX_LOOKUP_TABLE = ffi.new 'uint8_t[256]'
 local CHAR_LOOKUP_TABLE = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
 do
@@ -168,46 +181,44 @@ local function print_trie_table(s)
 	end
 	local lines = {}
 	for _, child in ipairs(s.children) do
-		local child_lines = print_trie_table(child)
+		local child_lines = print_trie_table(child, thicc)
 		for _, child_line in ipairs(child_lines) do
 			table.insert(lines, child_line)
 		end
 	end
 	local child_count = 0
-	for i, v in ipairs(lines) do
-		if v:match("^[%w%d]") then
+	for i, line in ipairs(lines) do
+		local line_parts = {}
+		if line:match("^%w") then
 			child_count = child_count + 1
 			if i == 1 then
-				lines[i] = mark.."─"..v
+				line_parts = {mark}
 			elseif i == #lines or child_count == #s.children then
-				lines[i] = "└──"..v
+				line_parts = {"└─"}
 			else
-				lines[i] = "├──"..v
+				line_parts = {"├─"}
 			end
 		else
 			if i == 1 then
-				lines[i] = mark.."─"..v
+				line_parts = {mark}
 			elseif #s.children > 1 and child_count ~= #s.children then
-				lines[i] = "│  "..v
+				line_parts = {"│ "}
 			else
-				lines[i] = "   "..v
+				line_parts = {"  "}
 			end
 		end
+		table.insert(line_parts, line)
+		lines[i] = table.concat(line_parts)
 	end
 	return lines
 end
 
-local function trie_destroy(trie)
+local function trie_to_string(trie)
 	if trie == nil then
-		return
-	end
-	for i = 0, 61 do
-		local child = trie.character[i]
-		if child ~= nil then
-			trie_destroy(child)
-		end
+		return 'nil'
 	end
-	ffi.C.free(trie)
+	local as_table = trie_as_table(trie)
+	return table.concat(print_trie_table(as_table), '\n')
 end
 
 local Trie_mt = {
@@ -224,52 +235,8 @@ local Trie_mt = {
 		longest_prefix = trie_longest_prefix;
 		extend = trie_extend;
 	};
-	__tostring = function(trie)
-		if trie == nil then
-			return 'nil'
-		end
-		return table.concat(print_trie_table(trie_as_table(trie)), '\n')
-	end;
+	__tostring = trie_to_string;
 	__gc = trie_destroy;
 }
 
 return ffi.metatype('struct Trie', Trie_mt)
-
--- local tests = {
--- 	"cat";
--- 	"car";
--- 	"celtic";
--- 	"carb";
--- 	"carb0";
--- 	"CART0";
--- 	"CaRT0";
--- 	"Cart0";
--- 	"931";
--- 	"191";
--- 	"121";
--- 	"cardio";
--- 	"call";
--- 	"calcium";
--- 	"calciur";
--- 	"carry";
--- 	"dog";
--- 	"catdog";
--- }
--- local trie = Trie()
--- for i, v in ipairs(tests) do
--- 	trie:insert(v)
--- end
-
--- print(trie)
--- print(trie.character[0])
--- print("catdo", trie:longest_prefix("catdo"))
--- print("catastrophic", trie:longest_prefix("catastrophic"))
-
--- local COLOR_MAP = vim.api.nvim_get_color_map()
--- local start = os.clock()
--- for k, v in pairs(COLOR_MAP) do
--- 	insert(trie, k)
--- end
--- print(os.clock() - start)
-
--- print(table.concat(print_trie_table(trie_as_table(trie)), '\n'))