feat: 106: Extract Function

Just the bare bones of the project getting started here.

Author: ThePrimeagen

lua/refactoring/106.lua

@@ -0,0 +1,138 @@
+local dev = require("refactoring.dev")
+dev.reload()
+
+local ts_utils = require("nvim-treesitter.ts_utils")
+local utils = require("refactoring.utils")
+local vim_helpers = require("refactoring.vim-helpers")
+
+local REFACTORING = {}
+local REFACTORING_OPTIONS = {
+    code_generation = {
+        lua = {
+            extract_function = function(opts)
+                return {
+                    create = table.concat(vim.tbl_flatten({
+                        string.format("local function %s(%s)", opts.name, table.concat(opts.args, ", ")),
+                        opts.body,
+                        "end",
+                        ""
+                    }), "\n"),
+
+                    call = string.format("%s(%s)", opts.name, table.concat(opts.args, ", ")),
+                }
+            end,
+        }
+    }
+}
+
+local function get_selection_range()
+    local _, start_row, start_col, _ = unpack(vim.fn.getpos("'<"))
+    local _, end_row, _, _ = unpack(vim.fn.getpos("'>"))
+    local end_col = vim.fn.col("'>")
+
+    -- end_col :: TS is 0 based, and '> on line selections is char_count + 1
+    -- I think - 2 is correct on
+    --
+    -- end_row : end_row is exclusive in TS, so we don't minus
+    return start_row, start_col, end_row, end_col
+end
+
+local function vim_range_to_ts_range(start_row, start_col, end_row, end_col)
+    return start_row - 1, start_col, end_row - 1, end_col
+end
+
+-- 106
+local function get_text_edits(selected_local_references, end_row, lang, start_col, start_row, end_col, scope_range, function_name)
+    -- local declaration within the selection range.
+    local lsp_text_edits = {}
+    local extract_function = REFACTORING_OPTIONS.code_generation[lang].extract_function({
+        args = vim.tbl_keys(selected_local_references),
+        body = vim.api.nvim_buf_get_lines(0, start_row, end_row, false),
+        name = function_name,
+    })
+    table.insert(lsp_text_edits, {
+        range = scope_range,
+        newText = string.format("\n%s", extract_function.create)
+    })
+    table.insert(lsp_text_edits, {
+        range = {
+            start = {line = start_row, character = start_col},
+            ["end"] = {line = end_row, character = end_col}
+        },
+        newText = string.format("\n%s", extract_function.call)
+    })
+    return lsp_text_edits
+end
+
+local function get_scope_range(scope)
+    -- vim_helpers.move_text(0, start_row, end_row, scope:range())
+    local scope_range = ts_utils.node_to_lsp_range(scope)
+    scope_range.start.line = scope_range.start.line - 1
+    scope_range["end"] = scope_range.start
+
+    return scope_range
+end
+
+local function get_local_definitions(local_defs, function_args)
+    local local_def_map = {}
+
+    for _, def in pairs(local_defs) do
+        local_def_map[ts_utils.get_node_text(def)[1]] = true
+    end
+    for _, def in pairs(function_args) do
+        local_def_map[ts_utils.get_node_text(def)[1]] = true
+    end
+
+    return local_def_map
+end
+
+REFACTORING.extract = function(bufnr)
+    -- lua 1  based index
+    -- vim apis are 1 based
+    -- treesitter is 0 based
+    -- first entry (1), line 1, row 0
+    bufnr = bufnr or 0
+
+    local lang = vim.bo.filetype
+    local start_row, start_col, end_row, end_col = get_selection_range()
+    local ts_start_row, ts_start_col, ts_end_row, ts_end_col =
+        vim_range_to_ts_range(start_row, start_col, end_row, end_col)
+    local root = utils.get_root(lang)
+    local scope = utils.get_scope_over_selection(root, start_row, start_col, end_row + 1, end_col, lang)
+
+    if scope == nil then
+        error("Scope is nil")
+    end
+
+    local local_defs = vim.tbl_filter(function(node)
+        return not utils.range_contains_node(node, ts_start_row, ts_start_col, ts_end_row, ts_end_col)
+    end, utils.get_locals_defs(scope, lang))
+
+    local function_args = utils.get_function_args(scope, lang)
+    local local_def_map = get_local_definitions(local_defs, function_args)
+    local local_references = utils.get_all_identifiers(scope, lang)
+    local selected_local_references = {}
+
+    for _, local_ref in pairs(local_references) do
+        local local_name = ts_utils.get_node_text(local_ref)[1]
+        if utils.range_contains_node(local_ref, ts_start_row, ts_start_col, ts_end_row, ts_end_col) and
+            local_def_map[local_name] then
+            selected_local_references[local_name] = true
+        end
+    end
+
+    -- TODO: Probably use text edit
+    local scope_range = get_scope_range(scope)
+
+    -- TODO: Polar, nvim_buf_get_lines doesn't actually get the highlighted
+    -- region, instead the highlighted rows
+    local function_name = vim.fn.input("106: Extract Function Name > ")
+
+    -- TODO: Polor, could you also make the variable that is returned the first
+    local text_edits = get_text_edits(selected_local_references, end_row, lang, start_col, start_row, end_col, scope_range, function_name)
+    vim.lsp.util.apply_text_edits(text_edits, 0)
+end
+
+REFACTORING.extract()
+
+return REFACTORING

lua/refactoring/106_spec.lua

@@ -0,0 +1,27 @@
+vim.cmd("set rtp+=" .. vim.loop.cwd())
+
+local eq = assert.are.same
+local refactoring = require("refactoring.106")
+
+describe("Refactoring", function()
+    it("should test selection grabbing", function(done)
+        vim.cmd(":new")
+        vim.bo.filetype = "lua"
+        vim.api.nvim_buf_set_lines(0, 0, 3, false, {
+            "local foo = 5",
+            "",
+            "foo = foo + 5 + foo",
+        })
+
+        eq(3, #vim.api.nvim_buf_get_lines(0, 0, 3, false))
+        vim.cmd(":norm! ggVjj")
+
+        eq(0, vim.fn.col("'<"), "Selections start at 0")
+        eq(#"foo = foo + 5 + foo" + 1, vim.fn.col("'>"), "Selection stops + 1 after last line.")
+
+        -- print("COL", vim.fn.col("'<"), vim.fn.col("'>"))
+    end)
+end)
+
+
+

lua/refactoring/dev.lua

@@ -0,0 +1,7 @@
+local M = {}
+M.reload = function()
+    require("plenary.reload").reload_module("refactoring")
+end
+
+return M
+

lua/refactoring/init.lua

@@ -0,0 +1,17 @@
+
+-- Array<node_wrapper>
+local intersect_nodes = function(nodes, row, col)
+    local found = {}
+    for idx = 1, #nodes do
+        local node = nodes[idx]
+        local sRow = node.dim.s.r
+        local sCol = node.dim.s.c
+        local eRow = node.dim.e.r
+        local eCol = node.dim.e.c
+        if utils.intersects(row, col, sRow, sCol, eRow, eCol) then
+            table.insert(found, node)
+        end
+    end
+    return found
+end
+

lua/refactoring/intersect.lua

@@ -0,0 +1,4 @@
+function find_identifiers()
+end
+
+

lua/refactoring/query.lua

@@ -0,0 +1,2 @@
+
+

lua/refactoring/refactoring.lua

@@ -0,0 +1,174 @@
+local ts_utils = require("nvim-treesitter.ts_utils")
+local ts_query = require("nvim-treesitter.query")
+local parsers = require("nvim-treesitter.parsers")
+local locals = require("nvim-treesitter.locals")
+
+local M = {}
+
+M.get_root = function(lang)
+    local parser = parsers.get_parser(0, lang)
+    return parser:parse()[1]:root()
+end
+
+M.get_bounded_query = function(query, lang, startR, stopR)
+    local success, parsed_query = pcall(function()
+        return vim.treesitter.parse_query(lang, query)
+    end)
+
+    if not success then
+        error("Unsuccessful successful first try")
+    end
+
+    local root = M.get_root(lang)
+
+    local out = {}
+    for match in ts_query.iter_prepared_matches(parsed_query, root, 0, startR - 1, stopR) do
+        locals.recurse_local_nodes(match, function(_, node, path)
+            table.insert(out, node)
+        end)
+    end
+    return out;
+end
+
+local refactor_constants = {
+    lua = {
+        scope = {
+            ["function"] = true,
+            ["function_definition"] = true
+        }
+    }
+}
+
+-- determines if a contains node b.
+-- @param a the containing node
+-- @param b the node to be contained
+function M.node_contains(a, b)
+    if a == nil or b == nil then
+        return false
+    end
+
+    local start_row, start_col, end_row, end_col = b:range()
+    return ts_utils.is_in_node_range(a, start_row, start_col) and
+           ts_utils.is_in_node_range(a, end_row, end_col)
+end
+
+-- determines if a node exists within a range.  Imagine a range selection
+-- across '<,'> and an identifier.  Does the identifier exist within the
+-- selection?
+--
+-- @param a the containing node
+-- @param b the node to be contained
+M.range_contains_node = function(node, start_row, start_col, end_row, end_col)
+    local node_start_row, node_start_col, node_end_row, node_end_col = node:range()
+
+    -- There are five possible conditions
+    -- 1. node start/end row are contained exclusively within the range.
+    -- 2. The range is a single line range
+    --   - the node start/end row must equal start_row and cols have to exist
+    --     within range, inclusive
+    -- 3. The node exists solely within the first line
+    --   - node_start_col has to be inclusive with start_col, end col doesn't
+    --     matter.
+    -- 4. The node exists solely within the last line
+    --   - node_start_col doesn't matter whereas node_end_col has to be
+    --     inclusive with end_col
+    -- 5. The node starts / ends on the same rows and has to have each column
+    --    considered
+    if start_row < node_start_row and end_row > node_end_row then
+        return true
+    elseif start_row == end_row then
+        return start_row == node_start_row and
+               end_row == node_end_row and
+               start_col <= node_start_col and
+               end_col >= node_end_col
+
+    elseif start_row == node_start_row and start_row == node_end_row then
+        return start_col <= node_start_col
+    elseif end_row == node_start_row and end_row == node_end_row then
+        return end_col >= node_end_col
+    elseif start_row <= node_start_row and end_row >= node_end_row then
+        return start_col <= node_start_col and end_col >= node_end_col
+    end
+
+    return false
+end
+
+M.get_scope_over_selection = function(root, start_line, start_col, end_line, end_col, lang)
+    local start_scope = M.get_scope(root, start_line, start_col, lang)
+    local end_scope = M.get_scope(root, end_line, end_col, lang)
+
+    if start_scope ~= end_scope then
+        error("Selection spans over two scopes, cannot determine scope")
+    end
+
+    return start_scope
+end
+
+M.get_scope = function(root, line, col, lang)
+    local function_scopes = {}
+    local query = vim.treesitter.get_query(lang, "locals")
+
+    for id, n, _ in query:iter_captures(root, 0, 0, -1) do
+        if query.captures[id] == "scope" and refactor_constants[lang].scope[n:type()] then
+            table.insert(function_scopes, n)
+        end
+    end
+
+    local out = nil
+    for _, scope in pairs(function_scopes) do
+        -- TODO: This is a confusing issue
+        -- should a scope that contains another scope but terminates at the
+        -- same point be the outer or inner?  Should potentially be considered
+        -- a list of scopes...
+        if ts_utils.is_in_node_range(scope, line, col) and
+            (out == nil or M.node_contains(out, scope)) then
+
+            out = scope
+        end
+    end
+
+    return out
+end
+
+local function get_refactoring_query(lang)
+    local query = vim.treesitter.get_query(lang, "refactoring")
+    if not query then
+        error("refactoring not supported in this language.  Please provide a queries/<lang>/refactoring.scm")
+    end
+    return query
+end
+
+local function pluck_by_capture(scope, lang, query, capture_name)
+    local local_defs = {}
+    local root = M.get_root(lang)
+    for id, node, _ in query:iter_captures(root, 0, 0, -1) do
+        if query.captures[id] == capture_name and M.node_contains(scope, node) then
+            table.insert(local_defs, node)
+        end
+    end
+
+    return local_defs
+end
+
+M.get_function_args = function(scope, lang)
+    return pluck_by_capture(scope, lang, get_refactoring_query(lang), "definition.function_argument")
+end
+
+M.get_locals_defs = function(scope, lang)
+    return pluck_by_capture(scope, lang, get_refactoring_query(lang), "definition.local_var")
+end
+
+M.get_all_identifiers  = function(scope, lang)
+    return pluck_by_capture(scope, lang, vim.treesitter.get_query(lang, "locals"), "reference")
+end
+
+-- is there a better way?
+M.range_to_table = function(node)
+    if node == nil then
+        return "range nil"
+    end
+    local a, b, c, d = node:range()
+    return {a, b, c, d}
+end
+
+return M