From f0c8b75cc423ff0decfdf8ebba0dc7c3015c9aac Mon Sep 17 00:00:00 2001 From: Ashkan Kiani Date: Thu, 17 Oct 2019 15:15:32 -0700 Subject: Initial commit. --- lua/colorizer.lua | 247 +++++++++++++++++++++++++++++++++++++++++++ lua/trie.lua | 306 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 553 insertions(+) create mode 100644 lua/colorizer.lua create mode 100644 lua/trie.lua diff --git a/lua/colorizer.lua b/lua/colorizer.lua new file mode 100644 index 0000000..03d2bbc --- /dev/null +++ b/lua/colorizer.lua @@ -0,0 +1,247 @@ +--- Highlights terminal CSI ANSI color codes. +-- @module terminal +local nvim = require 'nvim' +local Trie = require 'trie' + +--- Default namespace used in `highlight_buffer` and `attach_to_buffer`. +-- The name is "terminal_highlight" +-- @see highlight_buffer +-- @see attach_to_buffer +local DEFAULT_NAMESPACE = nvim.create_namespace 'colorizer' + +local COLOR_MAP +local COLOR_TRIE + +local function initialize_trie() + if not COLOR_TRIE then + COLOR_MAP = nvim.get_color_map() + COLOR_TRIE = Trie() + + for k in pairs(COLOR_MAP) do + COLOR_TRIE:insert(k) + end + end +end + +local b_a = string.byte('a') +local b_z = string.byte('z') +local b_A = string.byte('A') +local b_Z = string.byte('Z') +local b_0 = string.byte('0') +local b_9 = string.byte('9') +local b_f = string.byte('f') +local b_hash = string.byte('#') +local bnot = bit.bnot +local band, bor, bxor = bit.band, bit.bor, bit.bxor +local lshift, rshift, rol = bit.lshift, bit.rshift, bit.rol + +-- TODO use lookup table? +local function byte_is_hex(byte) + byte = bor(byte, 0x20) + return (byte >= b_0 and byte <= b_9) or (byte >= b_a and byte <= b_f) +end + +--- Determine whether to use black or white text +-- Ref: https://stackoverflow.com/a/1855903/837964 +-- https://stackoverflow.com/questions/596216/formula-to-determine-brightness-of-rgb-color +local function color_is_bright(r, g, b) + -- Counting the perceptive luminance - human eye favors green color + local luminance = (0.299*r + 0.587*g + 0.114*b)/255 + if luminance > 0.5 then + return true -- Bright colors, black font + else + return false -- Dark colors, white font + end +end + +local HIGHLIGHT_NAME_PREFIX = "colorizer_" +--- Make a deterministic name for a highlight given these attributes +local function make_highlight_name(rgb) + return table.concat {HIGHLIGHT_NAME_PREFIX, rgb} +end + +local highlight_cache = {} + +-- Ref: https://stackoverflow.com/questions/1252539/most-efficient-way-to-determine-if-a-lua-table-is-empty-contains-no-entries +local function table_is_empty(t) + return next(t) == nil +end + +local function create_highlight(rgb_hex, options) + -- TODO validate rgb format? + local highlight_name = highlight_cache[rgb_hex] + -- Look up in our cache. + if not highlight_name then + -- Create the highlight + highlight_name = make_highlight_name(rgb_hex) + if options.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) + local fg_color + if color_is_bright(r,g,b) then + fg_color = "Black" + else + fg_color = "White" + end + nvim.ex.highlight(highlight_name, "guifg="..fg_color, "guibg=#"..rgb_hex) + end + highlight_cache[rgb_hex] = highlight_name + end + return highlight_name +end + +--- Highlight a region in a buffer from the attributes specified +local function highlight_region(buf, ns, highlight_name, + region_line_start, region_byte_start, + region_line_end, region_byte_end) + -- TODO should I bother with highlighting normal regions? + if region_line_start == region_line_end then + nvim.buf_add_highlight(buf, ns, highlight_name, region_line_start, region_byte_start, region_byte_end) + else + nvim.buf_add_highlight(buf, ns, highlight_name, region_line_start, region_byte_start, -1) + for linenum = region_line_start + 1, region_line_end - 1 do + nvim.buf_add_highlight(buf, ns, highlight_name, linenum, 0, -1) + end + nvim.buf_add_highlight(buf, ns, highlight_name, region_line_end, 0, region_byte_end) + end +end + +--[[-- Highlight the buffer region. +Highlight starting from `line_start` (0-indexed) for each line described by `lines` in the +buffer `buf` and attach it to the namespace `ns`. + +@tparam integer buf buffer id. +@tparam[opt=DEFAULT_NAMESPACE] integer ns the namespace id. Create it with `vim.api.create_namespace` +@tparam {string,...} lines the lines to highlight from the buffer. +@tparam integer line_start should be 0-indexed +]] +local function highlight_buffer(buf, ns, lines, line_start, options) + options = options or {} + -- TODO do I have to put this here? + initialize_trie() + ns = ns or DEFAULT_NAMESPACE + for current_linenum, line in ipairs(lines) do + -- @todo it's possible to skip processing the new code if the attributes hasn't changed. + current_linenum = current_linenum - 1 + line_start + local i = 1 + while i < #line do + local byte = line:byte(i) + if byte == b_hash then + i = i + 1 + if #line >= i + 5 then + local invalid = false + for n = i, i+5 do + byte = line:byte(n) + if not byte_is_hex(byte) then + invalid = true + break + end + end + if not invalid then + local rgb_hex = line:sub(i, i+5) + -- TODO figure out proper background/foreground + local highlight_name = create_highlight(rgb_hex, options) + -- Subtract one because 0-indexed, subtract another 1 for the '#' + nvim.buf_add_highlight(buf, ns, highlight_name, current_linenum, i-1-1, i+6-1) + i = i + 5 + end + end + else + -- TODO skip if the remaining length is less than the shortest length + -- of an entry in our trie. + local prefix = COLOR_TRIE:longest_prefix(line:sub(i)) + if prefix then + local rgb = COLOR_MAP[prefix] + local rgb_hex = bit.tohex(rgb):sub(-6) + -- TODO figure out proper background/foreground + local highlight_name = create_highlight(rgb_hex, options) + nvim.buf_add_highlight(buf, ns, highlight_name, current_linenum, i-1, i+#prefix-1) + i = i + #prefix + else + i = i + 1 + end + end + end + end +end + +--- Attach to a buffer and continuously highlight changes. +-- @tparam[opt=0] integer buf A value of 0 implies the current buffer. +-- @see highlight_buffer +local function attach_to_buffer(buf, options) + local ns = DEFAULT_NAMESPACE + if buf == 0 or buf == nil then + buf = nvim.get_current_buf() + end + -- Already attached. + if pcall(vim.api.nvim_buf_get_var, buf, "colorizer_attached") then + return + end + nvim.buf_set_var(buf, "colorizer_attached", true) + do + nvim.buf_clear_namespace(buf, ns, 0, -1) + local lines = nvim.buf_get_lines(buf, 0, -1, true) + highlight_buffer(buf, ns, lines, 0, options) + end + -- send_buffer: true doesn't actually do anything in Lua (yet) + nvim.buf_attach(buf, false, { + on_lines = function(event_type, buf, changed_tick, firstline, lastline, new_lastline) + nvim.buf_clear_namespace(buf, ns, firstline, new_lastline) + local lines = nvim.buf_get_lines(buf, firstline, new_lastline, true) + highlight_buffer(buf, ns, lines, firstline, options) + end; + }) +end + +--- Easy to use function if you want the full setup without fine grained control. +-- Establishes an autocmd for `FileType terminal` +-- @usage require'colorizer'.setup() +local function auto_setup(filetypes, default_options) + if not nvim.o.termguicolors then + nvim.err_writeln("&termguicolors must be set") + return + end + initialize_trie() + local filetype_options = {} + function COLORIZER_SETUP_HOOK() + local filetype = nvim.bo.filetype + local options = filetype_options[filetype] or default_options + attach_to_buffer(nvim.get_current_buf(), options) + end + nvim.ex.augroup("ColorizerSetup") + nvim.ex.autocmd_() + if not filetypes then + nvim.ex.autocmd("FileType * lua COLORIZER_SETUP_HOOK()") + else + for k, v in pairs(filetypes) do + local filetype + local options = default_options or {} + if type(k) == 'string' then + filetype = k + if type(v) ~= 'table' then + nvim.err_writeln("colorizer: Invalid option type for filetype "..filetype) + else + options = vim.tbl_extend("keep", v, default_options) + end + else + filetype = v + end + filetype_options[filetype] = options + -- TODO What's the right mode for this? BufEnter? + nvim.ex.autocmd("FileType", filetype, "lua COLORIZER_SETUP_HOOK(%s)") + end + end + nvim.ex.augroup("END") +end + +--- @export +return { + DEFAULT_NAMESPACE = DEFAULT_NAMESPACE; + setup = auto_setup; + attach_to_buffer = attach_to_buffer; + highlight_buffer = highlight_buffer; + initialize = initialize_trie; +} + diff --git a/lua/trie.lua b/lua/trie.lua new file mode 100644 index 0000000..60750f9 --- /dev/null +++ b/lua/trie.lua @@ -0,0 +1,306 @@ +-- Trie implementation in luajit +-- Copyright © 2019 Ashkan Kiani + +-- This program is free software: you can redistribute it and/or modify +-- it under the terms of the GNU General Public License as published by +-- the Free Software Foundation, either version 3 of the License, or +-- (at your option) any later version. + +-- This program is distributed in the hope that it will be useful, +-- but WITHOUT ANY WARRANTY; without even the implied warranty of +-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +-- GNU General Public License for more details. + +-- You should have received a copy of the GNU General Public License +-- along with this program. If not, see . +local ffi = require 'ffi' + +local bnot = bit.bnot +local band, bor, bxor = bit.band, bit.bor, bit.bxor +local lshift, rshift, rol = bit.lshift, bit.rshift, bit.rol + +ffi.cdef [[ +struct Trie { + bool is_leaf; + struct Trie* character[62]; +}; +void *malloc(size_t size); +void free(void *ptr); +]] + +local Trie_t = ffi.typeof('struct Trie') +local Trie_ptr_t = ffi.typeof('$ *', Trie_t) + +local function byte_to_index(b) + -- 0-9 starts at string.byte('0') == 0x30 == 48 == 0b0011_0000 + -- A-Z starts at string.byte('A') == 0x41 == 65 == 0b0100_0001 + -- a-z starts at string.byte('a') == 0x61 == 97 == 0b0110_0001 + + -- This works for mapping characters to + -- 0-9 A-Z a-z in that order + -- Letters have bit 0x40 set, so we use that as an indicator for + -- an additional offset from the space of the digits, and then + -- add the 10 allocated for the range of digits. + -- Then, within that indicator for letters, we subtract another + -- (65 - 97) which is the difference between lower and upper case + -- and add back another 26 to allocate for the range of uppercase + -- letters. + -- return b - 0x30 + -- + rshift(b, 6) * ( + -- 0x30 - 0x41 + -- + 10 + -- + band(1, rshift(b, 5)) * ( + -- 0x61 - 0x41 + -- + 26 + -- )) + return b - 0x30 - rshift(b, 6) * (7 + band(1, rshift(b, 5)) * 6) +end + +local function insensitive_byte_to_index(b) + -- return b - 0x30 + -- + rshift(b, 6) * ( + -- 0x30 - 0x61 + -- + 10 + -- ) + b = bor(b, 0x20) + return b - 0x30 - rshift(b, 6) * 39 +end + +local function verify_byte_to_index() + local chars = '0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz' + for i = 1, #chars do + local c = chars:sub(i,i) + local index = byte_to_index(string.byte(c)) + assert((i-1) == index, vim.inspect{index=index,c=c}) + end +end + +local Trie_size = ffi.sizeof(Trie_t) + +local function new_trie() + local ptr = ffi.C.malloc(Trie_size) + ffi.fill(ptr, Trie_size) + return ffi.cast(Trie_ptr_t, ptr) +end + +local INDEX_LOOKUP_TABLE = ffi.new('uint8_t[256]') +local b_a = string.byte('a') +local b_z = string.byte('z') +local b_A = string.byte('A') +local b_Z = string.byte('Z') +local b_0 = string.byte('0') +local b_9 = string.byte('9') +for i = 0, 255 do + if i >= b_0 and i <= b_9 then + INDEX_LOOKUP_TABLE[i] = i - b_0 + elseif i >= b_A and i <= b_Z then + INDEX_LOOKUP_TABLE[i] = i - b_A + 10 + elseif i >= b_a and i <= b_z then + INDEX_LOOKUP_TABLE[i] = i - b_a + 10 + 26 + else + INDEX_LOOKUP_TABLE[i] = 255 + end +end + +local function insert(trie, value) + if trie == nil then return false end + local node = trie + for i = 1, #value do + local index = INDEX_LOOKUP_TABLE[value:byte(i)] + if index == 255 then + return false + end + if node.character[index] == nil then + node.character[index] = new_trie() + end + node = node.character[index] + end + node.is_leaf = true + return node, trie +end + +local function search(trie, value) + if trie == nil then return false end + local node = trie + for i = 1, #value do + local index = INDEX_LOOKUP_TABLE[value:byte(i)] + if index == 255 then + return + end + local child = node.character[index] + if child == nil then + return false + end + node = child + end + return node.is_leaf +end + +local function longest_prefix(trie, value) + if trie == nil then return false end + local node = trie + local last_i = nil + for i = 1, #value do + local index = INDEX_LOOKUP_TABLE[value:byte(i)] + if index == 255 then + break + end + local child = node.character[index] + if child == nil then + break + end + if child.is_leaf then + last_i = i + end + node = child + end + if last_i then + return value:sub(1, last_i) + end +end + +--- Printing utilities + +local function index_to_char(index) + if index < 10 then + return string.char(index + b_0) + elseif index < 36 then + return string.char(index - 10 + b_A) + else + return string.char(index - 26 - 10 + b_a) + end +end + +local function trie_structure(trie) + if trie == nil then + return nil + end + local children = {} + for i = 0, 61 do + local child = trie.character[i] + if child ~= nil then + local child_table = trie_structure(child) + child_table.c = index_to_char(i) + table.insert(children, child_table) + end + end + return { + is_leaf = trie.is_leaf; + children = children; + } +end + +local function print_structure(s) + local mark + if not s then + return nil + end + if s.c then + if s.is_leaf then + mark = s.c.."*" + else + mark = s.c.."─" + end + else + mark = "├─" + end + if #s.children == 0 then + return {mark} + end + local lines = {} + for _, child in ipairs(s.children) do + local child_lines = print_structure(child) + for _, child_line in ipairs(child_lines) do + table.insert(lines, child_line) + end + end + for i, v in ipairs(lines) do + if v:match("^[%w%d]") then + if i == 1 then + lines[i] = mark.."─"..v + elseif i == #lines then + lines[i] = "└──"..v + else + lines[i] = "├──"..v + end + else + if i == 1 then + lines[i] = mark.."─"..v + elseif #s.children > 1 then + lines[i] = "│ "..v + else + lines[i] = " "..v + end + end + end + return lines +end + +local function free_trie(trie) + for i = 0, 61 do + local child = trie.character[i] + if child ~= nil then + free_trie(child) + end + end + ffi.C.free(trie) +end + +local Trie_mt = { + __index = { + insert = insert; + search = search; + longest_prefix = longest_prefix; + }; + __tostring = function(trie) + local structure = trie_structure(trie) + if structure then + return table.concat(print_structure(structure), '\n') + else + return 'nil' + end + end; + __gc = free_trie; + } +local Trie = ffi.metatype(Trie_t, Trie_mt) + +return Trie + +-- 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_structure(trie_structure(trie)), '\n')) -- cgit v1.2.3-70-g09d2