Add color_picker and refactor.

1 files changed, 313 insertions, 25 deletions

diff --git a/lua/colorizer.lua b/lua/colorizer.lua
index e47e079..fbf2e11 100644
--- a/lua/colorizer.lua
+++ b/lua/colorizer.lua
@@ -5,6 +5,7 @@ local Trie = require 'colorizer/trie'
 local bit = require 'bit'
 local ffi = require 'ffi'
 
+local vim = vim
 local nvim_buf_add_highlight = vim.api.nvim_buf_add_highlight
 local nvim_buf_clear_namespace = vim.api.nvim_buf_clear_namespace
 local nvim_buf_get_lines = vim.api.nvim_buf_get_lines
@@ -167,6 +168,33 @@ 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 rgb_to_hsl(r, g, b)
+  r = r / 255
+  g = g / 255
+  b = b / 255
+  local c_max = max(r, g, b)
+  local c_min = min(r, g, b)
+  local chroma = c_max - c_min
+  if chroma == 0 then
+    return 0, 0, 0
+  end
+  local l = (c_max + c_min) / 2
+  local s = chroma / (1 - math.abs(2*l-1))
+  local h
+  if c_max == r then
+    h = ((g - b) / chroma) % 6
+  elseif c_max == g then
+    h = (b - r) / chroma + 2;
+  elseif c_max == b then
+    h = (r - g) / chroma + 4;
+  end
+
+  h = floor(h * 60)
+  s = floor(s * 100)
+  l = floor(l * 100)
+  return h, s, l
+end
+
 local function color_name_parser(line, i)
 	if i > 1 and byte_is_alphanumeric(line:byte(i-1)) then
 		return
@@ -185,6 +213,24 @@ local function color_name_parser(line, i)
 	end
 end
 
+-- Converts a number to its rgb parts
+local function num_to_rgb(n)
+  n = tonumber(n)
+  return band(rshift(n, 16), 0xFF),
+      band(rshift(n, 8), 0xFF),
+      band(n, 0xFF)
+end
+
+-- Converts a number to its rgb parts
+local function rgb_to_num(r, g, b)
+  return bor(lshift(band(r, 0xFF), 16), lshift(band(g, 0xFF), 8), band(b, 0xFF))
+end
+
+-- Converts a number to its rgb parts
+local function rgb_to_hex(r, g, b)
+  return tohex(rgb_to_num(r, g, b), 6)
+end
+
 local b_hash = ("#"):byte()
 local function rgb_hex_parser(line, i, minlen, maxlen)
 	if i > 1 and byte_is_alphanumeric(line:byte(i-1)) then
@@ -215,13 +261,17 @@ local function rgb_hex_parser(line, i, minlen, maxlen)
 	if length ~= 4 and length ~= 7 and length ~= 9 then return end
 	if alpha then
 		alpha = tonumber(alpha)/255
-		local r = floor(band(v, 0xFF)*alpha)
-		local g = floor(band(rshift(v, 8), 0xFF)*alpha)
-		local b = floor(band(rshift(v, 16), 0xFF)*alpha)
-		v = bor(lshift(r, 16), lshift(g, 8), b)
-		return 9, tohex(v, 6)
-	end
-	return length, line:sub(i+1, i+length-1)
+    local r, g, b = num_to_rgb(v)
+		return 9, rgb_to_hex(floor(r*alpha), floor(g*alpha), floor(b*alpha))
+	end
+  local rgb_hex = line:sub(i+1, i+length-1)
+  if length == 4 then
+    local x = tonumber(rgb_hex, 16)
+    local r,g,b = band(rshift(x, 8), 0xF), band(rshift(x, 4), 0xF), band(x, 0xF)
+    r, g, b = bor(r, lshift(r, 4)), bor(g, lshift(g, 4)), bor(b, lshift(b, 4))
+    return 4, rgb_to_hex(r,g,b)
+  end
+	return 7, rgb_hex
 end
 
 -- TODO consider removing the regexes here
@@ -240,8 +290,7 @@ do
 		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 = tohex(bor(lshift(r, 16), lshift(g, 8), b), 6)
-		return match_end - 1, rgb_hex
+		return match_end - 1, rgb_to_hex(r, g, b)
 	end
 	function css_fn.hsl(line, i)
 		if #line < i + CSS_HSL_FN_MINIMUM_LENGTH then return end
@@ -252,8 +301,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 = tohex(bor(lshift(floor(r), 16), lshift(floor(g), 8), floor(b)), 6)
-		return match_end - 1, rgb_hex
+		return match_end - 1, rgb_to_hex(floor(r), floor(g), floor(b))
 	end
 	function css_fn.rgba(line, i)
 		if #line < i + CSS_RGBA_FN_MINIMUM_LENGTH then return end
@@ -263,8 +311,7 @@ do
 		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 = tohex(bor(lshift(floor(r*a), 16), lshift(floor(g*a), 8), floor(b*a)), 6)
-		return match_end - 1, rgb_hex
+		return match_end - 1, rgb_to_hex(floor(r*a), floor(g*a), floor(b*a))
 	end
 	function css_fn.hsla(line, i)
 		if #line < i + CSS_HSLA_FN_MINIMUM_LENGTH then return end
@@ -276,8 +323,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 = tohex(bor(lshift(floor(r*a), 16), lshift(floor(g*a), 8), floor(b*a)), 6)
-		return match_end - 1, rgb_hex
+		return match_end - 1, rgb_to_hex(floor(r*a), floor(g*a), floor(b*a))
 	end
 end
 local css_function_parser, rgb_function_parser, hsl_function_parser
@@ -340,24 +386,17 @@ local function create_highlight(rgb_hex, options)
 	local mode = options.mode or 'background'
 	-- TODO validate rgb format?
 	rgb_hex = rgb_hex:lower()
-	local cache_key = table.concat({HIGHLIGHT_MODE_NAMES[mode], rgb_hex}, "_")
+	local cache_key = HIGHLIGHT_MODE_NAMES[mode].."_"..rgb_hex
 	local highlight_name = HIGHLIGHT_CACHE[cache_key]
 	-- Look up in our cache.
 	if not highlight_name then
-		if #rgb_hex == 3 then
-			rgb_hex = table.concat {
-				rgb_hex:sub(1,1):rep(2);
-				rgb_hex:sub(2,2):rep(2);
-				rgb_hex:sub(3,3):rep(2);
-			}
-		end
 		-- Create the highlight
 		highlight_name = make_highlight_name(rgb_hex, mode)
 		if mode == 'foreground' then
 			nvim.ex.highlight(highlight_name, "guifg=#"..rgb_hex)
 		else
-			local r, g, b = rgb_hex:sub(1,2), rgb_hex:sub(3,4), rgb_hex:sub(5,6)
-			r, g, b = tonumber(r,16), tonumber(g,16), tonumber(b,16)
+      -- Guess the foreground color based on the background color's brightness.
+			local r, g, b = num_to_rgb(tonumber(rgb_hex, 16))
 			local fg_color
 			if color_is_bright(r,g,b) then
 				fg_color = "Black"
@@ -372,6 +411,11 @@ local function create_highlight(rgb_hex, options)
 end
 
 local MATCHER_CACHE = {}
+-- Return a function which is called a "loop parse function," meaning that it
+-- can be used in a loop to check if there is a valid match in a string at the
+-- specified index for any known color functions as specified by {options}.
+--
+-- Returns: fn(line: string, index: int) -> (length, rgb_hex): (int, string)
 local function make_matcher(options)
 	local enable_names    = options.css or options.names
 	local enable_RGB      = options.css or options.RGB
@@ -633,9 +677,252 @@ local function get_buffer_options(buf)
 	return merge({}, BUFFER_OPTIONS[buf])
 end
 
+local partial_bar = "▅"
+local full_bar = "█"
+local empty_bar = "▁"
+
+-- starting: string
+-- finish: function(r,g,b)
+--
+-- returns: bufnr, winnr
+local function color_picker(starting, on_change)
+  if _picker then
+    print("There is already a color picker running.")
+    return
+  end
+  assert(type(on_change) == 'function')
+
+  local api = vim.api
+  local bufnr = api.nvim_create_buf(false, true)
+
+  local function progress_bar(x, m, fill)
+    m = m - 1
+    local pos = floor(x * m)
+    local cursor
+    if x == 0 then
+      cursor = empty_bar
+    elseif floor(x * m) == x * m then
+      cursor = full_bar
+    else
+      cursor = fill and partial_bar or full_bar
+    end
+    local pre = fill and full_bar or empty_bar
+    return pre:rep(pos)..cursor..empty_bar:rep(m - pos)
+  end
+
+  local function render_bars(focus, w, rgb, values, limits, styles)
+    local ns = DEFAULT_NAMESPACE
+    api.nvim_buf_clear_namespace(bufnr, ns, 0, -1)
+    local lines = {
+      "#"..rgb_to_hex(unpack(rgb));
+    }
+    for i = 1, #values do
+      local v = values[i]
+      local m = limits[i] or 1
+      local s = styles[i]
+      lines[#lines+1] = progress_bar(v/m, w, s).." = "..v
+    end
+    api.nvim_buf_set_lines(bufnr, 0, -1, false, lines)
+    api.nvim_buf_add_highlight(bufnr, ns, 'Underlined', focus + 1, 0, w*#empty_bar)
+  end
+
+  local bar_width = 10
+
+  local function clamp(x, x0, x1)
+    return min(max(x, x0), x1)
+  end
+
+  local mode = 0
+
+  local rgb = {
+    focus = 0;
+    values = {0, 0, 0};
+    limits = {255, 255, 255};
+    styles = {false, false, false};
+  }
+  local hsl = {
+    focus = 0;
+    values = {0, 0, 0};
+    limits = {360, 100, 100};
+    styles = {false, true, true};
+  }
+
+  function rgb.init(r,g,b)
+    rgb.values = {floor(r),floor(g),floor(b)}
+  end
+  function rgb.rgb()
+    return rgb.values
+  end
+
+  function hsl.init(r,g,b)
+    hsl.values = {rgb_to_hsl(r,g,b)}
+    -- hsl.h, hsl.s, hsl.l = rgb_to_hsl(r,g,b)
+    hsl.focus = hsl.focus or 0
+  end
+  function hsl.rgb()
+    local h,s,l = unpack(hsl.values)
+    return {hsl_to_rgb(h/360, s/100, l/100)}
+  end
+
+  local modes = { rgb; hsl; }
+
+  function _picker(S)
+    local cmode = modes[mode + 1]
+    local changed = false
+    if S.focus then
+      cmode.focus = clamp(cmode.focus+S.focus, 0, 2)
+    elseif S.value then
+      local i = cmode.focus+1
+      local values = cmode.values
+      values[i] = clamp(values[i]+S.value, 0, cmode.limits[i])
+      changed = true
+    elseif S.mode then
+      local values = cmode.rgb()
+      mode = (mode + S.mode) % 2
+      cmode = modes[mode + 1]
+      cmode.init(unpack(values))
+      changed = true
+    end
+    local rgbvals = cmode.rgb()
+    if changed then
+      on_change(rgbvals)
+    end
+    render_bars(cmode.focus, bar_width, rgbvals, cmode.values, cmode.limits, cmode.styles)
+    -- render_bars(cmode.focus, bar_width, cmode.rgb(), cmode.values, cmode.limits, cmode.styles)
+  end
+
+  api.nvim_buf_set_keymap(bufnr, 'n', 'j', '<cmd>lua _picker{focus=1}<cr>', {noremap=true})
+  api.nvim_buf_set_keymap(bufnr, 'n', 'k', '<cmd>lua _picker{focus=-1}<cr>', {noremap=true})
+  api.nvim_buf_set_keymap(bufnr, 'n', 'l', '<cmd>lua _picker{value=1}<cr>', {noremap=true})
+  api.nvim_buf_set_keymap(bufnr, 'n', 'h', '<cmd>lua _picker{value=-1}<cr>', {noremap=true})
+  api.nvim_buf_set_keymap(bufnr, 'n', 'L', '<cmd>lua _picker{value=10}<cr>', {noremap=true})
+  api.nvim_buf_set_keymap(bufnr, 'n', 'H', '<cmd>lua _picker{value=-10}<cr>', {noremap=true})
+  api.nvim_buf_set_keymap(bufnr, 'n', '<TAB>', '<cmd>lua _picker{mode=1}<cr>', {noremap=true})
+
+  api.nvim_buf_attach(bufnr, false, {
+    on_detach = function()
+      on_change(modes[mode+1].rgb(), true)
+      _picker = nil
+    end
+  })
+
+  if starting then
+    assert(type(starting) == 'string')
+    -- Make a matcher which works for all inputs.
+    local matcher = make_matcher{css=true}
+    local length, rgb_hex = matcher(starting, 1)
+    if length then
+      modes[mode+1].init(num_to_rgb(tonumber(rgb_hex, 16)))
+    else
+      print("Invalid starting color:", starting)
+      modes[mode+1].init(0, 0, 0)
+    end
+  else
+    modes[mode+1].init(0, 0, 0)
+  end
+
+  attach_to_buffer(bufnr)
+  api.nvim_buf_set_option(bufnr, 'undolevels', -1)
+  -- TODO(ashkan): skip sending first on_change?
+  _picker{}
+
+  local winnr = api.nvim_open_win(bufnr, true, {
+    style = 'minimal';
+    -- anchor = 'NW';
+    width = bar_width + 10;
+    relative = 'cursor';
+    row = 0; col = 0;
+    height = 4;
+  })
+  return bufnr, winnr
+end
+
+-- TODO(ashkan): Match the replacement type to the type of the input.
+local function color_picker_on_cursor(config)
+  config = config or {}
+  assert(type(config) == 'table')
+  local match_format = not config.rgb_hex
+
+  local api = vim.api
+  local bufnr = api.nvim_get_current_buf()
+  local pos = api.nvim_win_get_cursor(0)
+  local row, col = unpack(pos)
+  local line = api.nvim_get_current_line()
+  local matcher = make_matcher{css=true}
+  local start, length, rgb_hex
+
+  -- TODO(ashkan): How much should I backpedal? Is too much a problem? I don't
+  -- think it could be since the color must contain the cursor.
+  for i = col+1, max(col-50, 1), -1 do
+    local l, hex = matcher(line, i)
+    -- Check that col is bounded by i and i+l
+    if l and (i + l) > col+1 then
+      start, length, rgb_hex = i, l, hex
+      break
+    end
+  end
+
+  local function startswith(s, n)
+    return s:sub(1, #n) == n
+  end
+
+  -- TODO(ashkan): make insertion on no color found configurable.
+  -- Currently, it doesn't insert unless you modify something, which is pretty
+  -- nice.
+  start = start or col+1
+  local prefix = line:sub(1, start-1)
+  local suffix = line:sub(start+(length or 0))
+  local matched = line:sub(start, start+length)
+  local formatter = function(rgb)
+    return "#"..rgb_to_hex(unpack(rgb))
+  end
+  if match_format then
+    -- TODO(ashkan): make matching the result optional?
+    if startswith(matched, "rgba") then
+      -- TODO(ashkan): support alpha?
+      formatter = function(rgb)
+        return string.format("rgba(%d, %d, %d, 1)", unpack(rgb))
+      end
+    elseif startswith(matched, "rgb") then
+      formatter = function(rgb)
+        return string.format("rgb(%d, %d, %d)", unpack(rgb))
+      end
+    elseif startswith(matched, "hsla") then
+      formatter = function(rgb)
+        return string.format("hsla(%d, %d%%, %d%%, 1)", rgb_to_hsl(unpack(rgb)))
+      end
+    elseif startswith(matched, "hsl") then
+      formatter = function(rgb)
+        return string.format("hsl(%d, %d%%, %d%%)", rgb_to_hsl(unpack(rgb)))
+      end
+      -- elseif startswith(matched, "#") and length == 4 then
+      -- elseif startswith(matched, "#") and length == 7 then
+      -- else
+    end
+  end
+  -- Disable live previews on long lines.
+  -- TODO(ashkan): enable this when you can to nvim_buf_set_text instead of set_lines.
+  -- TODO(ashkan): is 200 a fair number?
+  if #line > 200 then
+    return color_picker(rgb_hex and "#"..rgb_hex, vim.schedule_wrap(function(rgb, is_last)
+      if is_last then
+        api.nvim_buf_set_lines(bufnr, row-1, row, true, {prefix..formatter(rgb)..suffix})
+      end
+    end))
+  end
+  return color_picker(rgb_hex and "#"..rgb_hex, vim.schedule_wrap(function(rgb, is_last)
+    -- Since we're modifying it perpetually, we don't need is_last, and this
+    -- avoids modifying when nothing has changed.
+    if is_last then return end
+    api.nvim_buf_set_lines(bufnr, row-1, row, true, {prefix..formatter(rgb)..suffix})
+  end))
+end
+
 --- @export
 return {
 	DEFAULT_NAMESPACE = DEFAULT_NAMESPACE;
+  color_picker = color_picker;
+  color_picker_on_cursor = color_picker_on_cursor;
 	setup = setup;
 	is_buffer_attached = is_buffer_attached;
 	attach_to_buffer = attach_to_buffer;
@@ -643,5 +930,6 @@ return {
 	highlight_buffer = highlight_buffer;
 	reload_all_buffers = reload_all_buffers;
 	get_buffer_options = get_buffer_options;
+  create_highlight = create_highlight;
 }