let return and elvis return support returned variable

AntoniRokitnicki · Feb 12, 2024 · 2fc82a8 · 2fc82a8
1 parent ee203d4
commit 2fc82a8
Show file tree

Hide file tree

Showing 7 changed files with 183 additions and 543 deletions.
diff --git a/src/com/intellij/advancedExpressionFolding/expression/custom/ElvisReturnNull.kt b/src/com/intellij/advancedExpressionFolding/expression/custom/ElvisReturnNull.kt
@@ -14,9 +14,9 @@ if (expression == null) {
 }</fold>
  */
 class ElvisReturnNull(
-        element: PsiElement, textRange: TextRange,
-        private val declaration: PsiElement, private val declarationRange: TextRange,
-        private val letElement: PsiElement, private val letRange: TextRange,
+    element: PsiElement, textRange: TextRange,
+    private val declaration: PsiElement, private val declarationRange: TextRange,
+    private val letElement: PsiElement, private val letRange: TextRange, val foldVariable: Boolean,
         ) : Expression(element, textRange) {
     override fun supportsFoldRegions(document: Document,
                                      parent: Expression?): Boolean {
@@ -25,9 +25,11 @@ class ElvisReturnNull(
 
     override fun buildFoldRegions(element: PsiElement, document: Document, parent: Expression?): Array<FoldingDescriptor> {
         val descriptors = mutableListOf<FoldingDescriptor>()
-        descriptors.add(
+        if (foldVariable) {
+            descriptors.add(
                 FoldingDescriptor(declaration.node, declarationRange,
-                        FoldingGroup.newGroup(ElvisReturnNull::class.java.name), ""))
+                    FoldingGroup.newGroup(ElvisReturnNull::class.java.name), ""))
+        }
         descriptors.add(
                 FoldingDescriptor(letElement.node, letRange,
                         FoldingGroup.newGroup(ElvisReturnNull::class.java.name), " ?: return null"))

diff --git a/src/com/intellij/advancedExpressionFolding/expression/custom/LetReturnIt.kt b/src/com/intellij/advancedExpressionFolding/expression/custom/LetReturnIt.kt
@@ -14,9 +14,9 @@ if (expression != null) {
 }</fold>
  */
 class LetReturnIt(
-        element: PsiElement, textRange: TextRange,
-        private val declaration: PsiElement, private val declarationRange: TextRange,
-        private val letElement: PsiElement, private val letRange: TextRange,
+    element: PsiElement, textRange: TextRange,
+    private val declaration: PsiElement, private val declarationRange: TextRange,
+    private val letElement: PsiElement, private val letRange: TextRange, val foldVariable: Boolean,
         ) : Expression(element, textRange) {
     override fun supportsFoldRegions(document: Document,
                                      parent: Expression?): Boolean {
@@ -25,9 +25,11 @@ class LetReturnIt(
 
     override fun buildFoldRegions(element: PsiElement, document: Document, parent: Expression?): Array<FoldingDescriptor> {
         val descriptors = mutableListOf<FoldingDescriptor>()
-        descriptors.add(
+        if (foldVariable) {
+            descriptors.add(
                 FoldingDescriptor(declaration.node, declarationRange,
-                        FoldingGroup.newGroup(LetReturnIt::class.java.name), ""))
+                    FoldingGroup.newGroup(LetReturnIt::class.java.name), ""))
+        }
         descriptors.add(
                 FoldingDescriptor(letElement.node, letRange,
                         FoldingGroup.newGroup(LetReturnIt::class.java.name), "?.let { return it }"))

diff --git a/src/com/intellij/advancedExpressionFolding/extension/LetReturnExt.kt b/src/com/intellij/advancedExpressionFolding/extension/LetReturnExt.kt
@@ -8,32 +8,32 @@ import com.intellij.psi.*
 import com.intellij.refactoring.suggested.endOffset
 import com.intellij.refactoring.suggested.startOffset
 
+private val PsiElement.prevRealSibling: PsiElement?
+    get() {
+        return generateSequence(this.prevSibling) { it.prevSibling }
+            .firstOrNull { it !is PsiWhiteSpace }
+    }
+
 class LetReturnExt : IExtension {
     companion object {
         @JvmStatic
         fun getIfExpression(element: PsiIfStatement): Expression? {
             if (!isKotlinQuickReturn()) {
                 return null
             }
-
-            val containingBlock = element.parent as? PsiCodeBlock ?: return null
-            val array =
-                containingBlock.statements.asInstance(PsiDeclarationStatement::class.java, PsiIfStatement::class.java)
-                    ?: return null
             /*
             Expression expression = ForStatementExpressionExt.getForStatementExpression((PsiForStatement) element, document);
             if (expression != null) {
                 return expression;
             }
             */
-            val declaration = array[0] as PsiDeclarationStatement
-            val ifNotNullStatement = array[1] as PsiIfStatement
-            val psiBinaryExpression = ifNotNullStatement.condition as? PsiBinaryExpression
+            val declaration = element.prevRealSibling as? PsiDeclarationStatement ?: return null
+            val psiBinaryExpression = element.condition as? PsiBinaryExpression ?: return null
 
             // 1) if (expression != null) {
 
             // 1a) expression
-            val variable = (psiBinaryExpression?.lOperand as? PsiReferenceExpression)?.resolve() as? PsiLocalVariable
+            val variable = (psiBinaryExpression.lOperand as? PsiReferenceExpression)?.resolve() as? PsiLocalVariable
                 ?: return null
 
             val token = psiBinaryExpression.operationTokenType
@@ -58,11 +58,11 @@ class LetReturnExt : IExtension {
 
             // 2) return expression;
             // 2a) no else
-            ifNotNullStatement.elseBranch?.let { return null }
+            element.elseBranch?.let { return null }
             // 2b) then
             // return local var
             val returnValue =
-                ifNotNullStatement.thenBranch.asInstance<PsiBlockStatement>()?.codeBlock?.statements?.asInstance(
+                element.thenBranch.asInstance<PsiBlockStatement>()?.codeBlock?.statements?.asInstance(
                     PsiReturnStatement::class.java
                 )
                     ?.firstOrNullIfNotEmpty()
@@ -98,25 +98,59 @@ class LetReturnExt : IExtension {
                 .asInstance<PsiLocalVariable>()?.initializer.asInstance<PsiMethodCallExpression>()
                 ?: return null
 
-            val ifParent = ifNotNullStatement.parent
+            val block = element.parent.asInstance<PsiCodeBlock>() ?: return null
+            val index = block.statements.indexOf(element)
+
+            val foldVariable = !isLocalVariableUsed(block, variable, index)
+
+            val ifParent = element.parent
             val methodCallComma = methodCall.nextSibling
 
             val declarationRange = (declaration.startOffset..methodCall.startOffset).toTextRange()
-            val letRange = (methodCallComma.startOffset..ifNotNullStatement.endOffset).toTextRange()
-            if (notEquals) {
-                return LetReturnIt(
-                    ifNotNullStatement, ifNotNullStatement.textRange,
+            val letRange = (methodCallComma.startOffset..element.endOffset).toTextRange()
+            return if (notEquals) {
+                LetReturnIt(
+                    element, element.textRange,
                     declaration, declarationRange,
-                    ifParent, letRange
+                    ifParent, letRange, foldVariable
                 )
             } else {
-                return ElvisReturnNull(
-                    ifNotNullStatement, ifNotNullStatement.textRange,
+                ElvisReturnNull(
+                    element, element.textRange,
                     declaration, declarationRange,
-                    ifParent, letRange
+                    ifParent, letRange, foldVariable
                 )
             }
         }
+
+        private fun isLocalVariableUsed(block: PsiCodeBlock, variable: PsiLocalVariable, afterIndex: Int): Boolean {
+            var used = false
+            var startChecking = false
+            block.accept(object : JavaRecursiveElementVisitor() {
+                override fun visitStatement(statement: PsiStatement) {
+                    val indexOf = block.statements.indexOf(statement)
+                    if (afterIndex < indexOf) {
+                        startChecking = true
+                    }
+                    if (!used) {
+                        super.visitStatement(statement)
+                    }
+                }
+
+                override fun visitReferenceExpression(expression: PsiReferenceExpression) {
+                    if (!startChecking) {
+                        return
+                    }
+                    if (expression.resolve() == variable) {
+                        used = true
+                    }
+                    if (!used) {
+                        super.visitReferenceExpression(expression)
+                    }
+                }
+            })
+            return used
+        }
     }
 
 }
diff --git a/test/com/intellij/advancedExpressionFolding/FoldingTest.java b/test/com/intellij/advancedExpressionFolding/FoldingTest.java
@@ -268,8 +268,4 @@ public void testLetReturnIt() {
         doFoldingTest();
     }
 
-    public void testElvisReturnNull() {
-        AdvancedExpressionFoldingSettings.getInstance().enableAll();
-        doFoldingTest();
-    }
 }
diff --git a/test/data/LetReturnIt.java b/test/data/LetReturnIt.java
@@ -0,0 +1,63 @@
+package data;
+
+import java.time.LocalDate;
+
+class LetReturnIt {
+    static String buildExpression(String someString) {
+        String var1 = getData(someString);
+        if (var1 != null) {
+            return var1;
+        }
+        String var2 = getData(someString);
+        if (var2 == null) {
+            return null;
+        }
+        String var4 = getData(someString);
+        if (var4 != null) {
+            return var4;
+        }
+        var4.toString();
+        String var5 = getData(someString);
+        if (var5 == null) {
+            return null;
+        }
+        System.out.println(var5 + "1");
+
+
+        String var6 = getData(someString);
+        if (var6 == null) {
+            return null;
+        }
+        while (true) {
+            if (LocalDate.now().isAfter(LocalDate.now())) {
+                if (var6 == null) {
+                    System.out.println("1");
+                }
+            }
+            break;
+        }
+
+
+        String var7 = getData(someString);
+        if (var7 != null) {
+            return var7;
+        }
+        new Thread(new Runnable() {
+            @Override
+            public void run() {
+                System.out.println(var7 + "1");
+            }
+        });
+        return null;
+    }
+
+    private static String getData(String someString) {
+        try {
+            return ClassLoader.getSystemResource("a").toString();
+        } catch (Exception e) {
+            return null;
+        }
+    }
+
+
+}