Inline math highlighting (Continued from PR #426) (#428)

Co-authored-by: Kristijan Husak <[email protected]>

Author: gzagatti

lua/orgmode/colors/markup_highlighter.lua

@@ -2,33 +2,64 @@ local config = require('orgmode.config')
 local ts_utils = require('nvim-treesitter.ts_utils')
 local query = nil
 
-local valid_pre_marker_chars = { ' ', '(', '-', "'", '"', '{' }
-local valid_post_marker_chars = { ' ', ')', '-', '}', '"', "'", ':', ';', '!', '\\', '[', ',', '.', '?' }
+local valid_pre_marker_chars = { ' ', '(', '-', "'", '"', '{', '*', '/', '_', '+' }
+local valid_post_marker_chars =
+  { ' ', ')', '-', '}', '"', "'", ':', ';', '!', '\\', '[', ',', '.', '?', '*', '/', '_', '+' }
 
 local markers = {
   ['*'] = {
     hl_name = 'org_bold',
     hl_cmd = 'hi def org_bold term=bold cterm=bold gui=bold',
+    nestable = true,
+    type = 'text',
   },
   ['/'] = {
     hl_name = 'org_italic',
     hl_cmd = 'hi def org_italic term=italic cterm=italic gui=italic',
+    nestable = true,
+    type = 'text',
   },
   ['_'] = {
     hl_name = 'org_underline',
     hl_cmd = 'hi def org_underline term=underline cterm=underline gui=underline',
+    nestable = true,
+    type = 'text',
   },
   ['+'] = {
     hl_name = 'org_strikethrough',
     hl_cmd = 'hi def org_strikethrough term=strikethrough cterm=strikethrough gui=strikethrough',
+    nestable = true,
+    type = 'text',
   },
   ['~'] = {
     hl_name = 'org_code',
     hl_cmd = 'hi def link org_code String',
+    nestable = false,
+    type = 'text',
   },
   ['='] = {
     hl_name = 'org_verbatim',
     hl_cmd = 'hi def link org_verbatim String',
+    nestable = false,
+    type = 'text',
+  },
+  ['\\('] = {
+    hl_name = 'org_latex',
+    hl_cmd = 'hi def link org_latex OrgTSLatex',
+    nestable = false,
+    type = 'latex',
+  },
+  ['\\{'] = {
+    hl_name = 'org_latex',
+    hl_cmd = 'hi def link org_latex OrgTSLatex',
+    nestable = false,
+    type = 'latex',
+  },
+  ['\\s'] = {
+    hl_name = 'org_latex',
+    hl_cmd = 'hi def link org_latex OrgTSLatex',
+    nestable = false,
+    type = 'latex',
   },
 }
 
@@ -71,23 +102,18 @@ local get_tree = ts_utils.memoize_by_buf_tick(function(bufnr)
   return tree[1]:root()
 end)
 
-local function get_predicate_nodes(match)
+local function get_predicate_nodes(match, n)
+  local total = n or 2
   local counter = 1
-  local start_node = nil
-  local end_node = nil
+  local nodes = {}
   for i, node in pairs(match) do
-    if counter == 1 then
-      start_node = node
-    end
-    if counter == 2 then
-      end_node = node
-    end
+    nodes[counter] = node
     counter = counter + 1
+    if counter > total then
+      break
+    end
   end
-  if not start_node or not end_node then
-    return false
-  end
-  return start_node, end_node
+  return unpack(nodes)
 end
 
 local function is_valid_markup_range(match, _, source, _)
@@ -96,9 +122,11 @@ local function is_valid_markup_range(match, _, source, _)
     return
   end
 
-  -- Ignore conflicts with hyperlink
-  if start_node:type() == '[' or end_node:type() == ']' then
-    return true
+  -- Ignore conflicts with hyperlink or math
+  for _, char in ipairs({ '[', '\\' }) do
+    if start_node:type() == char or end_node:type() == char then
+      return true
+    end
   end
 
   local start_line = start_node:range()
@@ -146,6 +174,48 @@ local function is_valid_hyperlink_range(match, _, source, _)
   return is_valid_start and is_valid_end
 end
 
+local function is_valid_latex_range(match, _, source, _)
+  local start_node_left, start_node_right, end_node = get_predicate_nodes(match, 3)
+  -- Ignore conflicts with markup
+  if start_node_left:type() ~= '\\' then
+    return true
+  end
+  if not start_node_right or not end_node then
+    return
+  end
+
+  local start_line = start_node_left:range()
+  local end_line = start_node_left:range()
+
+  if start_line ~= end_line then
+    return false
+  end
+
+  local _, start_left_col_end = start_node_left:end_()
+  local _, start_right_col_end = start_node_right:end_()
+  local start_text = get_node_text(start_node_left, source, 0, start_right_col_end - start_left_col_end)
+
+  if start_text == '\\(' then
+    local end_text = get_node_text(end_node, source, -1, 0)
+    if end_text == '\\)' then
+      return true
+    end
+  else
+    -- we have to deal with two cases here either \foo{bar} or \bar
+    local char_after_start = get_node_text(start_node_right, source, 0, 1):sub(-1)
+    local end_text = get_node_text(end_node, source, 0, 0)
+    -- if \foo{bar}
+    if char_after_start == '{' and end_text == '}' then
+      return true
+    end
+    -- elseif \bar
+    if not start_text:sub(2):match('%A') and end_text ~= '}' then
+      return true
+    end
+  end
+  return false
+end
+
 local function load_deps()
   -- Already defined
   if query then
@@ -154,6 +224,7 @@ local function load_deps()
   query = vim.treesitter.get_query('org', 'markup')
   vim.treesitter.query.add_predicate('org-is-valid-markup-range?', is_valid_markup_range)
   vim.treesitter.query.add_predicate('org-is-valid-hyperlink-range?', is_valid_hyperlink_range)
+  vim.treesitter.query.add_predicate('org-is-valid-latex-range?', is_valid_latex_range)
 end
 
 ---@param bufnr? number
@@ -173,15 +244,19 @@ local function get_matches(bufnr, first_line, last_line)
   for _, match, _ in query:iter_matches(root, bufnr, first_line, last_line) do
     for _, node in pairs(match) do
       local char = node:type()
-      local range = ts_utils.node_to_lsp_range(node)
-      local linenr = tostring(range.start.line)
-      taken_locations[linenr] = taken_locations[linenr] or {}
-      if not taken_locations[linenr][range.start.character] then
-        table.insert(ranges, {
-          type = char,
-          range = range,
-        })
-        taken_locations[linenr][range.start.character] = true
+      -- saves unnecessary parsing, since \\ is not used below
+      if char ~= '\\' then
+        local range = ts_utils.node_to_lsp_range(node)
+        local linenr = tostring(range.start.line)
+        taken_locations[linenr] = taken_locations[linenr] or {}
+        if not taken_locations[linenr][range.start.character] then
+          table.insert(ranges, {
+            type = char,
+            range = range,
+            node = node,
+          })
+          taken_locations[linenr][range.start.character] = true
+        end
       end
     end
   end
@@ -197,30 +272,74 @@ local function get_matches(bufnr, first_line, last_line)
   local seek_link = {}
   local result = {}
   local link_result = {}
+  local latex_result = {}
+
+  local nested = {}
+  local can_nest = true
+
+  local type_map = {
+    ['('] = '\\(',
+    [')'] = '\\(',
+    ['}'] = '\\{',
+  }
 
   for _, item in ipairs(ranges) do
+    if item.type == '(' then
+      item.range.start.character = item.range.start.character - 1
+    elseif item.type == 'str' then
+      item.range.start.character = item.range.start.character - 1
+      local char = get_node_text(item.node, bufnr, 0, 1):sub(-1)
+      if char == '{' then
+        item.type = '\\{'
+      else
+        item.type = '\\s'
+      end
+    end
+
+    item.type = type_map[item.type] or item.type
+
     if markers[item.type] then
       if seek[item.type] then
         local from = seek[item.type]
-        table.insert(result, {
-          type = item.type,
-          from = from.range,
-          to = item.range,
-        })
-
-        seek[item.type] = nil
-
-        for t, pos in pairs(seek) do
-          if
-            pos.range.start.line == from.range.start.line
-            and pos.range.start.character > from.range['end'].character
-            and pos.range.start.character < item.range.start.character
-          then
-            seek[t] = nil
+        if nested[#nested] == nil or nested[#nested] == from.type then
+          local target_result = result
+          if markers[item.type].type == 'latex' then
+            target_result = latex_result
+          end
+
+          table.insert(target_result, {
+            type = item.type,
+            from = from.range,
+            to = item.range,
+          })
+
+          seek[item.type] = nil
+          nested[#nested] = nil
+          can_nest = true
+
+          for t, pos in pairs(seek) do
+            if
+              pos.range.start.line == from.range.start.line
+              and pos.range.start.character > from.range['end'].character
+              and pos.range.start.character < item.range.start.character
+            then
+              seek[t] = nil
+            end
           end
         end
-      else
-        seek[item.type] = item
+      elseif can_nest then
+        -- escaped strings have no pairs, their markup info is self-contained
+        if item.type == '\\s' then
+          table.insert(result, {
+            type = item.type,
+            from = item.range,
+            to = item.range,
+          })
+        else
+          seek[item.type] = item
+          nested[#nested + 1] = item.type
+          can_nest = markers[item.type].nestable
+        end
       end
     end
 
@@ -237,7 +356,7 @@ local function get_matches(bufnr, first_line, last_line)
     end
   end
 
-  return result, link_result
+  return result, link_result, latex_result
 end
 
 local function apply(namespace, bufnr, _, first_line, last_line, _)
@@ -252,7 +371,7 @@ local function apply(namespace, bufnr, _, first_line, last_line, _)
   if #visible_lines == 0 then
     return
   end
-  local ranges, link_ranges = get_matches(bufnr, visible_lines[1], visible_lines[#visible_lines])
+  local ranges, link_ranges, latex_ranges = get_matches(bufnr, visible_lines[1], visible_lines[#visible_lines])
   local hide_markers = config.org_hide_emphasis_markers
 
   for _, range in ipairs(ranges) do
@@ -301,6 +420,15 @@ local function apply(namespace, bufnr, _, first_line, last_line, _)
       conceal = '',
     })
   end
+
+  for _, latex_range in ipairs(latex_ranges) do
+    vim.api.nvim_buf_set_extmark(bufnr, namespace, latex_range.from.start.line, latex_range.from.start.character, {
+      ephemeral = true,
+      end_col = latex_range.to['end'].character,
+      hl_group = markers[latex_range.type].hl_name,
+      priority = 110 + latex_range.from.start.character,
+    })
+  end
 end
 
 local function setup()

queries/org/markup.scm

@@ -13,6 +13,11 @@
  (expr "+" @strikethrough.start "+" @strikethrough.end (#org-is-valid-markup-range? @strikethrough.start @strikethrough.end))
  ((expr "[" @hyperlink.start) (expr "]" @hyperlink.end) (#org-is-valid-hyperlink-range? @hyperlink.start @hyperlink.end))
  (expr "[" @hyperlink.start "]" @hyperlink.end (#org-is-valid-hyperlink-range? @hyperlink.start @hyperlink.end))
+ ((expr ("\\" @text.math.start.left "(" @text.math.start.right)) (expr ("\\" ")" @text.math.end)) (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left "(" @text.math.start.right "\\" ")" @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ ((expr ("\\" @text.math.start.left ("str")+ @text.math.start.right "{")) (expr "}" @text.math.end) (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left ("str")+ @text.math.start.right "{" "}" @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left ("str")+ @text.math.start.right @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
 ])
 
 (item [
@@ -30,6 +35,11 @@
  (expr "+" @strikethrough.start "+" @strikethrough.end (#org-is-valid-markup-range? @strikethrough.start @strikethrough.end))
  ((expr "[" @hyperlink.start) (expr "]" @hyperlink.end) (#org-is-valid-hyperlink-range? @hyperlink.start @hyperlink.end))
  (expr "[" @hyperlink.start "]" @hyperlink.end (#org-is-valid-hyperlink-range? @hyperlink.start @hyperlink.end))
+ ((expr ("\\" @text.math.start.left "(" @text.math.start.right)) (expr ("\\" ")" @text.math.end)) (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left "(" @text.math.start.right "\\" ")" @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ ((expr ("\\" @text.math.start.left ("str")+ @text.math.start.right "{")) (expr "}" @text.math.end) (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left ("str")+ @text.math.start.right "{" "}" @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left ("str")+ @text.math.start.right @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
 ])
 
 (cell (contents [
@@ -47,6 +57,11 @@
  (expr "+" @strikethrough.start "+" @strikethrough.end (#org-is-valid-markup-range? @strikethrough.start @strikethrough.end))
  ((expr "[" @hyperlink.start) (expr "]" @hyperlink.end) (#org-is-valid-hyperlink-range? @hyperlink.start @hyperlink.end))
  (expr "[" @hyperlink.start "]" @hyperlink.end (#org-is-valid-hyperlink-range? @hyperlink.start @hyperlink.end))
+ ((expr ("\\" @text.math.start.left "(" @text.math.start.right)) (expr ("\\" ")" @text.math.end)) (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left "(" @text.math.start.right "\\" ")" @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ ((expr ("\\" @text.math.start.left ("str")+ @text.math.start.right "{")) (expr "}" @text.math.end) (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left ("str")+ @text.math.start.right "{" "}" @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left ("str")+ @text.math.start.right @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
 ]))
 
 (drawer (contents [
@@ -64,4 +79,9 @@
  (expr "+" @strikethrough.start "+" @strikethrough.end (#org-is-valid-markup-range? @strikethrough.start @strikethrough.end))
  ((expr "[" @hyperlink.start) (expr "]" @hyperlink.end) (#org-is-valid-hyperlink-range? @hyperlink.start @hyperlink.end))
  (expr "[" @hyperlink.start "]" @hyperlink.end (#org-is-valid-hyperlink-range? @hyperlink.start @hyperlink.end))
+ ((expr ("\\" @text.math.start.left "(" @text.math.start.right)) (expr ("\\" ")" @text.math.end)) (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left "(" @text.math.start.right "\\" ")" @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ ((expr ("\\" @text.math.start.left ("str")+ @text.math.start.right "{")) (expr "}" @text.math.end) (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left ("str")+ @text.math.start.right "{" "}" @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
+ (expr "\\" @text.math.start.left ("str")+ @text.math.start.right @text.math.end (#org-is-valid-latex-range? @text.math.start.left @text.math.start.right @text.math.end))
 ]))