From 86a600849f7c3b3b1a8c92509205154158e5f117 Mon Sep 17 00:00:00 2001
From: Vhyrro <vhyrro@gmail.com>
Date: Thu, 19 Aug 2021 16:43:35 +0200
Subject: [PATCH 1/6] feat: rewrite the entire core

---
 lua/neogen.lua        | 288 +++++++++++++++++++++++++++++++++++-------
 lua/neogen/config.lua |   5 -
 stylua.toml           |   6 +
 3 files changed, 251 insertions(+), 48 deletions(-)
 delete mode 100644 lua/neogen/config.lua
 create mode 100644 stylua.toml

diff --git a/lua/neogen.lua b/lua/neogen.lua
index 01ad38d..d843339 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -1,67 +1,269 @@
-local ts_utils = require("nvim-treesitter.ts_utils")
+local ok, ts_utils = pcall(require, "nvim-treesitter.ts_utils")
 
-neogen = {}
+assert(ok, "neogen requires nvim-treesitter to operate :(")
 
-local configuration = require('neogen.config')
+local neogen = {
+    utility = {
+        wrap = function(name, ...)
+            local args = { ... }
 
-neogen.generate = function ()
-    local comment = {}
+            return function()
+                name(table.unpack(args))
+            end
+        end,
 
-    -- Try to find the upper function
-    local cursor = ts_utils.get_node_at_cursor(0)
-    local function_node = cursor
-    while function_node ~= nil do
-        if function_node:type() == "function_definition" then break end
-        if function_node:type() == "function" then break end
-        if function_node:type() == "local_function" then break end
-        function_node = function_node:parent()
-    end
-    local line = ts_utils.get_node_range(function_node)
+        extract_children = function(_, name)
+            return function(node)
+                local result = {}
+                local split = vim.split(name, "|", true)
 
-    -- find the starting position in the line function
-    local line_content = vim.api.nvim_buf_get_lines(0, line, line+1, false)[1]
-    local offset = line_content:match("^%s+") or ""
+                for child in node:iter_children() do
+                    if vim.tbl_contains(split, child:type()) then
+                        table.insert(result, ts_utils.get_node_text(child)[1])
+                    end
+                end
 
-    local return_comment = offset .. "---@return "
-    local param_comment = offset .. "---@param "
+                return result
+            end
+        end,
 
-    -- Parse and iterate over each found query
-    local returned = vim.treesitter.get_query("lua", "neogen")
-    for id, node in returned:iter_captures(function_node) do
+        extract_children_from = function(self, name, nodes)
+            return function(node)
+                local result = {}
 
-        -- Try to add params
-        if returned.captures[id] == "params" then 
-            local params = ts_utils.get_node_text(node)[1]:sub(2,-2)
-            for p in string.gmatch(params, '[^,]+') do
-                p = p:gsub("%s+", "") -- remove trailing spaces
-                table.insert(comment, param_comment .. p .. " ")
+                for i, value in ipairs(nodes) do
+                    local child_node = node:named_child(i - 1)
+
+                    if value == "extract" then
+                        return self:extract_children(name)(child_node)
+                    else
+                        return self:extract_children_from(name, value)(node)
+                    end
+                end
+
+                return result
+            end
+        end,
+    },
+
+    default_locator = function(node_info, nodes_to_match)
+        if vim.tbl_contains(nodes_to_match, node_info.current:type()) then
+            return node_info.current
+        end
+
+        while node_info.current and not vim.tbl_contains(nodes_to_match, node_info.current:type()) do
+            node_info.current = node_info.current:parent()
+        end
+
+        return node_info.current
+    end,
+
+    default_granulator = function(parent_node, node_data)
+        local result = {}
+
+        for parent_type, child_data in pairs(node_data) do
+            local matches = vim.split(parent_type, "|", true)
+            if vim.tbl_contains(matches, parent_node:type()) then
+                for i, _ in pairs(node_data[parent_type]) do
+                    local data = child_data[i]
+
+                    local child_node = parent_node:named_child(tonumber(i) - 1)
+
+                    if not child_node then
+                        return
+                    end
+
+                    if child_node:type() == data.match or not data.match then
+                        local extract = {}
+
+                        if data.extract then
+                            extract = data.extract(child_node)
+
+                            if data.type then
+                                -- Extract information into a one-dimensional array
+                                local one_dimensional_arr = {}
+
+                                for _, values in pairs(extract) do
+                                    table.insert(one_dimensional_arr, values)
+                                end
+
+                                result[data.type] = one_dimensional_arr
+                            else
+                                for type, extracted_data in pairs(extract) do
+                                    result[type] = extracted_data
+                                end
+                            end
+                        else
+                            extract = ts_utils.get_node_text(child_node)
+                            result[data.type] = extract
+                        end
+                    end
+                end
             end
         end
 
-        -- Try to add return statement
-        if returned.captures[id] == "return" then
-            table.insert(comment, return_comment)
+        return result
+    end,
+
+    default_generator = function(parent, data, template)
+        local start_row, start_column, _, _ = ts_utils.get_node_range(parent)
+        local commentstring, generated_template = vim.trim(vim.api.nvim_buf_get_option(0, "commentstring"):format(""))
+
+        if not template then
+            generated_template = {
+                { nil, "" },
+                { "name", " @Summary " },
+                { "parameters", " @Param " },
+                { "return", " @Return " },
+            }
+        elseif type(template) == "function" then
+            generated_template = template(parent, commentstring, data)
+        else
+            generated_template = template
         end
-    end
 
-    -- At the end, add description annotation
-    table.insert(comment, 1, offset .. "---")
+        local function parse_generated_template()
+            local result = {}
+            local prefix = (" "):rep(start_column) .. commentstring
 
-    if #comment == 0 then return end
+            for _, values in ipairs(generated_template) do
+                local type = values[1]
 
-    -- Write on top of function
-    vim.fn.append(line, comment)
-    vim.fn.cursor(line+1, #comment[1])
-    vim.api.nvim_command('startinsert!')
+                if not type then
+                    table.insert(result, prefix .. values[2]:format(""))
+                else
+                    if data[type] then
+                        if #vim.tbl_values(data[type]) == 1 then
+                            table.insert(result, prefix .. values[2]:format(data[type][1]))
+                        else
+                            for _, value in ipairs(data[type]) do
+                                table.insert(result, prefix .. values[2]:format(value))
+                            end
+                        end
+                    end
+                end
+            end
+
+            return result
+        end
+
+        return start_row, parse_generated_template()
+    end,
+}
+
+-- TODO: Move code here
+neogen.generate = function(searcher, generator) end
+
+neogen.auto_generate = function(custom_template)
+    vim.treesitter.get_parser(0):for_each_tree(function(tree, language_tree)
+        local searcher = neogen.configuration.languages[language_tree:lang()]
+
+        if searcher then
+            searcher.locator = searcher.locator or neogen.default_locator
+            searcher.granulator = searcher.granulator or neogen.default_granulator
+            searcher.generator = searcher.generator or neogen.default_generator
+
+            local located_parent_node = searcher.locator({
+                root = tree:root(),
+                current = ts_utils.get_node_at_cursor(0),
+            }, searcher.parent)
+
+            if not located_parent_node then
+                return
+            end
+
+            local data = searcher.granulator(located_parent_node, searcher.data)
+
+            if data and not vim.tbl_isempty(data) then
+                local to_place, content = searcher.generator(
+                    located_parent_node,
+                    data,
+                    custom_template or searcher.template
+                )
+
+                vim.fn.append(to_place, content)
+            end
+        end
+    end)
 end
 
 function neogen.generate_command()
-    vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").generate()')
+    vim.api.nvim_command('command! -range -bar Neogen lua require("neogen").auto_generate()')
 end
 
 neogen.setup = function(opts)
-    local config = opts or configuration
-    if config.enabled == true then neogen.generate_command() end
+    neogen.configuration = vim.tbl_deep_extend("keep", opts or {}, {
+        -- DEFAULT CONFIGURATION
+        languages = {
+            lua = {
+                -- Search for these nodes
+                parent = { "function", "local_function", "local_variable_declaration", "field" },
+
+                -- Traverse down these nodes and extract the information as necessary
+                data = {
+                    ["function|local_function"] = {
+                        ["2"] = {
+                            match = "parameters",
+
+                            extract = function(node)
+                                local regular_params = neogen.utility:extract_children("identifier")(node)
+                                local varargs = neogen.utility:extract_children("spread")(node)
+
+                                return {
+                                    parameters = regular_params,
+                                    vararg = varargs,
+                                }
+                            end,
+                        },
+                    },
+                    ["local_variable_declaration|field"] = {
+                        ["2"] = {
+                            match = "function_definition",
+
+                            extract = function(node)
+                                local regular_params = neogen.utility:extract_children_from("identifier", {
+                                    [1] = "extract",
+                                })(node)
+
+                                local varargs = neogen.utility:extract_children_from("spread", {
+                                    [1] = "extract",
+                                })(node)
+
+                                return {
+                                    parameters = regular_params,
+                                    vararg = varargs,
+                                }
+                            end,
+                        },
+                    },
+                },
+
+                -- Custom lua locator that escapes from comments
+                locator = function(node_info, nodes_to_match)
+                    -- We're dealing with a lua comment and we need to escape its grasp
+                    if node_info.current:type() == "source" then
+                        local start_row, _, _, _ = ts_utils.get_node_range(node_info.current)
+                        vim.api.nvim_win_set_cursor(0, { start_row, 0 })
+                        node_info.current = ts_utils.get_node_at_cursor()
+                    end
+
+                    return neogen.default_locator(node_info, nodes_to_match)
+                end,
+
+                -- Use default granulator and generator
+                granulator = nil,
+                generator = nil,
+
+                template = {
+                    { nil, "-" },
+                    { "parameters", "-@param %s any" },
+                    { "vararg", "-@vararg any" },
+                },
+            },
+        },
+    })
+
+    neogen.generate_command()
 end
 
 return neogen
diff --git a/lua/neogen/config.lua b/lua/neogen/config.lua
deleted file mode 100644
index 78ea978..0000000
--- a/lua/neogen/config.lua
+++ /dev/null
@@ -1,5 +0,0 @@
-neogen.configuration = {
-    enabled = false,
-}
-
-return neogen
diff --git a/stylua.toml b/stylua.toml
new file mode 100644
index 0000000..bb258b9
--- /dev/null
+++ b/stylua.toml
@@ -0,0 +1,6 @@
+column_width = 120
+line_endings = "Unix"
+indent_type = "Spaces"
+indent_width = 4
+quote_style = "AutoPreferDouble"
+no_call_parentheses = false

From c4f3725af77faca45d6fcb9327d65e253738255b Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Thu, 19 Aug 2021 22:42:57 +0200
Subject: [PATCH 2/6] Refactor

This is a first iteration of refactoring the codebase, and generating
annotations along the way.
Added "variable_declaration" as parent node for lua
---
 lua/neogen.lua                     | 221 +----------------------------
 lua/neogen/configurations/lua.lua  |  65 +++++++++
 lua/neogen/generators/default.lua  |  46 ++++++
 lua/neogen/granulators/default.lua |  48 +++++++
 lua/neogen/locators/default.lua    |  17 +++
 lua/neogen/utility.lua             |  50 +++++++
 6 files changed, 233 insertions(+), 214 deletions(-)
 create mode 100644 lua/neogen/configurations/lua.lua
 create mode 100644 lua/neogen/generators/default.lua
 create mode 100644 lua/neogen/granulators/default.lua
 create mode 100644 lua/neogen/locators/default.lua
 create mode 100644 lua/neogen/utility.lua

diff --git a/lua/neogen.lua b/lua/neogen.lua
index d843339..2f30332 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -1,155 +1,7 @@
 local ok, ts_utils = pcall(require, "nvim-treesitter.ts_utils")
-
 assert(ok, "neogen requires nvim-treesitter to operate :(")
 
-local neogen = {
-    utility = {
-        wrap = function(name, ...)
-            local args = { ... }
-
-            return function()
-                name(table.unpack(args))
-            end
-        end,
-
-        extract_children = function(_, name)
-            return function(node)
-                local result = {}
-                local split = vim.split(name, "|", true)
-
-                for child in node:iter_children() do
-                    if vim.tbl_contains(split, child:type()) then
-                        table.insert(result, ts_utils.get_node_text(child)[1])
-                    end
-                end
-
-                return result
-            end
-        end,
-
-        extract_children_from = function(self, name, nodes)
-            return function(node)
-                local result = {}
-
-                for i, value in ipairs(nodes) do
-                    local child_node = node:named_child(i - 1)
-
-                    if value == "extract" then
-                        return self:extract_children(name)(child_node)
-                    else
-                        return self:extract_children_from(name, value)(node)
-                    end
-                end
-
-                return result
-            end
-        end,
-    },
-
-    default_locator = function(node_info, nodes_to_match)
-        if vim.tbl_contains(nodes_to_match, node_info.current:type()) then
-            return node_info.current
-        end
-
-        while node_info.current and not vim.tbl_contains(nodes_to_match, node_info.current:type()) do
-            node_info.current = node_info.current:parent()
-        end
-
-        return node_info.current
-    end,
-
-    default_granulator = function(parent_node, node_data)
-        local result = {}
-
-        for parent_type, child_data in pairs(node_data) do
-            local matches = vim.split(parent_type, "|", true)
-            if vim.tbl_contains(matches, parent_node:type()) then
-                for i, _ in pairs(node_data[parent_type]) do
-                    local data = child_data[i]
-
-                    local child_node = parent_node:named_child(tonumber(i) - 1)
-
-                    if not child_node then
-                        return
-                    end
-
-                    if child_node:type() == data.match or not data.match then
-                        local extract = {}
-
-                        if data.extract then
-                            extract = data.extract(child_node)
-
-                            if data.type then
-                                -- Extract information into a one-dimensional array
-                                local one_dimensional_arr = {}
-
-                                for _, values in pairs(extract) do
-                                    table.insert(one_dimensional_arr, values)
-                                end
-
-                                result[data.type] = one_dimensional_arr
-                            else
-                                for type, extracted_data in pairs(extract) do
-                                    result[type] = extracted_data
-                                end
-                            end
-                        else
-                            extract = ts_utils.get_node_text(child_node)
-                            result[data.type] = extract
-                        end
-                    end
-                end
-            end
-        end
-
-        return result
-    end,
-
-    default_generator = function(parent, data, template)
-        local start_row, start_column, _, _ = ts_utils.get_node_range(parent)
-        local commentstring, generated_template = vim.trim(vim.api.nvim_buf_get_option(0, "commentstring"):format(""))
-
-        if not template then
-            generated_template = {
-                { nil, "" },
-                { "name", " @Summary " },
-                { "parameters", " @Param " },
-                { "return", " @Return " },
-            }
-        elseif type(template) == "function" then
-            generated_template = template(parent, commentstring, data)
-        else
-            generated_template = template
-        end
-
-        local function parse_generated_template()
-            local result = {}
-            local prefix = (" "):rep(start_column) .. commentstring
-
-            for _, values in ipairs(generated_template) do
-                local type = values[1]
-
-                if not type then
-                    table.insert(result, prefix .. values[2]:format(""))
-                else
-                    if data[type] then
-                        if #vim.tbl_values(data[type]) == 1 then
-                            table.insert(result, prefix .. values[2]:format(data[type][1]))
-                        else
-                            for _, value in ipairs(data[type]) do
-                                table.insert(result, prefix .. values[2]:format(value))
-                            end
-                        end
-                    end
-                end
-            end
-
-            return result
-        end
-
-        return start_row, parse_generated_template()
-    end,
-}
+neogen = { }
 
 -- TODO: Move code here
 neogen.generate = function(searcher, generator) end
@@ -195,75 +47,16 @@ neogen.setup = function(opts)
     neogen.configuration = vim.tbl_deep_extend("keep", opts or {}, {
         -- DEFAULT CONFIGURATION
         languages = {
-            lua = {
-                -- Search for these nodes
-                parent = { "function", "local_function", "local_variable_declaration", "field" },
-
-                -- Traverse down these nodes and extract the information as necessary
-                data = {
-                    ["function|local_function"] = {
-                        ["2"] = {
-                            match = "parameters",
-
-                            extract = function(node)
-                                local regular_params = neogen.utility:extract_children("identifier")(node)
-                                local varargs = neogen.utility:extract_children("spread")(node)
-
-                                return {
-                                    parameters = regular_params,
-                                    vararg = varargs,
-                                }
-                            end,
-                        },
-                    },
-                    ["local_variable_declaration|field"] = {
-                        ["2"] = {
-                            match = "function_definition",
-
-                            extract = function(node)
-                                local regular_params = neogen.utility:extract_children_from("identifier", {
-                                    [1] = "extract",
-                                })(node)
-
-                                local varargs = neogen.utility:extract_children_from("spread", {
-                                    [1] = "extract",
-                                })(node)
-
-                                return {
-                                    parameters = regular_params,
-                                    vararg = varargs,
-                                }
-                            end,
-                        },
-                    },
-                },
-
-                -- Custom lua locator that escapes from comments
-                locator = function(node_info, nodes_to_match)
-                    -- We're dealing with a lua comment and we need to escape its grasp
-                    if node_info.current:type() == "source" then
-                        local start_row, _, _, _ = ts_utils.get_node_range(node_info.current)
-                        vim.api.nvim_win_set_cursor(0, { start_row, 0 })
-                        node_info.current = ts_utils.get_node_at_cursor()
-                    end
-
-                    return neogen.default_locator(node_info, nodes_to_match)
-                end,
-
-                -- Use default granulator and generator
-                granulator = nil,
-                generator = nil,
-
-                template = {
-                    { nil, "-" },
-                    { "parameters", "-@param %s any" },
-                    { "vararg", "-@vararg any" },
-                },
-            },
+            lua = require('neogen.configurations.lua'),
         },
     })
 
     neogen.generate_command()
 end
 
+require('neogen.utility')
+require('neogen.locators.default')
+require('neogen.granulators.default')
+require('neogen.generators.default')
+
 return neogen
diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua
new file mode 100644
index 0000000..8ab997e
--- /dev/null
+++ b/lua/neogen/configurations/lua.lua
@@ -0,0 +1,65 @@
+return {
+    -- Search for these nodes
+    parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration"},
+
+    -- Traverse down these nodes and extract the information as necessary
+    data = {
+        ["function|local_function"] = {
+            ["2"] = {
+                match = "parameters",
+
+                extract = function(node)
+                    local regular_params = neogen.utility:extract_children("identifier")(node)
+                    local varargs = neogen.utility:extract_children("spread")(node)
+
+                    return {
+                        parameters = regular_params,
+                        vararg = varargs,
+                    }
+                end,
+            },
+        },
+        ["local_variable_declaration|field|variable_declaration"] = {
+            ["2"] = {
+                match = "function_definition",
+
+                extract = function(node)
+                    local regular_params = neogen.utility:extract_children_from("identifier", {
+                        [1] = "extract",
+                    })(node)
+
+                    local varargs = neogen.utility:extract_children_from("spread", {
+                        [1] = "extract",
+                    })(node)
+
+                    return {
+                        parameters = regular_params,
+                        vararg = varargs,
+                    }
+                end,
+            },
+        },
+    },
+
+    -- Custom lua locator that escapes from comments
+    locator = function(node_info, nodes_to_match)
+        -- We're dealing with a lua comment and we need to escape its grasp
+        if node_info.current:type() == "source" then
+            local start_row, _, _, _ = ts_utils.get_node_range(node_info.current)
+            vim.api.nvim_win_set_cursor(0, { start_row, 0 })
+            node_info.current = ts_utils.get_node_at_cursor()
+        end
+
+        return neogen.default_locator(node_info, nodes_to_match)
+    end,
+
+    -- Use default granulator and generator
+    granulator = nil,
+    generator = nil,
+
+    template = {
+        { nil, "- " },
+        { "parameters", "- @param %s any" },
+        { "vararg", "- @vararg any" },
+    },
+}
diff --git a/lua/neogen/generators/default.lua b/lua/neogen/generators/default.lua
new file mode 100644
index 0000000..f2d03de
--- /dev/null
+++ b/lua/neogen/generators/default.lua
@@ -0,0 +1,46 @@
+local ts_utils = require("nvim-treesitter.ts_utils")
+
+neogen.default_generator = function(parent, data, template)
+    local start_row, start_column, _, _ = ts_utils.get_node_range(parent)
+    local commentstring, generated_template = vim.trim(vim.api.nvim_buf_get_option(0, "commentstring"):format(""))
+
+    if not template then
+        generated_template = {
+            { nil, "" },
+            { "name", " @Summary " },
+            { "parameters", " @Param " },
+            { "return", " @Return " },
+        }
+    elseif type(template) == "function" then
+        generated_template = template(parent, commentstring, data)
+    else
+        generated_template = template
+    end
+
+    local function parse_generated_template()
+        local result = {}
+        local prefix = (" "):rep(start_column) .. commentstring
+
+        for _, values in ipairs(generated_template) do
+            local type = values[1]
+
+            if not type then
+                table.insert(result, prefix .. values[2]:format(""))
+            else
+                if data[type] then
+                    if #vim.tbl_values(data[type]) == 1 then
+                        table.insert(result, prefix .. values[2]:format(data[type][1]))
+                    else
+                        for _, value in ipairs(data[type]) do
+                            table.insert(result, prefix .. values[2]:format(value))
+                        end
+                    end
+                end
+            end
+        end
+
+        return result
+    end
+
+    return start_row, parse_generated_template()
+end
diff --git a/lua/neogen/granulators/default.lua b/lua/neogen/granulators/default.lua
new file mode 100644
index 0000000..fc5e1a4
--- /dev/null
+++ b/lua/neogen/granulators/default.lua
@@ -0,0 +1,48 @@
+local ts_utils = require("nvim-treesitter.ts_utils")
+
+neogen.default_granulator = function(parent_node, node_data)
+    local result = {}
+
+    for parent_type, child_data in pairs(node_data) do
+        local matches = vim.split(parent_type, "|", true)
+        if vim.tbl_contains(matches, parent_node:type()) then
+            for i, _ in pairs(node_data[parent_type]) do
+                local data = child_data[i]
+
+                local child_node = parent_node:named_child(tonumber(i) - 1)
+
+                if not child_node then
+                    return
+                end
+
+                if child_node:type() == data.match or not data.match then
+                    local extract = {}
+
+                    if data.extract then
+                        extract = data.extract(child_node)
+
+                        if data.type then
+                            -- Extract information into a one-dimensional array
+                            local one_dimensional_arr = {}
+
+                            for _, values in pairs(extract) do
+                                table.insert(one_dimensional_arr, values)
+                            end
+
+                            result[data.type] = one_dimensional_arr
+                        else
+                            for type, extracted_data in pairs(extract) do
+                                result[type] = extracted_data
+                            end
+                        end
+                    else
+                        extract = ts_utils.get_node_text(child_node)
+                        result[data.type] = extract
+                    end
+                end
+            end
+        end
+    end
+
+    return result
+end
diff --git a/lua/neogen/locators/default.lua b/lua/neogen/locators/default.lua
new file mode 100644
index 0000000..b8cd697
--- /dev/null
+++ b/lua/neogen/locators/default.lua
@@ -0,0 +1,17 @@
+---The default locator tries to find one of the nodes to match in the current node
+--If it does not find one, will fetch the parents until he finds one
+---@param node_info table
+---@param nodes_to_match table
+neogen.default_locator = function(node_info, nodes_to_match)
+    -- If we find one of the wanted nodes in current one, return the current node
+    if vim.tbl_contains(nodes_to_match, node_info.current:type()) then
+        return node_info.current
+    end
+
+    -- Else, loop to parents until we find one of the nodes to match
+    while node_info.current and not vim.tbl_contains(nodes_to_match, node_info.current:type()) do
+        node_info.current = node_info.current:parent()
+    end
+
+    return node_info.current
+end
diff --git a/lua/neogen/utility.lua b/lua/neogen/utility.lua
new file mode 100644
index 0000000..de15db7
--- /dev/null
+++ b/lua/neogen/utility.lua
@@ -0,0 +1,50 @@
+local ts_utils = require("nvim-treesitter.ts_utils")
+
+neogen.utility = {
+    --- Return a function to extract content of required children from a node
+    --- @param _ any self
+    --- @param name string the children we want to extract (if multiple childrens, separate each one with "|")
+    --- @return function cb function taking a node and getting the content of each children we want from name
+    extract_children = function(_, name)
+        return function(node)
+            local result = {}
+            local split = vim.split(name, "|", true)
+
+            for child in node:iter_children() do
+                if vim.tbl_contains(split, child:type()) then
+                    table.insert(result, ts_utils.get_node_text(child)[1])
+                end
+            end
+
+            return result
+        end
+    end,
+
+    --- Extract content from specified children from a set of nodes
+    --- the nodes parameter can be a nested { [key] = value} with key being the 
+    --- * key: is which children we want to extract the values from (e.g first children is 1)
+    --- * value: "extract" or { [key] = value }. If value is "extract", it will extract the key child node
+    --- Example (extract the first child node from the first child node of the parent node):
+    --- [1] = {
+    ---  [1] = "extract"
+    --- }
+    --- @param name string the children we want to extract (if multiple children, separate each one with "|")
+    --- @param nodes table see description
+    extract_children_from = function(self, name, nodes)
+        return function(node)
+            local result = {}
+
+            for i, value in ipairs(nodes) do
+                local child_node = node:named_child(i - 1)
+
+                if value == "extract" then
+                    return self:extract_children(name)(child_node)
+                else
+                    return self:extract_children_from(name, value)(node)
+                end
+            end
+
+            return result
+        end
+    end,
+}

From f42e14f633ca9ccdcd922c6699f2612f38573024 Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Sat, 21 Aug 2021 07:31:57 +0200
Subject: [PATCH 3/6] Add documentation

Added docs in multiple places, refactored some functions and changed
places.
---
 lua/neogen.lua                     | 24 +++++++++++-------------
 lua/neogen/configurations/lua.lua  | 23 +++++++++--------------
 lua/neogen/granulators/default.lua | 10 +++++++---
 lua/neogen/locators/default.lua    |  9 +++++----
 lua/neogen/locators/lua.lua        | 12 ++++++++++++
 lua/neogen/utility.lua             | 12 ++++++------
 6 files changed, 50 insertions(+), 40 deletions(-)
 create mode 100644 lua/neogen/locators/lua.lua

diff --git a/lua/neogen.lua b/lua/neogen.lua
index 2f30332..d52341d 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -3,34 +3,32 @@ assert(ok, "neogen requires nvim-treesitter to operate :(")
 
 neogen = { }
 
--- TODO: Move code here
-neogen.generate = function(searcher, generator) end
-
 neogen.auto_generate = function(custom_template)
     vim.treesitter.get_parser(0):for_each_tree(function(tree, language_tree)
-        local searcher = neogen.configuration.languages[language_tree:lang()]
+        local language = neogen.configuration.languages[language_tree:lang()]
 
-        if searcher then
-            searcher.locator = searcher.locator or neogen.default_locator
-            searcher.granulator = searcher.granulator or neogen.default_granulator
-            searcher.generator = searcher.generator or neogen.default_generator
+        if language then
+            language.locator = language.locator or neogen.default_locator
+            language.granulator = language.granulator or neogen.default_granulator
+            language.generator = language.generator or neogen.default_generator
 
-            local located_parent_node = searcher.locator({
+            -- Use the language locator to locate one the the required parent nodes above the cursor
+            local located_parent_node = language.locator({
                 root = tree:root(),
                 current = ts_utils.get_node_at_cursor(0),
-            }, searcher.parent)
+            }, language.parent)
 
             if not located_parent_node then
                 return
             end
 
-            local data = searcher.granulator(located_parent_node, searcher.data)
+            local data = language.granulator(located_parent_node, language.data)
 
             if data and not vim.tbl_isempty(data) then
-                local to_place, content = searcher.generator(
+                local to_place, content = language.generator(
                     located_parent_node,
                     data,
-                    custom_template or searcher.template
+                    custom_template or language.template
                 )
 
                 vim.fn.append(to_place, content)
diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua
index 8ab997e..13cd137 100644
--- a/lua/neogen/configurations/lua.lua
+++ b/lua/neogen/configurations/lua.lua
@@ -1,3 +1,5 @@
+local ts_utils = require("nvim-treesitter.ts_utils")
+
 return {
     -- Search for these nodes
     parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration"},
@@ -5,7 +7,9 @@ return {
     -- Traverse down these nodes and extract the information as necessary
     data = {
         ["function|local_function"] = {
+            -- Get second child from the parent node
             ["2"] = {
+                -- It has to be of type "parameters"
                 match = "parameters",
 
                 extract = function(node)
@@ -24,13 +28,13 @@ return {
                 match = "function_definition",
 
                 extract = function(node)
-                    local regular_params = neogen.utility:extract_children_from("identifier", {
+                    local regular_params = neogen.utility:extract_children_from({
                         [1] = "extract",
-                    })(node)
+                    }, "identifier")(node)
 
-                    local varargs = neogen.utility:extract_children_from("spread", {
+                    local varargs = neogen.utility:extract_children_from({
                         [1] = "extract",
-                    })(node)
+                    }, "spread")(node)
 
                     return {
                         parameters = regular_params,
@@ -42,16 +46,7 @@ return {
     },
 
     -- Custom lua locator that escapes from comments
-    locator = function(node_info, nodes_to_match)
-        -- We're dealing with a lua comment and we need to escape its grasp
-        if node_info.current:type() == "source" then
-            local start_row, _, _, _ = ts_utils.get_node_range(node_info.current)
-            vim.api.nvim_win_set_cursor(0, { start_row, 0 })
-            node_info.current = ts_utils.get_node_at_cursor()
-        end
-
-        return neogen.default_locator(node_info, nodes_to_match)
-    end,
+    locator = require('neogen.locators.lua'),
 
     -- Use default granulator and generator
     granulator = nil,
diff --git a/lua/neogen/granulators/default.lua b/lua/neogen/granulators/default.lua
index fc5e1a4..81e6ec9 100644
--- a/lua/neogen/granulators/default.lua
+++ b/lua/neogen/granulators/default.lua
@@ -5,10 +5,11 @@ neogen.default_granulator = function(parent_node, node_data)
 
     for parent_type, child_data in pairs(node_data) do
         local matches = vim.split(parent_type, "|", true)
-        if vim.tbl_contains(matches, parent_node:type()) then
-            for i, _ in pairs(node_data[parent_type]) do
-                local data = child_data[i]
 
+        -- Look if the parent node is one of the matches
+        if vim.tbl_contains(matches, parent_node:type()) then
+            -- For each child_data in the matched parent node
+            for i, data in pairs(child_data) do
                 local child_node = parent_node:named_child(tonumber(i) - 1)
 
                 if not child_node then
@@ -19,8 +20,10 @@ neogen.default_granulator = function(parent_node, node_data)
                     local extract = {}
 
                     if data.extract then
+                        -- Extract content from it { [type] = { data } }
                         extract = data.extract(child_node)
 
+                        -- All extracted values are created added in result, like so: [data.type] = { extract }
                         if data.type then
                             -- Extract information into a one-dimensional array
                             local one_dimensional_arr = {}
@@ -36,6 +39,7 @@ neogen.default_granulator = function(parent_node, node_data)
                             end
                         end
                     else
+                        -- if not extract function, get the text from the node (required: data.type)
                         extract = ts_utils.get_node_text(child_node)
                         result[data.type] = extract
                     end
diff --git a/lua/neogen/locators/default.lua b/lua/neogen/locators/default.lua
index b8cd697..dcb99c9 100644
--- a/lua/neogen/locators/default.lua
+++ b/lua/neogen/locators/default.lua
@@ -1,7 +1,8 @@
----The default locator tries to find one of the nodes to match in the current node
---If it does not find one, will fetch the parents until he finds one
----@param node_info table
----@param nodes_to_match table
+--- The default locator tries to find one of the nodes to match in the current node
+--- If it does not find one, will fetch the parents until he finds one
+--- @param node_info table a node informations
+--- @param nodes_to_match table a list of parent nodes to match
+--- @return node_info.current node one of the nodes to match directly above the given node
 neogen.default_locator = function(node_info, nodes_to_match)
     -- If we find one of the wanted nodes in current one, return the current node
     if vim.tbl_contains(nodes_to_match, node_info.current:type()) then
diff --git a/lua/neogen/locators/lua.lua b/lua/neogen/locators/lua.lua
new file mode 100644
index 0000000..07b3956
--- /dev/null
+++ b/lua/neogen/locators/lua.lua
@@ -0,0 +1,12 @@
+local ts_utils = require("nvim-treesitter.ts_utils")
+
+return function(node_info, nodes_to_match)
+    -- We're dealing with a lua comment and we need to escape its grasp
+    if node_info.current:type() == "source" then
+        local start_row, _, _, _ = ts_utils.get_node_range(node_info.current)
+        vim.api.nvim_win_set_cursor(0, { start_row, 0 })
+        node_info.current = ts_utils.get_node_at_cursor()
+    end
+
+    return neogen.default_locator(node_info, nodes_to_match)
+end
diff --git a/lua/neogen/utility.lua b/lua/neogen/utility.lua
index de15db7..c495d5a 100644
--- a/lua/neogen/utility.lua
+++ b/lua/neogen/utility.lua
@@ -21,26 +21,26 @@ neogen.utility = {
     end,
 
     --- Extract content from specified children from a set of nodes
-    --- the nodes parameter can be a nested { [key] = value} with key being the 
+    --- the tree parameter can be a nested { [key] = value} with key being the 
     --- * key: is which children we want to extract the values from (e.g first children is 1)
     --- * value: "extract" or { [key] = value }. If value is "extract", it will extract the key child node
     --- Example (extract the first child node from the first child node of the parent node):
     --- [1] = {
     ---  [1] = "extract"
     --- }
+    --- @param tree table see description
     --- @param name string the children we want to extract (if multiple children, separate each one with "|")
-    --- @param nodes table see description
-    extract_children_from = function(self, name, nodes)
+    extract_children_from = function(self, tree, name)
         return function(node)
             local result = {}
 
-            for i, value in ipairs(nodes) do
+            for i, subtree in ipairs(tree) do
                 local child_node = node:named_child(i - 1)
 
-                if value == "extract" then
+                if subtree == "extract" then
                     return self:extract_children(name)(child_node)
                 else
-                    return self:extract_children_from(name, value)(node)
+                    return self:extract_children_from(subtree, name)(node)
                 end
             end
 

From ebeb5aa329d6a9b8319bf7a71552a89de1bdbf8e Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Sat, 21 Aug 2021 16:47:49 +0200
Subject: [PATCH 4/6] Refactor

More documentation, changed function signature
---
 lua/neogen.lua                     | 5 ++++-
 lua/neogen/configurations/lua.lua  | 1 +
 lua/neogen/generators/default.lua  | 8 ++++++++
 lua/neogen/granulators/default.lua | 3 +++
 lua/neogen/utility.lua             | 2 +-
 5 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/lua/neogen.lua b/lua/neogen.lua
index d52341d..e8b8e45 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -12,7 +12,7 @@ neogen.auto_generate = function(custom_template)
             language.granulator = language.granulator or neogen.default_granulator
             language.generator = language.generator or neogen.default_generator
 
-            -- Use the language locator to locate one the the required parent nodes above the cursor
+            -- Use the language locator to locate one of the required parent nodes above the cursor
             local located_parent_node = language.locator({
                 root = tree:root(),
                 current = ts_utils.get_node_at_cursor(0),
@@ -22,15 +22,18 @@ neogen.auto_generate = function(custom_template)
                 return
             end
 
+            -- Use the language granulator to get the required content inside the node found with the locator
             local data = language.granulator(located_parent_node, language.data)
 
             if data and not vim.tbl_isempty(data) then
+                -- Will try to generate the documentation from a template and the data found from the granulator 
                 local to_place, content = language.generator(
                     located_parent_node,
                     data,
                     custom_template or language.template
                 )
 
+                -- Append the annotation in required place
                 vim.fn.append(to_place, content)
             end
         end
diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua
index 13cd137..305ad58 100644
--- a/lua/neogen/configurations/lua.lua
+++ b/lua/neogen/configurations/lua.lua
@@ -36,6 +36,7 @@ return {
                         [1] = "extract",
                     }, "spread")(node)
 
+
                     return {
                         parameters = regular_params,
                         vararg = varargs,
diff --git a/lua/neogen/generators/default.lua b/lua/neogen/generators/default.lua
index f2d03de..eb73d78 100644
--- a/lua/neogen/generators/default.lua
+++ b/lua/neogen/generators/default.lua
@@ -1,10 +1,17 @@
 local ts_utils = require("nvim-treesitter.ts_utils")
 
+---Default Generator:
+---Uses the provided template to format the annotations with data found by the granulator
+--- @param parent userdata the node used to generate the annotations
+--- @param data table the data from the granulator, which is a set of [type] = results
+--- @param template table a template from the configuration
+--- @return table { line, content }, with line being the line to append the content
 neogen.default_generator = function(parent, data, template)
     local start_row, start_column, _, _ = ts_utils.get_node_range(parent)
     local commentstring, generated_template = vim.trim(vim.api.nvim_buf_get_option(0, "commentstring"):format(""))
 
     if not template then
+        -- Default template
         generated_template = {
             { nil, "" },
             { "name", " @Summary " },
@@ -12,6 +19,7 @@ neogen.default_generator = function(parent, data, template)
             { "return", " @Return " },
         }
     elseif type(template) == "function" then
+        -- You can also pass a function as a template
         generated_template = template(parent, commentstring, data)
     else
         generated_template = template
diff --git a/lua/neogen/granulators/default.lua b/lua/neogen/granulators/default.lua
index 81e6ec9..74bcb84 100644
--- a/lua/neogen/granulators/default.lua
+++ b/lua/neogen/granulators/default.lua
@@ -1,5 +1,8 @@
 local ts_utils = require("nvim-treesitter.ts_utils")
 
+--- Tries to use the configuration to find all required content nodes from the parent node 
+--- @param parent_node userdata the node found by the locator
+--- @param node_data table the data from configurations[lang].data
 neogen.default_granulator = function(parent_node, node_data)
     local result = {}
 
diff --git a/lua/neogen/utility.lua b/lua/neogen/utility.lua
index c495d5a..91f4849 100644
--- a/lua/neogen/utility.lua
+++ b/lua/neogen/utility.lua
@@ -20,7 +20,7 @@ neogen.utility = {
         end
     end,
 
-    --- Extract content from specified children from a set of nodes
+    --- Extract content from specified children from a tree
     --- the tree parameter can be a nested { [key] = value} with key being the 
     --- * key: is which children we want to extract the values from (e.g first children is 1)
     --- * value: "extract" or { [key] = value }. If value is "extract", it will extract the key child node

From ec09985e9dec1e07c50c12525b601dc9e106218b Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Sat, 21 Aug 2021 18:05:03 +0200
Subject: [PATCH 5/6] Add return statement for lua

At the moment, it only extracts return statements from direct children
of the function_definition
---
 after/queries/lua/neogen.scm       |  7 -------
 lua/neogen.lua                     | 19 +++++++++++--------
 lua/neogen/configurations/lua.lua  |  7 +++++--
 lua/neogen/granulators/default.lua |  2 +-
 lua/neogen/utility.lua             |  2 +-
 5 files changed, 18 insertions(+), 19 deletions(-)
 delete mode 100644 after/queries/lua/neogen.scm

diff --git a/after/queries/lua/neogen.scm b/after/queries/lua/neogen.scm
deleted file mode 100644
index 6d72a7f..0000000
--- a/after/queries/lua/neogen.scm
+++ /dev/null
@@ -1,7 +0,0 @@
-(function (parameters) @params)
-(function_definition (parameters) @params)
-(local_function (parameters) @params)
-
-(function (return_statement) @return)
-(function_definition (if_statement (return_statement) @return))
-(local_function (return_statement) @return)
diff --git a/lua/neogen.lua b/lua/neogen.lua
index e8b8e45..8349ce1 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -1,7 +1,15 @@
 local ok, ts_utils = pcall(require, "nvim-treesitter.ts_utils")
 assert(ok, "neogen requires nvim-treesitter to operate :(")
 
-neogen = { }
+neogen = {}
+
+require("neogen.utility")
+
+-- Require defaults
+require("neogen.locators.default")
+require("neogen.granulators.default")
+require("neogen.generators.default")
+
 
 neogen.auto_generate = function(custom_template)
     vim.treesitter.get_parser(0):for_each_tree(function(tree, language_tree)
@@ -26,7 +34,7 @@ neogen.auto_generate = function(custom_template)
             local data = language.granulator(located_parent_node, language.data)
 
             if data and not vim.tbl_isempty(data) then
-                -- Will try to generate the documentation from a template and the data found from the granulator 
+                -- Will try to generate the documentation from a template and the data found from the granulator
                 local to_place, content = language.generator(
                     located_parent_node,
                     data,
@@ -48,16 +56,11 @@ neogen.setup = function(opts)
     neogen.configuration = vim.tbl_deep_extend("keep", opts or {}, {
         -- DEFAULT CONFIGURATION
         languages = {
-            lua = require('neogen.configurations.lua'),
+            lua = require("neogen.configurations.lua"),
         },
     })
 
     neogen.generate_command()
 end
 
-require('neogen.utility')
-require('neogen.locators.default')
-require('neogen.granulators.default')
-require('neogen.generators.default')
-
 return neogen
diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua
index 305ad58..f391678 100644
--- a/lua/neogen/configurations/lua.lua
+++ b/lua/neogen/configurations/lua.lua
@@ -2,7 +2,7 @@ local ts_utils = require("nvim-treesitter.ts_utils")
 
 return {
     -- Search for these nodes
-    parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration"},
+    parent = { "function", "local_function", "local_variable_declaration", "field", "variable_declaration" },
 
     -- Traverse down these nodes and extract the information as necessary
     data = {
@@ -36,10 +36,12 @@ return {
                         [1] = "extract",
                     }, "spread")(node)
 
+                    local return_statement = neogen.utility:extract_children("return_statement")(node)
 
                     return {
                         parameters = regular_params,
                         vararg = varargs,
+                        return_statement = return_statement
                     }
                 end,
             },
@@ -47,7 +49,7 @@ return {
     },
 
     -- Custom lua locator that escapes from comments
-    locator = require('neogen.locators.lua'),
+    locator = require("neogen.locators.lua"),
 
     -- Use default granulator and generator
     granulator = nil,
@@ -57,5 +59,6 @@ return {
         { nil, "- " },
         { "parameters", "- @param %s any" },
         { "vararg", "- @vararg any" },
+        { "return_statement", "- @return any" }
     },
 }
diff --git a/lua/neogen/granulators/default.lua b/lua/neogen/granulators/default.lua
index 74bcb84..2a373aa 100644
--- a/lua/neogen/granulators/default.lua
+++ b/lua/neogen/granulators/default.lua
@@ -1,6 +1,6 @@
 local ts_utils = require("nvim-treesitter.ts_utils")
 
---- Tries to use the configuration to find all required content nodes from the parent node 
+--- Tries to use the configuration to find all required content nodes from the parent node
 --- @param parent_node userdata the node found by the locator
 --- @param node_data table the data from configurations[lang].data
 neogen.default_granulator = function(parent_node, node_data)
diff --git a/lua/neogen/utility.lua b/lua/neogen/utility.lua
index 91f4849..8fa0f8a 100644
--- a/lua/neogen/utility.lua
+++ b/lua/neogen/utility.lua
@@ -21,7 +21,7 @@ neogen.utility = {
     end,
 
     --- Extract content from specified children from a tree
-    --- the tree parameter can be a nested { [key] = value} with key being the 
+    --- the tree parameter can be a nested { [key] = value} with key being the
     --- * key: is which children we want to extract the values from (e.g first children is 1)
     --- * value: "extract" or { [key] = value }. If value is "extract", it will extract the key child node
     --- Example (extract the first child node from the first child node of the parent node):

From 96319356fa9dc319b8e974ea30839fe273b7f007 Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Sat, 21 Aug 2021 18:20:29 +0200
Subject: [PATCH 6/6] Add input_after_comment configuration

When set to true (default true), the cursor will go back to the start of
the added comments and trigger insert mode
---
 lua/neogen.lua                    | 9 ++++++++-
 lua/neogen/generators/default.lua | 2 +-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/lua/neogen.lua b/lua/neogen.lua
index 8349ce1..7e67e17 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -35,7 +35,7 @@ neogen.auto_generate = function(custom_template)
 
             if data and not vim.tbl_isempty(data) then
                 -- Will try to generate the documentation from a template and the data found from the granulator
-                local to_place, content = language.generator(
+                local to_place, start_column, content = language.generator(
                     located_parent_node,
                     data,
                     custom_template or language.template
@@ -43,6 +43,12 @@ neogen.auto_generate = function(custom_template)
 
                 -- Append the annotation in required place
                 vim.fn.append(to_place, content)
+
+                -- Place cursor after annotations ans start editing
+                if neogen.configuration.input_after_comment == true then
+                    vim.fn.cursor(to_place+1, start_column)
+                    vim.api.nvim_command('startinsert!')
+                end
             end
         end
     end)
@@ -54,6 +60,7 @@ end
 
 neogen.setup = function(opts)
     neogen.configuration = vim.tbl_deep_extend("keep", opts or {}, {
+        input_after_comment = true,
         -- DEFAULT CONFIGURATION
         languages = {
             lua = require("neogen.configurations.lua"),
diff --git a/lua/neogen/generators/default.lua b/lua/neogen/generators/default.lua
index eb73d78..aa5e80f 100644
--- a/lua/neogen/generators/default.lua
+++ b/lua/neogen/generators/default.lua
@@ -50,5 +50,5 @@ neogen.default_generator = function(parent, data, template)
         return result
     end
 
-    return start_row, parse_generated_template()
+    return start_row, start_column, parse_generated_template()
 end