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):