diff --git a/lua/neogen.lua b/lua/neogen.lua
index 8e52492..6d3f3af 100644
--- a/lua/neogen.lua
+++ b/lua/neogen.lua
@@ -8,6 +8,7 @@ M.generate = function ()
     (function (parameters) @params)
     (function (return_statement) @return)
     (function_definition (parameters) @params)
+    (function_definition (if_statement (return_statement) @return))
     ]]
 
     -- Try to find the upper function