From ce765cc43ca4f086741fb131ed3c0bcc701baa70 Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Wed, 25 Aug 2021 12:22:36 +0200
Subject: [PATCH] (lua) Add support for returns in local_functions

The extractor function being the same, I refactored the lua config file
a little bit
---
 lua/neogen/configurations/lua.lua | 74 +++++++++++++------------------
 1 file changed, 31 insertions(+), 43 deletions(-)

diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua
index 0f886ab..973bb04 100644
--- a/lua/neogen/configurations/lua.lua
+++ b/lua/neogen/configurations/lua.lua
@@ -1,56 +1,44 @@
+local common_function_extractor = function(node)
+    local tree = {
+        {
+            retrieve = "first",
+            node_type = "parameters",
+            subtree = {
+                { retrieve = "all", node_type = "identifier", extract = true },
+                { retrieve = "all", node_type = "spread", extract = true },
+            },
+        },
+        { retrieve = "first", node_type = "return_statement", extract = true },
+    }
+
+    local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
+    local res = neogen.utilities.extractors:extract_from_matched(nodes)
+
+    return {
+        parameters = res.identifier,
+        vararg = res.spread,
+        return_statement = res.return_statement,
+    }
+end
+
 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"] = {
-            -- Get second child from the parent node
-            ["2"] = {
-                -- It has to be of type "parameters"
-                match = "parameters",
-
-                extract = function(node)
-                    local tree = {
-                        { retrieve = "all", node_type = "identifier", extract = true },
-                        { retrieve = "all", node_type = "spread", extract = true },
-                    }
-                    local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
-                    local res = neogen.utilities.extractors:extract_from_matched(nodes)
-
-                    return {
-                        parameters = res.identifier,
-                        vararg = res.spread,
-                    }
-                end,
-            },
-        },
+        -- When the function is inside one of those
         ["local_variable_declaration|field|variable_declaration"] = {
             ["2"] = {
                 match = "function_definition",
 
-                extract = function(node)
-                    local tree = {
-                        {
-                            retrieve = "first",
-                            node_type = "parameters",
-                            subtree = {
-                                { retrieve = "all", node_type = "identifier", extract = true },
-                                { retrieve = "all", node_type = "spread", extract = true },
-                            },
-                        },
-                        { retrieve = "first", node_type = "return_statement", extract = true },
-                    }
+                extract = common_function_extractor,
+            },
+        },
+        -- When the function is in the root tree
+        ["function_definition|function|local_function"] = {
+            ["0"] = {
 
-                    local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
-                    local res = neogen.utilities.extractors:extract_from_matched(nodes)
-
-                    return {
-                        parameters = res.identifier,
-                        vararg = res.spread,
-                        return_statement = res.return_statement,
-                    }
-                end,
+                extract = common_function_extractor,
             },
         },
     },