[groovy] IDEA-202693 Support lambda in control flow builder

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/api/GrFunctionalExpression.java

@@ -1,10 +1,12 @@
 // Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package org.jetbrains.plugins.groovy.lang.psi.api;
 
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiType;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrParameterListOwner;
+import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrStatement;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrExpression;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.params.GrParameter;
 
@@ -17,4 +19,10 @@ public interface GrFunctionalExpression extends GrExpression, GrParameterListOwn
 
   @Nullable
   PsiType getReturnType();
+
+  @NotNull
+  GrStatement[] getStatements();
+
+  @Nullable
+  PsiElement getArrow();
 }

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/api/GrLambdaBody.java

@@ -1,10 +1,19 @@
 // Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package org.jetbrains.plugins.groovy.lang.psi.api;
 
+import com.intellij.psi.PsiType;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.groovy.lang.psi.GrControlFlowOwner;
 
 /**
  * Represents a Groovy lambda expression body.
  */
 public interface GrLambdaBody extends GrControlFlowOwner {
+
+  @NotNull
+  GrLambdaExpression getLambdaExpression();
+
+  @Nullable
+  PsiType getReturnType();
 }

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/api/GrLambdaExpression.java

@@ -1,6 +1,8 @@
 // Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package org.jetbrains.plugins.groovy.lang.psi.api;
 
+import com.intellij.psi.*;
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 /**
@@ -14,4 +16,8 @@ public interface GrLambdaExpression extends GrFunctionalExpression {
    */
   @Nullable
   GrLambdaBody getBody();
+
+  @Override
+  @NotNull
+  PsiElement getArrow();
 }

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/api/statements/blocks/GrClosableBlock.java

@@ -40,6 +40,7 @@ public interface GrClosableBlock extends GrFunctionalExpression, GrCodeBlock {
   @Override
   GrParameter[] getAllParameters();
 
+  @Override
   @Nullable
   PsiElement getArrow();
 

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/controlFlow/Instruction.java

@@ -23,6 +23,7 @@ import org.jetbrains.annotations.Nullable;
  * @author ven
  */
 public interface Instruction {
+  Instruction[] EMPTY_ARRAY = new Instruction[0];
 
   @NotNull
   Iterable<Instruction> successors(@NotNull CallEnvironment environment);

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/controlFlow/impl/ControlFlowBuilder.java

@@ -17,11 +17,11 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.plugins.groovy.codeInspection.utils.ControlFlowUtils;
 import org.jetbrains.plugins.groovy.lang.lexer.GroovyTokenTypes;
+import org.jetbrains.plugins.groovy.lang.psi.GrControlFlowOwner;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyFileBase;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyPsiElement;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyRecursiveElementVisitor;
-import org.jetbrains.plugins.groovy.lang.psi.api.GrInExpression;
-import org.jetbrains.plugins.groovy.lang.psi.api.GrTryResourceList;
+import org.jetbrains.plugins.groovy.lang.psi.api.*;
 import org.jetbrains.plugins.groovy.lang.psi.api.auxiliary.GrCondition;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.*;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.blocks.GrClosableBlock;
@@ -99,8 +99,9 @@ public class ControlFlowBuilder extends GroovyRecursiveElementVisitor {
     final PsiElement lbrace = block.getLBrace();
     if (lbrace != null && parent instanceof GrMethod) {
       for (GrParameter parameter : ((GrMethod)parent).getParameters()) {
-        if (myPolicy.isVariableInitialized(parameter)) {
-          addNode(new ReadWriteVariableInstruction(parameter.getName(), parameter, ReadWriteVariableInstruction.WRITE));
+        String parameterName = parameter.getName();
+        if (myPolicy.isVariableInitialized(parameter) && parameterName != null) {
+          addNode(new ReadWriteVariableInstruction(parameterName, parameter, ReadWriteVariableInstruction.WRITE));
         }
       }
     }
@@ -114,6 +115,12 @@ public class ControlFlowBuilder extends GroovyRecursiveElementVisitor {
     }
   }
 
+  @Override
+  public void visitLambdaBody(@NotNull GrLambdaBody body) {
+    addFunctionalExpressionParameters(body.getLambdaExpression());
+    addControlFlowInstructions(body);
+  }
+
   @Override
   public void visitFile(@NotNull GroovyFileBase file) {
     super.visitFile(file);
@@ -133,14 +140,20 @@ public class ControlFlowBuilder extends GroovyRecursiveElementVisitor {
   }
 
   public Instruction[] buildControlFlow(GroovyPsiElement scope) {
+    if (scope instanceof GrLambdaExpression) {
+      GrLambdaBody body = ((GrLambdaExpression)scope).getBody();
+      return body != null ? body.getControlFlow() : Instruction.EMPTY_ARRAY;
+    }
     myFinallyCount = 0;
     myInstructionNumber = 0;
 
     myScope = scope;
 
     startNode(null);
+
     if (scope instanceof GrClosableBlock) {
-      buildFlowForClosure((GrClosableBlock)scope);
+      addFunctionalExpressionParameters((GrFunctionalExpression)scope);
+      addControlFlowInstructions((GrControlFlowOwner)scope);
     }
     else {
       scope.accept(this);
@@ -150,7 +163,7 @@ public class ControlFlowBuilder extends GroovyRecursiveElementVisitor {
     checkPending(end); //collect return edges
 
 
-    return assertValidPsi(myInstructions.toArray(new Instruction[0]));
+    return assertValidPsi(myInstructions.toArray(Instruction.EMPTY_ARRAY));
   }
 
   public static Instruction[] assertValidPsi(Instruction[] instructions) {
@@ -163,16 +176,9 @@ public class ControlFlowBuilder extends GroovyRecursiveElementVisitor {
     return instructions;
   }
 
-  private void buildFlowForClosure(final GrClosableBlock closure) {
-    for (GrParameter parameter : closure.getAllParameters()) {
-      if (myPolicy.isVariableInitialized(parameter)) {
-        addNode(new ReadWriteVariableInstruction(parameter.getName(), parameter, ReadWriteVariableInstruction.WRITE));
-      }
-    }
 
-    addNode(new ReadWriteVariableInstruction("owner", closure.getLBrace(), ReadWriteVariableInstruction.WRITE));
-
-    PsiElement child = closure.getFirstChild();
+  private void addControlFlowInstructions(final GrControlFlowOwner owner) {
+    PsiElement child = owner.getFirstChild();
     while (child != null) {
       if (child instanceof GroovyPsiElement) {
         ((GroovyPsiElement)child).accept(this);
@@ -180,12 +186,25 @@ public class ControlFlowBuilder extends GroovyRecursiveElementVisitor {
       child = child.getNextSibling();
     }
 
-    final GrStatement[] statements = closure.getStatements();
+    final GrStatement[] statements = owner.getStatements();
     if (statements.length > 0) {
       handlePossibleReturn(statements[statements.length - 1]);
     }
   }
 
+  private void addFunctionalExpressionParameters(GrFunctionalExpression expression) {
+    for (GrParameter parameter : expression.getAllParameters()) {
+      String parameterName = parameter.getName();
+      if (myPolicy.isVariableInitialized(parameter) && parameterName != null) {
+        addNode(new ReadWriteVariableInstruction(parameterName, parameter, ReadWriteVariableInstruction.WRITE));
+      }
+    }
+
+    PsiElement anchor = expression.getArrow();
+    if (expression instanceof GrClosableBlock) anchor = ((GrClosableBlock)expression).getLBrace();
+    addNode(new ReadWriteVariableInstruction("owner", anchor, ReadWriteVariableInstruction.WRITE));
+  }
+
   private <T extends InstructionImpl> T addNode(T instruction) {
     instruction.setNumber(myInstructionNumber++);
     myInstructions.add(instruction);
@@ -211,30 +230,37 @@ public class ControlFlowBuilder extends GroovyRecursiveElementVisitor {
     }
   }
 
+  @Override
+  public void visitLambdaExpression(@NotNull GrLambdaExpression expression) {
+    GrLambdaBody body = expression.getBody();
+    if (body == null) return;
+    ReadWriteVariableInstruction[] reads = ControlFlowBuilderUtil.getReadsWithoutPriorWrites(body.getControlFlow(), false);
+    addReadFromNestedControlFlow(expression, reads);
+  }
+
   @Override
   public void visitClosure(@NotNull GrClosableBlock closure) {
     //do not go inside closures except gstring injections
     if (closure.getParent() instanceof GrStringInjection) {
-      for (GrParameter parameter : closure.getAllParameters()) {
-        if (myPolicy.isVariableInitialized(parameter)) {
-          addNode(new ReadWriteVariableInstruction(parameter.getName(), parameter, ReadWriteVariableInstruction.WRITE));
-        }
-      }
-      addNode(new ReadWriteVariableInstruction("owner", closure.getLBrace(), ReadWriteVariableInstruction.WRITE));
+      addFunctionalExpressionParameters(closure);
 
       super.visitClosure(closure);
       return;
     }
 
     ReadWriteVariableInstruction[] reads = ControlFlowBuilderUtil.getReadsWithoutPriorWrites(closure.getControlFlow(), false);
+    addReadFromNestedControlFlow(closure, reads);
+  }
+
+  private void addReadFromNestedControlFlow(@NotNull PsiElement anchor, @NotNull ReadWriteVariableInstruction[] reads) {
     for (ReadWriteVariableInstruction read : reads) {
       PsiElement element = read.getElement();
       if (!(element instanceof GrReferenceExpression) || myPolicy.isReferenceAccepted((GrReferenceExpression)element)) {
-        addNodeAndCheckPending(new ReadWriteVariableInstruction(read.getVariableName(), closure, ReadWriteVariableInstruction.READ));
+        addNodeAndCheckPending(new ReadWriteVariableInstruction(read.getVariableName(), anchor, ReadWriteVariableInstruction.READ));
       }
     }
 
-    addNodeAndCheckPending(new InstructionImpl(closure));
+    addNodeAndCheckPending(new InstructionImpl(anchor));
   }
 
   @Override
@@ -1186,13 +1212,7 @@ public class ControlFlowBuilder extends GroovyRecursiveElementVisitor {
 
     final Set<ReadWriteVariableInstruction> vars = collectUsedVariableWithoutInitialization(typeDefinition);
 
-    for (ReadWriteVariableInstruction var : vars) {
-      PsiElement element = var.getElement();
-      if (!(element instanceof GrReferenceExpression) || myPolicy.isReferenceAccepted((GrReferenceExpression)element)) {
-        addNodeAndCheckPending(new ReadWriteVariableInstruction(var.getVariableName(), typeDefinition, ReadWriteVariableInstruction.READ));
-      }
-    }
-    addNodeAndCheckPending(new InstructionImpl(typeDefinition));
+    addReadFromNestedControlFlow(typeDefinition, vars.toArray(ReadWriteVariableInstruction.EMPTY_ARRAY));
   }
 
   private static Set<ReadWriteVariableInstruction> collectUsedVariableWithoutInitialization(GrTypeDefinition typeDefinition) {

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/GrLambdaBodyBlockImpl.java

@@ -1,30 +0,0 @@
-// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
-package org.jetbrains.plugins.groovy.lang.psi.impl;
-
-import com.intellij.psi.tree.IElementType;
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.plugins.groovy.lang.psi.GroovyElementVisitor;
-import org.jetbrains.plugins.groovy.lang.psi.api.GrLambdaBody;
-import org.jetbrains.plugins.groovy.lang.psi.impl.statements.blocks.GrBlockImpl;
-
-public class GrLambdaBodyBlockImpl extends GrBlockImpl implements GrLambdaBody {
-
-  public GrLambdaBodyBlockImpl(@NotNull IElementType type, CharSequence buffer) {
-    super(type, buffer);
-  }
-
-  @Override
-  public void accept(@NotNull GroovyElementVisitor visitor) {
-    visitor.visitLambdaBody(this);
-  }
-
-  @Override
-  public String toString() {
-    return "Lambda body block";
-  }
-
-  @Override
-  public boolean isTopControlFlowOwner() {
-    return true;
-  }
-}

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/GrLambdaBodyBlockImpl.kt

@@ -0,0 +1,24 @@
+// Copyright 2000-2019 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
+package org.jetbrains.plugins.groovy.lang.psi.impl
+
+import com.intellij.psi.PsiType
+import com.intellij.psi.tree.IElementType
+import com.intellij.psi.util.parentOfType
+import org.jetbrains.plugins.groovy.lang.psi.GroovyElementVisitor
+import org.jetbrains.plugins.groovy.lang.psi.api.GrLambdaBody
+import org.jetbrains.plugins.groovy.lang.psi.api.GrLambdaExpression
+import org.jetbrains.plugins.groovy.lang.psi.impl.statements.blocks.GrBlockImpl
+import org.jetbrains.plugins.groovy.lang.resolve.MethodTypeInferencer
+
+class GrLambdaBodyBlockImpl(type: IElementType, buffer: CharSequence?) : GrBlockImpl(type, buffer), GrLambdaBody {
+
+  override fun getReturnType(): PsiType? = GroovyPsiManager.inferType(this, MethodTypeInferencer(this))
+
+  override fun isTopControlFlowOwner(): Boolean = true
+
+  override fun getLambdaExpression(): GrLambdaExpression = requireNotNull(parentOfType())
+
+  override fun accept(visitor: GroovyElementVisitor) = visitor.visitLambdaBody(this)
+
+  override fun toString(): String = "Lambda body block"
+}

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/GrLambdaBodyExpressionImpl.kt

@@ -2,15 +2,21 @@
 package org.jetbrains.plugins.groovy.lang.psi.impl
 
 import com.intellij.lang.ASTNode
+import com.intellij.psi.PsiType
+import com.intellij.psi.util.parentOfType
 import com.intellij.util.IncorrectOperationException
 import org.jetbrains.plugins.groovy.lang.psi.GroovyElementVisitor
 import org.jetbrains.plugins.groovy.lang.psi.api.GrLambdaBody
+import org.jetbrains.plugins.groovy.lang.psi.api.GrLambdaExpression
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrStatement
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.expressions.GrExpression
 import org.jetbrains.plugins.groovy.lang.psi.controlFlow.Instruction
 import org.jetbrains.plugins.groovy.lang.psi.controlFlow.impl.ControlFlowBuilder
 
 class GrLambdaBodyExpressionImpl(node: ASTNode) : GroovyPsiElementImpl(node), GrLambdaBody {
+  override fun getReturnType(): PsiType? = getExpression().type
+
+  override fun getLambdaExpression(): GrLambdaExpression = requireNotNull(parentOfType())
 
   fun getExpression() : GrExpression = findNotNullChildByClass(GrExpression::class.java)
 

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/GrLambdaExpressionImpl.kt

@@ -9,20 +9,19 @@ import com.intellij.psi.scope.PsiScopeProcessor
 import com.intellij.psi.util.CachedValueProvider.Result.create
 import com.intellij.psi.util.CachedValuesManager.getCachedValue
 import com.intellij.psi.util.PsiModificationTracker
+import org.jetbrains.plugins.groovy.lang.psi.GroovyElementTypes
 import org.jetbrains.plugins.groovy.lang.psi.GroovyElementVisitor
 import org.jetbrains.plugins.groovy.lang.psi.api.GrLambdaBody
 import org.jetbrains.plugins.groovy.lang.psi.api.GrLambdaExpression
+import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrStatement
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.params.GrParameter
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.params.GrParameterList
-import org.jetbrains.plugins.groovy.lang.psi.dataFlow.types.TypeInferenceHelper
 import org.jetbrains.plugins.groovy.lang.psi.impl.statements.expressions.GrExpressionImpl
 import org.jetbrains.plugins.groovy.lang.psi.impl.statements.params.GrParameterListImpl
 
 class GrLambdaExpressionImpl(node: ASTNode) : GrExpressionImpl(node), GrLambdaExpression {
 
-  override fun getParameters(): Array<GrParameter> {
-    return parameterList.parameters
-  }
+  override fun getParameters(): Array<GrParameter> = parameterList.parameters
 
   override fun getParameterList(): GrParameterList = findNotNullChildByClass(GrParameterListImpl::class.java)
 
@@ -47,11 +46,13 @@ class GrLambdaExpressionImpl(node: ASTNode) : GrExpressionImpl(node), GrLambdaEx
     }
   }
 
-  override fun getReturnType(): PsiType? = TypeInferenceHelper.getCurrentContext().getExpressionType(this, ::calculateReturnType)
+  override fun getArrow(): PsiElement = findNotNullChildByType(GroovyElementTypes.T_ARROW)
 
-  override fun getType(): PsiType? {
-    return GrClosureType.create(this, true)
-  }
+  override fun getStatements(): Array<GrStatement> = body?.statements ?: emptyArray()
+
+  override fun getReturnType(): PsiType? = body?.returnType
+
+  override fun getType(): PsiType? = GrClosureType.create(this, true)
 
   override fun toString(): String = "Lambda expression"
 }

plugins/groovy/groovy-psi/src/org/jetbrains/plugins/groovy/lang/psi/impl/GrReassignedLocalVarsChecker.java

@@ -34,6 +34,7 @@ import org.jetbrains.plugins.groovy.lang.psi.GrControlFlowOwner;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyFile;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyPsiElement;
 import org.jetbrains.plugins.groovy.lang.psi.GroovyRecursiveElementVisitor;
+import org.jetbrains.plugins.groovy.lang.psi.api.GrLambdaBody;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.GrVariable;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.blocks.GrClosableBlock;
 import org.jetbrains.plugins.groovy.lang.psi.api.statements.blocks.GrCodeBlock;
@@ -88,6 +89,13 @@ public class GrReassignedLocalVarsChecker {
             }
           }
 
+          @Override
+          public void visitLambdaBody(@NotNull GrLambdaBody body) {
+            if (getUsedVarsInsideBlock(body).contains(name)) {
+              isReassigned.set(true);
+            }
+          }
+
           @Override
           public void visitElement(@NotNull GroovyPsiElement element) {
             if (isReassigned.get()) return;
@@ -144,11 +152,11 @@ public class GrReassignedLocalVarsChecker {
   }
 
   @NotNull
-  private static Set<String> getUsedVarsInsideBlock(@NotNull final GrCodeBlock block) {
-      return CachedValuesManager.getCachedValue(block, () -> {
+  private static Set<String> getUsedVarsInsideBlock(@NotNull final GroovyPsiElement element) {
+      return CachedValuesManager.getCachedValue(element, () -> {
         final Set<String> result = ContainerUtil.newHashSet();
 
-        block.acceptChildren(new GroovyRecursiveElementVisitor() {
+        element.acceptChildren(new GroovyRecursiveElementVisitor() {
 
           @Override
           public void visitOpenBlock(@NotNull GrOpenBlock openBlock) {
@@ -160,6 +168,11 @@ public class GrReassignedLocalVarsChecker {
             result.addAll(getUsedVarsInsideBlock(closure));
           }
 
+          @Override
+          public void visitLambdaBody(@NotNull GrLambdaBody body) {
+            result.addAll(getUsedVarsInsideBlock(body));
+          }
+
           @Override
           public void visitReferenceExpression(@NotNull GrReferenceExpression referenceExpression) {
             if (referenceExpression.getQualifier() == null && referenceExpression.getReferenceName() != null) {
@@ -167,7 +180,7 @@ public class GrReassignedLocalVarsChecker {
             }
           }
         });
-        return CachedValueProvider.Result.create(result, block);
+        return CachedValueProvider.Result.create(result, element);
       });
   }