From 7b8631697d493be5323d04ae6e76624f01fa8ee5 Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Sun, 12 Sep 2021 17:45:36 +0200
Subject: [PATCH] (lua) Fix nested functions sometimes not being found

---
 lua/neogen.lua                    |  2 +-
 lua/neogen/configurations/lua.lua | 41 ++++++++++++++++++++-----------
 lua/neogen/utilities/cursor.lua   |  2 +-
 3 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/lua/neogen.lua b/lua/neogen.lua
index 9ca50bc..f6482dd 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -95,7 +95,7 @@ neogen.generate = function(opts)
                         end
 
                         -- Creates extmark for the end of the content
-                        neogen.utilities.cursor.create(to_place + #content+ 1, 0)
+                        neogen.utilities.cursor.create(to_place + #content + 1, 0)
 
                         neogen.utilities.cursor.jump({ first_time = true })
                     end
diff --git a/lua/neogen/configurations/lua.lua b/lua/neogen/configurations/lua.lua
index 1779904..b208f64 100644
--- a/lua/neogen/configurations/lua.lua
+++ b/lua/neogen/configurations/lua.lua
@@ -1,22 +1,30 @@
-local common_function_extractor = function(node)
+local function_extractor = function(node, type)
+    if not vim.tbl_contains({ "local", "function" }, type) then
+        error("Incorrect call for function_extractor")
+    end
+
     local tree = {
         {
-            retrieve = "first_recursive",
-            node_type = "function_definition",
+            retrieve = "first",
+            node_type = "parameters",
             subtree = {
-                {
-                    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 },
+                { retrieve = "all", node_type = "identifier", extract = true },
+                { retrieve = "all", node_type = "spread", extract = true },
             },
         },
+        { retrieve = "first", node_type = "return_statement", extract = true },
     }
 
+    if type == "local" then
+        tree = {
+            {
+                retrieve = "first_recursive",
+                node_type = "function_definition",
+                subtree = tree,
+            },
+        }
+    end
+
     local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
     local res = neogen.utilities.extractors:extract_from_matched(nodes)
 
@@ -59,14 +67,17 @@ return {
             -- When the function is inside one of those
             ["local_variable_declaration|field|variable_declaration"] = {
                 ["0"] = {
-                    extract = common_function_extractor,
+                    extract = function(node)
+                        return function_extractor(node, "local")
+                    end,
                 },
             },
             -- When the function is in the root tree
             ["function|local_function"] = {
                 ["0"] = {
-
-                    extract = common_function_extractor,
+                    extract = function(node)
+                        return function_extractor(node, "function")
+                    end,
                 },
             },
         },
diff --git a/lua/neogen/utilities/cursor.lua b/lua/neogen/utilities/cursor.lua
index 1845c60..5d6bc63 100644
--- a/lua/neogen/utilities/cursor.lua
+++ b/lua/neogen/utilities/cursor.lua
@@ -62,7 +62,7 @@ neogen.utilities.cursor.jumpable = function()
     end
     local cursor = vim.api.nvim_win_get_cursor(0)
     if cursor[1] > extm_list[#extm_list][2] or cursor[1] < extm_list[1][2] then
-      return false
+        return false
     end
     if #extm_list > 2 then
         return true