Update

lua/util/latex.lua

@@ -1,5 +1,42 @@
 local M = {}
 
+local has_treesitter, ts = pcall(require, "vim.treesitter")
+local _, query = pcall(require, "vim.treesitter.query")
+
+local M = {}
+
+local function get_node_at_cursor()
+  local cursor = vim.api.nvim_win_get_cursor(0)
+  local cursor_range = { cursor[1] - 1, cursor[2] }
+  local buf = vim.api.nvim_get_current_buf()
+  local ok, parser = pcall(ts.get_parser, buf, "latex")
+  if not ok or not parser then
+    return
+  end
+  local root_tree = parser:parse()[1]
+  local root = root_tree and root_tree:root()
+
+  if not root then
+    return
+  end
+
+  return root:named_descendant_for_range(
+    cursor_range[1],
+    cursor_range[2],
+    cursor_range[1],
+    cursor_range[2]
+  )
+end
+
+local MATH_ENVIRONMENTS = {
+  displaymath = true,
+  equation = true,
+  eqnarray = true,
+  align = true,
+  math = true,
+  array = true,
+}
+
 local MATH_NODES = {
   displayed_equation = true,
   inline_formula = true,
@@ -28,62 +65,33 @@ M.in_env = function(env)
   return pos[1] ~= 0 or pos[2] ~= 0
 end
 
--- For markdown
-M.in_mathzone_md = function()
-  local row, col = unpack(vim.api.nvim_win_get_cursor(0))
-  local node = vim.treesitter.get_node({ bufnr = 0, pos = { row - 1, col } })
-  while node do
-    if MATH_NODES[node:type()] then
-      return true
-    end
-    node = node:parent()
-  end
-  return false
-end
-M.in_text_md = function()
-  return not M.in_mathzone_md()
-end
-
 -- For typst
-M.in_mathzone_typ = function()
-  local row, col = unpack(vim.api.nvim_win_get_cursor(0))
-  local node = vim.treesitter.get_node({ bufnr = 0, pos = { row - 1, col - 1} })
-  while node do
-    if node:type() == "math" then
-      return true
+M.in_mathzone= function()
+  if has_treesitter then
+    local buf = vim.api.nvim_get_current_buf()
+    local node = get_node_at_cursor()
+    while node do
+      if MATH_NODES[node:type()] then
+        return true
+      elseif node:type() == "math_environment" or node:type() == "generic_environment" then
+        local begin = node:child(0)
+        local names = begin and begin:field("name")
+        if names and names[1] and MATH_ENVIRONMENTS[query.get_node_text(names[1], buf):match("[A-Za-z]+")] then
+          return true
+        end
+      end
+      node = node:parent()
     end
-    node = node:parent()
-  end
-  return false
-end
-
-M.in_mathzone_tex = function()
-  local row, col = unpack(vim.api.nvim_win_get_cursor(0))
-  local node = vim.treesitter.get_node({ bufnr = 0, pos = { row - 1, col - 1} })
-  while node do
-    if MATH_NODES[node:type()] then
-      return true
-    end
-    node = node:parent()
-  end
-
-  return false
-end
-
-M.in_mathzone = function()
-  local ft = vim.bo.filetype
-  if ft == "tex" then
-    return M.in_mathzone_tex()
-  elseif ft == "markdown" then
-    return M.in_mathzone_md()
-  elseif ft == "typst" then
-    return M.in_mathzone_typ()
+    return false
   end
 end
 
 M.in_text = function()
   return not M.in_mathzone()
 end
+M.in_text_md = function()
+  return not M.in_mathzone()
+end
 
 M.in_item = function()
   return M.in_env("itemize") or M.in_env("enumerate")
@@ -97,39 +105,5 @@ end
 M.in_tikz = function()
   return M.in_env("tikzpicture")
 end
-M.in_quantikz = function()
-  return M.in_env("quantikz")
-end
-M.in_algo = function()
-  return M.in_env("algorithmic")
-end
-
--- M.clean = function()
---   local current_dir = vim.fn.expand("%:p:h")
---   local file_types = { "aux", "log", "out", "fls", "fdb_latexmk", "bcf", "run.xml", "toc", "DS_Store", "bak*", "dvi" }
---   for _, file_type in ipairs(file_types) do
---     local command = "rm " .. current_dir .. "/*." .. file_type
---     vim.api.nvim_call_function("system", { command })
---   end
--- end
---
--- M.format = function()
---   local current_file = vim.fn.expand("%:p")
---   local latexindent = "latexindent -g /dev/null " .. current_file .. " -wd -l ~/Documents/Latex/latexindent.yaml"
---   local build = "pdflatex " .. current_file
---   vim.api.nvim_call_function("system", { build })
---   vim.cmd("w")
---   M.clean()
---   vim.api.nvim_call_function("system", { latexindent })
---   vim.cmd("e")
---   vim.cmd("normal! zz")
---   -- vim.cmd("TexlabForward")
--- end
---
--- M.sympy_calc = function()
---   local selected_text = vim.fn.getreg("v")
---   print(selected_text)
---   vim.api.nvim_out_write(selected_text)
--- end
 
 return M