From c17dd6f8645770ff34ca4d592297a51167d81dd3 Mon Sep 17 00:00:00 2001
From: Philippe Massicotte <pmassicotte@hotmail.com>
Date: Sun, 23 Feb 2025 15:13:24 -0500
Subject: [PATCH] refactor(numbers.lua, quarto.lua): simplify code chunk
 extraction by removing root node dependency and improve error handling

---
 lua/r/format/numbers.lua | 11 ++++++-----
 lua/r/quarto.lua         | 22 +++++++++++++---------
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/lua/r/format/numbers.lua b/lua/r/format/numbers.lua
index 5147a90a..07c135dd 100644
--- a/lua/r/format/numbers.lua
+++ b/lua/r/format/numbers.lua
@@ -1,6 +1,5 @@
 local inform = require("r.log").inform
 local get_code_chunk = require("r.quarto").get_code_chunk
-local get_root_node = require("r.utils").get_root_node
 
 local M = {}
 
@@ -64,11 +63,13 @@ M.formatnum = function(bufnr)
         return
     end
 
-    local root_node = get_root_node(bufnr)
-    if not root_node then return end
-
     if filetype == "quarto" or filetype == "rmd" then
-        local r_chunks_content = get_code_chunk(root_node, bufnr, nil, "r")
+        local r_chunks_content = get_code_chunk(bufnr, "r")
+
+        if not r_chunks_content then
+            error("Failed to extract code chunks.")
+            return
+        end
 
         for _, r_chunk in ipairs(r_chunks_content) do
             find_and_replace_float(r_chunk.content, bufnr, r_chunk.start_row, 0)
diff --git a/lua/r/quarto.lua b/lua/r/quarto.lua
index c0423d76..bd5be3ea 100644
--- a/lua/r/quarto.lua
+++ b/lua/r/quarto.lua
@@ -21,12 +21,16 @@ M.command = function(what)
     send_cmd(cmd)
 end
 
--- Helper function to get R content/code block from Quarto document
--- @param root The root node of the parsed tree
--- @param bufnr The buffer number
--- @param cursor_pos The cursor position (optional)
--- @return A table containing R code blocks with their content, start row and end row
-M.get_code_chunk = function(root, bufnr, cursor_pos, code_fence_language)
+--- Find and replace floating point numbers in the given R content.
+---@param bufnr  integer The buffer number.
+---@param lang string The language of the code chunk.
+---@param row integer|nil The row number. If nil, all code chunks are returned.
+---@return table|nil
+M.get_code_chunk = function(bufnr, lang, row)
+    local root = require("r.utils").get_root_node()
+
+    if not root then return nil end
+
     local query = vim.treesitter.query.parse(
         "markdown",
         string.format(
@@ -35,7 +39,7 @@ M.get_code_chunk = function(root, bufnr, cursor_pos, code_fence_language)
                 (info_string (language) @lang (#eq? @lang "%s"))
                 (code_fence_content) @content)
             ]],
-            code_fence_language
+            lang
         )
     )
 
@@ -45,13 +49,13 @@ M.get_code_chunk = function(root, bufnr, cursor_pos, code_fence_language)
     for _, node, _ in query:iter_captures(root, bufnr, 0, -1) do
         if node:type() == "code_fence_content" then
             local start_row, _, end_row, _ = node:range()
-            if not cursor_pos or (cursor_pos >= start_row and cursor_pos <= end_row) then
+            if not row or (row >= start_row and row <= end_row) then
                 table.insert(r_contents, {
                     content = vim.treesitter.get_node_text(node, bufnr),
                     start_row = start_row,
                     end_row = end_row,
                 })
-                if cursor_pos then break end
+                if row then break end
             end
         end
     end