use control flow for collecting method return statements

Author: Haehnchen

src/main/java/fr/adrienbrault/idea/symfony2plugin/util/PhpElementsUtil.java

@@ -17,6 +17,10 @@
 import com.jetbrains.php.PhpClassHierarchyUtils;
 import com.jetbrains.php.PhpIndex;
 import com.jetbrains.php.codeInsight.PhpCodeInsightUtil;
+import com.jetbrains.php.codeInsight.controlFlow.PhpControlFlowUtil;
+import com.jetbrains.php.codeInsight.controlFlow.PhpInstructionProcessor;
+import com.jetbrains.php.codeInsight.controlFlow.instructions.PhpCallInstruction;
+import com.jetbrains.php.codeInsight.controlFlow.instructions.PhpReturnInstruction;
 import com.jetbrains.php.completion.PhpLookupElement;
 import com.jetbrains.php.lang.PhpLangUtil;
 import com.jetbrains.php.lang.PhpLanguage;
@@ -634,26 +638,27 @@ static public String getMethodReturnAsString(@NotNull PhpClass phpClass, @NotNul
      * Find a string return value of a method context "function() { return 'foo'}"
      */
     @NotNull
-    static public Collection<String> getMethodReturnAsStrings(@NotNull PhpClass phpClass, @NotNull String methodName) {
-
+    public static Collection<String> getMethodReturnAsStrings(@NotNull PhpClass phpClass, @NotNull String methodName) {
         Method method = phpClass.findMethodByName(methodName);
         if(method == null) {
             return Collections.emptyList();
         }
 
-        final Set<String> values = new HashSet<>();
-        method.acceptChildren(new PsiRecursiveElementWalkingVisitor() {
-            @Override
-            public void visitElement(PsiElement element) {
+        Set<String> values = new HashSet<>();
 
+        PhpControlFlowUtil.processFlow(method.getControlFlow(), new PhpInstructionProcessor() {
+            @Override
+            public boolean processReturnInstruction(PhpReturnInstruction instruction) {
+                PsiElement element = instruction.getArgument();
+                
                 if(PhpElementsUtil.getMethodReturnPattern().accepts(element)) {
                     String value = PhpElementsUtil.getStringValue(element);
-                    if(value != null && StringUtils.isNotBlank(value)) {
+                    if(StringUtils.isNotBlank(value)) {
                         values.add(value);
                     }
                 }
 
-                super.visitElement(element);
+                return super.processReturnInstruction(instruction);
             }
         });
 

src/test/java/fr/adrienbrault/idea/symfony2plugin/tests/util/PhpElementsUtilTest.java

@@ -1,6 +1,9 @@
 package fr.adrienbrault.idea.symfony2plugin.tests.util;
 
+import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.php.PhpIndex;
 import com.jetbrains.php.lang.PhpFileType;
@@ -350,6 +353,14 @@ public void testGetMethodParameterListStringPattern() {
 
         assertTrue(PhpElementsUtil.getMethodParameterListStringPattern().accepts(psiElement3));
     }
+    public void testGetMethodWithFirstStringOrNamedArgumentPattern2() {
+        PsiFile psiFile = myFixture.configureByFile("PhpElementsUtilReturn.php");
+        PhpClass phpClass = PsiTreeUtil.collectElementsOfType(psiFile, PhpClass.class).iterator().next();
+
+        assertContainsElements(PhpElementsUtil.getMethodReturnAsStrings(phpClass, "foo1"), "foo1_1", "foo1_2", "foo1_3", "foo1_4_const", "Foo");
+        assertContainsElements(PhpElementsUtil.getMethodReturnAsStrings(phpClass, "foo2"), "foo2_1", "foo2_2", "foo2_3", "foo2_4", "foo2_x");
+        assertContainsElements(PhpElementsUtil.getMethodReturnAsStrings(phpClass, "foo3"), "foo3_1", "foo3_2", "foo3_3");
+    }
 
     public void testGetMethodWithFirstStringOrNamedArgumentPattern() {
         StringLiteralExpression psiElement = (StringLiteralExpression) myFixture.configureByText(PhpFileType.INSTANCE, "<?php" +

src/test/java/fr/adrienbrault/idea/symfony2plugin/tests/util/fixtures/PhpElementsUtilReturn.php

@@ -0,0 +1,51 @@
+<?php
+
+class Foo
+{
+    const FOO = 'foo1_4_const';
+
+    public function foo1()
+    {
+        return 'foo1_1';
+
+        return 'foo1_2';
+
+        return 'foo1_3';
+
+        return self::FOO;
+
+        return Foo::class;
+    }
+
+    public function foo2()
+    {
+        if (true) {
+            return 'foo2_1';
+        }
+
+
+        if (true) {
+            return 'foo2_2';
+            if (true) {
+                return 'foo2_3';
+                if (true) {
+                    return 'foo2_4';
+                }
+            }
+        }
+
+        return 'foo2_x';
+    }
+
+    public function foo3()
+    {
+        switch (true) {
+            case 0:
+                return 'foo3_1';
+            default:
+                return 'foo3_2';
+        }
+
+        return 'foo3_3';
+    }
+}