From 14424a08872653421f51be4cbe5c4fe9a8dc75e1 Mon Sep 17 00:00:00 2001
From: Daniel Mathiot <d.danymat@gmail.com>
Date: Tue, 24 Aug 2021 12:44:33 +0200
Subject: [PATCH] (js) Add support for functions in statements

---
 lua/neogen/configurations/javascript.lua | 31 +++++++++++++++++++-----
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/lua/neogen/configurations/javascript.lua b/lua/neogen/configurations/javascript.lua
index 700174b..c796665 100644
--- a/lua/neogen/configurations/javascript.lua
+++ b/lua/neogen/configurations/javascript.lua
@@ -1,18 +1,37 @@
+local function_tree = {
+    { retrieve = "first", node_type = "formal_parameters", subtree = {
+        { retrieve = "all", node_type = "identifier", extract = true }
+    } },
+    { retrieve = "first", node_type = "statement_block", subtree = {
+        { retrieve = "first", node_type = "return_statement", extract = true }
+    } }
+}
 return {
-    parent = { "function_declaration" },
+    parent = { "function_declaration", "expression_statement" },
 
     data = {
         ["function_declaration"] = {
             ["0"] = {
 
+                extract = function (node)
+                    local results = {}
+                    local tree = function_tree
+                    local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)
+                    local res = neogen.utilities.extractors:extract_from_matched(nodes)
+
+                    results.parameters = res.identifier
+                    results.return_statement = res.return_statement
+                    return results
+                end
+            }
+        },
+        ["expression_statement"] = {
+            ["0"] = {
                 extract = function (node)
                     local results = {}
                     local tree = {
-                        { retrieve = "first", node_type = "formal_parameters", subtree = {
-                            { retrieve = "all", node_type = "identifier", extract = true }
-                        } },
-                        { retrieve = "first", node_type = "statement_block", subtree = {
-                            { retrieve = "first", node_type = "return_statement", extract = true }
+                        { retrieve = "all", node_type = "assignment_expression", subtree = {
+                            { retrieve = "all", node_type = "function", subtree = function_tree }
                         } }
                     }
                     local nodes = neogen.utilities.nodes:matching_nodes_from(node, tree)