refactor: extract ts parse logic into a single function

Author: TheLeoP

lua/refactoring/debug/cleanup.lua

@@ -30,6 +30,7 @@ end
 function M.cleanup(range_type, config)
   local get_extracted_range = require("refactoring.utils").get_extracted_range
   local apply_text_edits = require("refactoring.utils").apply_text_edits
+  local get_ts_info = require("refactoring.utils").get_ts_info
 
   local opts = config.debug.cleanup
 
@@ -58,16 +59,8 @@ function M.cleanup(range_type, config)
       return
     end
 
-    local comments = {} ---@type TSNode[]
-    for _, tree in ipairs(nested_lang_tree:trees()) do
-      for _, match, _ in query:iter_matches(tree:root(), buf) do
-        for capture_id, nodes in pairs(match) do
-          local name = query.captures[capture_id]
-
-          if name == "comment" then table.insert(comments, nodes[1]) end
-        end
-      end
-    end
+    local ts_info = get_ts_info(buf, nested_lang_tree, query)
+    local comments = ts_info.comments
 
     table.sort(comments, node_comp_asc)
     ---@type vim.Range[]

lua/refactoring/debug/print_loc.lua

@@ -32,6 +32,7 @@ function M.print_loc(range_type, config)
   local code_gen_error = require("refactoring.utils").code_gen_error
   local indent = require("refactoring.utils").indent
   local apply_text_edits = require("refactoring.utils").apply_text_edits
+  local get_ts_info = require("refactoring.utils").get_ts_info
 
   local opts = config.debug.print_loc
   local code_generation = opts.code_generation
@@ -64,33 +65,9 @@ function M.print_loc(range_type, config)
     local get_print_loc = code_generation.print_loc[lang]
     if not get_print_loc then return code_gen_error("print_loc", lang) end
 
-    -- TODO: change to a better name everywhere (debug_path_element?)
-    local debug_paths = {} ---@type refactor.DebugPath[]
-    local output_statements = {} ---@type refactor.OutputStatement[]
-    for _, tree in ipairs(nested_lang_tree:trees()) do
-      for _, match, metadata in query:iter_matches(tree:root(), buf) do
-        local output_statement ---@type nil|refactor.OutputStatement
-        for capture_id, nodes in pairs(match) do
-          local name = query.captures[capture_id]
-          if name == "debug_path" then
-            for i, node in ipairs(nodes) do
-              local text = type(metadata.text) == "string" and metadata.text
-                or ts.get_node_text(match[metadata.text][i], buf)
-              table.insert(debug_paths, { debug_path = node, text = text })
-            end
-          end
-
-          if name == "output_statement" then
-            output_statement = output_statement or {}
-            output_statement.output_statement = nodes[1]
-          elseif name == "output_statement.inside" then
-            output_statement = output_statement or {}
-            output_statement.inside = nodes[1]
-          end
-        end
-        if output_statement then table.insert(output_statements, output_statement) end
-      end
-    end
+    local ts_info = get_ts_info(buf, nested_lang_tree, query)
+    local debug_paths = ts_info.debug_paths
+    local output_statements = ts_info.output_statements
 
     local extracted_range_api = { extracted_range:to_extmark() }
     -- NOTE: treesitter nodes usualy do not include leading whitespace

lua/refactoring/debug/print_var.lua

@@ -44,6 +44,7 @@ function M.print_var(range_type, config)
   local scopes_for_range = require("refactoring.utils").scopes_for_range
   local get_declaration_scope = require("refactoring.utils").get_declaration_scope
   local indent = require("refactoring.utils").indent
+  local get_ts_info = require("refactoring.utils").get_ts_info
 
   local opts = config.debug.print_var
   local code_generation = opts.code_generation
@@ -79,41 +80,20 @@ function M.print_var(range_type, config)
     local get_print_var = code_generation.print_var[lang]
     if not get_print_var then return code_gen_error("print_var", lang) end
 
-    local references = {} ---@type refactor.ReferenceInfo[]
-    local output_statements = {} ---@type refactor.OutputStatement[]
-    local scopes = {} ---@type TSNode[]
-    for _, tree in ipairs(nested_lang_tree:trees()) do
-      for _, match, metadata in query:iter_matches(tree:root(), buf) do
-        local output_statement ---@type nil|refactor.OutputStatement
-        for capture_id, nodes in pairs(match) do
-          local name = query.captures[capture_id]
-          if name == "reference.identifier" then
-            for i, node in ipairs(nodes) do
-              table.insert(references, {
-                identifier = node,
-                reference_type = metadata.reference_type,
-                type = metadata.types and metadata.types[i],
-                declaration = metadata.declaration ~= nil,
-              })
-            end
-          end
-
-          if name == "output_statement" then
-            output_statement = output_statement or {}
-            output_statement.output_statement = nodes[1]
-          elseif name == "output_statement.inside" then
-            output_statement = output_statement or {}
-            output_statement.inside = nodes[1]
-          end
-          if name == "scope" then
-            for _, node in ipairs(nodes) do
-              table.insert(scopes, node)
-            end
-          end
+    local ts_info = get_ts_info(buf, nested_lang_tree, query)
+    local references = ts_info.references
+    local output_statements = ts_info.output_statements
+    -- TODO: modify the util functions that use `scopes` as TSNode[] to use
+    -- refactor.Scope[] instead?
+    ---@type TSNode[]
+    local scopes = iter(ts_info.scopes)
+      :map(
+        ---@param scope refactor.Scope
+        function(scope)
+          return scope.scope
         end
-        if output_statement then table.insert(output_statements, output_statement) end
-      end
-    end
+      )
+      :totable()
 
     local extracted_range_api = { extracted_range:to_extmark() }
     -- NOTE: treesitter nodes usualy do not include leading whitespace
@@ -187,7 +167,6 @@ function M.print_var(range_type, config)
       end
     end
 
-    local output_start_pos = pos(output_range.start_row, output_range.start_col, { buf = output_range.buf })
     -- TODO: I also compute `declarations_before_output_range` in
     -- `extract_func`. Is there a cleaner wat to do all this in both places?
     local declarations_by_scope = get_declarations_by_scope(references, scopes, buf)

lua/refactoring/refactor/extract_func.lua

@@ -59,33 +59,10 @@ end
 ---@return {method: boolean?, singleton: boolean?, struct_name: string?, struct_var_name: string?}
 local function get_output_node(nested_lang_tree, query, buf, extracted_range)
   local is_first_closer = require("refactoring.utils").is_first_closer
+  local get_ts_info = require("refactoring.utils").get_ts_info
 
-  local outputs = {} ---@type refactor.Output[]
-  for _, tree in ipairs(nested_lang_tree:trees()) do
-    for _, match, metadata in query:iter_matches(tree:root(), buf) do
-      local output ---@type table|refactor.Output|nil
-      for capture_id, nodes in pairs(match) do
-        local name = query.captures[capture_id]
-
-        -- TODO: split input.info and output location
-        if name == "output.comment" then
-          output = output or {}
-          output.comment = nodes
-        elseif name == "output.function" then
-          output = output or {}
-          output.fn = nodes[1]
-          output.method = metadata.method ~= nil
-          output.singleton = metadata.singleton ~= nil
-
-          local struct_name = metadata.struct_name
-          if struct_name then output.struct_name = ts.get_node_text(match[struct_name][1], buf) end
-          local struct_var_name = metadata.struct_var_name
-          if struct_var_name then output.struct_var_name = ts.get_node_text(match[struct_var_name][1], buf) end
-        end
-      end
-      if output then table.insert(outputs, output) end
-    end
-  end
+  local ts_info = get_ts_info(buf, nested_lang_tree, query)
+  local outputs = ts_info.outputs
 
   local extracted_start_pos = pos(extracted_range.start_row, extracted_range.start_col, { buf = extracted_range.buf })
   ---@type refactor.Output|nil
@@ -180,6 +157,7 @@ local function extract_func(opts)
   local get_declarations_by_scope = require("refactoring.utils").get_declarations_by_scope
   local scopes_for_range = require("refactoring.utils").scopes_for_range
   local get_declaration_scope = require("refactoring.utils").get_declaration_scope
+  local get_ts_info = require("refactoring.utils").get_ts_info
 
   local code_generation = opts.config_opts.code_generation
 
@@ -220,30 +198,19 @@ local function extract_func(opts)
     output_range = range.extmark(0, 0, 0, 0, { buf = out_buf })
   end
 
-  local references_info = {} ---@type refactor.ReferenceInfo[]
-  local scopes = {} ---@type TSNode[]
-  for _, tree in ipairs(nested_lang_tree:trees()) do
-    for _, match, metadata in in_query:iter_matches(tree:root(), in_buf) do
-      for capture_id, nodes in pairs(match) do
-        local name = in_query.captures[capture_id]
-        if name == "reference.identifier" then
-          for i, node in ipairs(nodes) do
-            table.insert(references_info, {
-              identifier = node,
-              reference_type = metadata.reference_type,
-              type = metadata.types and metadata.types[i],
-              declaration = metadata.declaration ~= nil,
-            })
-          end
-        elseif name == "scope" then
-          for _, node in ipairs(nodes) do
-            table.insert(scopes, node)
-          end
-        end
+  local in_ts_info = get_ts_info(in_buf, nested_lang_tree, in_query)
+  local references_info = in_ts_info.references
+  -- TODO: modify the util functions that use `scopes` as TSNode[] to use
+  -- refactor.Scope[] instead?
+  ---@type TSNode[]
+  local scopes = iter(in_ts_info.scopes)
+    :map(
+      ---@param scope refactor.Scope
+      function(scope)
+        return scope.scope
       end
-    end
-  end
-  -- TODO: maybe check that all the treesitter captures are not empty(?
+    )
+    :totable()
 
   local scopes_for_extracted_range = scopes_for_range(in_buf, scopes, extracted_range)
 
@@ -272,7 +239,6 @@ local function extract_func(opts)
     ---@param a refactor.ReferenceInfo
     ---@param b refactor.ReferenceInfo
     function(a, b)
-      -- TODO: don't I already have a function to sort nodes in utils?
       local a_srow, a_scol, a_erow, a_ecol = a.identifier:range()
       local a_range = range(a_srow, a_scol, a_erow, a_ecol, { buf = in_buf })
       local b_srow, b_scol, b_erow, b_ecol = b.identifier:range()

lua/refactoring/refactor/extract_var.lua

@@ -65,6 +65,7 @@ function M.extract_var(range_type, config)
   local apply_text_edits = require("refactoring.utils").apply_text_edits
   local input = require("refactoring.utils").input
   local code_gen_error = require("refactoring.utils").code_gen_error
+  local get_ts_info = require("refactoring.utils").get_ts_info
 
   local opts = config.refactor.extract_var
   local code_generation = opts.code_generation
@@ -118,30 +119,14 @@ function M.extract_var(range_type, config)
 
     local extracted_significant_text = significant_text(encompassing_node, buf)
     local matching_nodes = {} ---@type TSNode[]
-    local scopes = {} ---@type refactor.Scope[]
     for _, tree in ipairs(nested_lang_tree:trees()) do
       for _, node in encompasing_query:iter_captures(tree:root(), buf) do
         local node_significant_text = significant_text(node, buf)
         if node_significant_text == extracted_significant_text then table.insert(matching_nodes, node) end
       end
-      for _, match in query:iter_matches(tree:root(), buf) do
-        local match_info ---@type refactor.Scope|nil
-        for capture_id, nodes in pairs(match) do
-          local name = query.captures[capture_id]
-          if name == "scope" then
-            match_info = match_info or {}
-            match_info.scope = nodes[1]
-          elseif name == "scope.inside" then
-            match_info = match_info or {}
-            match_info.inside = nodes[1]
-          elseif name == "scope.outside" then
-            match_info = match_info or {}
-            match_info.outside = nodes[1]
-          end
-        end
-        if match_info then table.insert(scopes, match_info) end
-      end
     end
+    local ts_info = get_ts_info(buf, nested_lang_tree, query)
+    local scopes = ts_info.scopes
 
     ---@type {[integer]: refactor.TextEdit[]}
     local text_edits_by_buf = {}