From 68a82aaa5d75d532b597d8389e981334baa5cab3 Mon Sep 17 00:00:00 2001
From: v3natio <72844565+v3natio@users.noreply.github.com>
Date: Thu, 3 Apr 2025 10:01:56 +0200
Subject: [PATCH] Added treesitter check for math/text mode

---
 lua/obsidian/completion/refs.lua         |  5 ++
 lua/obsidian/treesitter.lua              | 66 ++++++++++++++++++++++++
 test/fixtures/notes/note_with_MathJax.md | 19 +++++++
 3 files changed, 90 insertions(+)
 create mode 100644 lua/obsidian/treesitter.lua
 create mode 100644 test/fixtures/notes/note_with_MathJax.md

diff --git a/lua/obsidian/completion/refs.lua b/lua/obsidian/completion/refs.lua
index 70e009c1..5cf1aba9 100644
--- a/lua/obsidian/completion/refs.lua
+++ b/lua/obsidian/completion/refs.lua
@@ -1,4 +1,5 @@
 local util = require "obsidian.util"
+local ts_util = require "obsidian.treesitter"
 
 local M = {}
 
@@ -13,6 +14,10 @@ M.RefType = {
 ---@param input string
 ---@return string|?, string|?, obsidian.completion.RefType|?
 local find_search_start = function(input)
+  if ts_util.in_mathzone() then
+    return nil
+  end
+
   for i = string.len(input), 1, -1 do
     local substr = string.sub(input, i)
     if vim.startswith(substr, "]") or vim.endswith(substr, "]") then
diff --git a/lua/obsidian/treesitter.lua b/lua/obsidian/treesitter.lua
new file mode 100644
index 00000000..c0c95b8c
--- /dev/null
+++ b/lua/obsidian/treesitter.lua
@@ -0,0 +1,66 @@
+local M = {}
+
+local ts = require "vim.treesitter"
+
+local MATH_NODES = {
+  displayed_equation = true,
+  inline_formula = true,
+  math_environment = true,
+}
+
+local TEXT_NODES = {
+  text_mode = true,
+  label_definition = true,
+  label_reference = true,
+}
+
+local CODE_BLOCK_NODES = {
+  fenced_code_block = true,
+  indented_code_block = true,
+}
+
+function M.in_text(check_parent)
+  local node = ts.get_node { ignore_injections = false }
+  while node do
+    local node_type = node:type()
+    -- if in code block, always consider it as text
+    if CODE_BLOCK_NODES[node_type] then
+      return true
+    end
+    if node_type == "text_mode" then
+      if check_parent then
+        local parent = node:parent()
+        if parent and MATH_NODES[parent:type()] then
+          return false
+        end
+      end
+      return true
+    end
+    if MATH_NODES[node_type] then
+      return false
+    end
+    node = node:parent()
+  end
+  return true
+end
+
+function M.in_mathzone()
+  local node = ts.get_node { ignore_injections = false }
+  while node do
+    local node_type = node:type()
+    -- if in code block, don't consider it math.
+    if CODE_BLOCK_NODES[node_type] then
+      return false
+    end
+    if TEXT_NODES[node_type] then
+      return false
+    end
+    if MATH_NODES[node_type] then
+      return true
+    end
+    node = node:parent()
+  end
+  return false
+end
+
+return M
diff --git a/test/fixtures/notes/note_with_MathJax.md b/test/fixtures/notes/note_with_MathJax.md
new file mode 100644
index 00000000..a59ffb27
--- /dev/null
+++ b/test/fixtures/notes/note_with_MathJax.md
@@ -0,0 +1,19 @@
+---
+id: note_with_MathJax
+---
+
+# Inline and Display Modes
+
+By the law of large numbers we have that $\widehat{\mathbb{V}[\theta_{i}]} \to \mathbb{V}[\theta_{i}]$, hence
+$$
+\begin{flalign*}
+  \hat{\beta}_{1} &\approx \beta_{1} + \frac{\frac{1}{n} \sum_{i=1}^{n} \left( X_{i} - \mu_{X} \right) U_{i}}{\frac{1}{n} \sum_{i=1}^{n} \left( X_{i} - \mu_{X} \right)^{2}} &&\\
+  \mathbb{V} \left[ \hat{\beta}_{1} \right] &= \mathbb{V} \left[ \beta_{1} + \frac{\frac{1}{n} \sum_{i=1}^{n} \left( X_{i} - \mu_{X} \right) U_{i}}{\sigma_{X}^{2}} \right] &&\\
+  &= \mathbb{V} \left[ \beta_{1} \right] + \mathbb{V} \left[ \frac{\frac{1}{n} \sum_{i=1}^{n} \left( X_{i} - \mu_{X} \right) U_{i}}{\sigma_{X}^{2}} \right] &&\\
+  &= \left( \frac{1}{\sigma_{X}^{2}} \right)^{2} \cdot \left( \frac{1}{n} \right)^{2} \cdot \underbrace{\mathbb{V} \left[ \sum_{i=1}^{n} \left( X_{i} - \mu_{X} \right) U_{i} \right]}_{\textcolor{yellow}{*_{1}}} &&\\
+  &= \left( \frac{1}{\sigma_{X}^{2}} \right)^{2} \cdot \left( \frac{1}{n} \right)^{2} \cdot n \mathbb{V} \left[ \left( X_{i} - \mu_{X} \right) U_{i} \right] &&\\
+  &= \frac{1}{n} \cdot \frac{\mathbb{V} \left[ \left( X_{i} - \mu_{X} \right) U_{i} \right]}{\left( \sigma_{X}^{2} \right)^{2}} &&\\
+\end{flalign*}
+$$
+
+Therefore, $\hat{\beta}_{1} \overset{\text{CLT}}{\sim} \mathcal{N} \left( \beta_{1}, \sigma_{\hat{\beta}_{1}}^{2} \right)$, where $\sigma_{\hat{\beta}_{1}}^{2} = \frac{\mathbb{V} \left[ \left( X_{i} - \mu_{X} \right) U_{i} \right]}{n \cdot \left( \sigma_{X}^{2} \right)^{2}}$.