Java slicing: support -> this/-> paramN contracts

Fixes IDEA-121571 Feature request: A Contract Annotation that indicates that a method returns its argument

java/java-analysis-impl/src/com/intellij/codeInspection/dataFlow/JavaMethodContractUtil.java

@@ -130,7 +130,9 @@ public class JavaMethodContractUtil {
    * @return the expression which is always returned by this method if it completes successfully,
    * null if method may return something less trivial or its contract is unknown.
    */
-  public static PsiExpression findReturnedValue(PsiMethodCallExpression call) {
+  @Nullable
+  @Contract("null -> null")
+  public static PsiExpression findReturnedValue(@Nullable PsiMethodCallExpression call) {
     if (call == null) return null;
     PsiMethod method = call.resolveMethod();
     if (method == null) return null;

java/java-impl/src/com/intellij/slicer/SliceForwardUtil.java

@@ -15,6 +15,7 @@
  */
 package com.intellij.slicer;
 
+import com.intellij.codeInspection.dataFlow.JavaMethodContractUtil;
 import com.intellij.lang.java.JavaLanguage;
 import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.util.Pair;
@@ -25,7 +26,9 @@ import com.intellij.psi.search.searches.ReferencesSearch;
 import com.intellij.psi.util.MethodSignatureUtil;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ArrayUtilRt;
+import com.intellij.util.ObjectUtils;
 import com.intellij.util.Processor;
+import com.siyeh.ig.psiutils.ExpressionUtils;
 import gnu.trove.THashSet;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -42,6 +45,13 @@ class SliceForwardUtil {
   static boolean processUsagesFlownFromThe(@NotNull PsiElement element,
                                            @NotNull final JavaSliceUsage parent,
                                            @NotNull final Processor<SliceUsage> processor) {
+    PsiExpression expression = getMethodCallTarget(element);
+    if (expression != null) {
+      SliceUsage usage = SliceUtil.createSliceUsage(expression, parent, parent.getSubstitutor(), parent.indexNesting, "");
+      if (!processor.process(usage)) {
+        return false;
+      }
+    }
     Pair<PsiElement, PsiSubstitutor> pair = getAssignmentTarget(element, parent);
     if (pair != null) {
       PsiElement target = pair.getFirst();
@@ -196,7 +206,7 @@ class SliceForwardUtil {
     }
     Pair<PsiElement, PsiSubstitutor> pair = getAssignmentTarget(element, parent);
     if (pair != null) {
-      SliceUsage usage = SliceUtil.createSliceUsage(element, parent, pair.getSecond(),parent.indexNesting, "");
+      SliceUsage usage = SliceUtil.createSliceUsage(element, parent, pair.getSecond(), parent.indexNesting, "");
       return processor.process(usage);
     }
     if (parent.params.showInstanceDereferences && isDereferenced(element)) {
@@ -206,6 +216,16 @@ class SliceForwardUtil {
     return true;
   }
 
+  private static PsiExpression getMethodCallTarget(PsiElement element) {
+    element = complexify(element);
+    PsiMethodCallExpression call = null;
+    if (element.getParent() instanceof PsiExpressionList) {
+      call = ObjectUtils.tryCast(element.getParent().getParent(), PsiMethodCallExpression.class);
+    }
+    PsiExpression value = JavaMethodContractUtil.findReturnedValue(call);
+    return value == element ? call : null;
+  }
+
   private static boolean isDereferenced(@NotNull PsiElement element) {
     if (!(element instanceof PsiReferenceExpression)) return false;
     PsiElement parent = element.getParent();
@@ -259,6 +279,13 @@ class SliceForwardUtil {
         target = PsiTreeUtil.getParentOfType(statement, PsiMethod.class);
       }
     }
+    else if (element instanceof PsiExpression){
+      PsiMethodCallExpression call = ExpressionUtils.getCallForQualifier((PsiExpression)element);
+      PsiExpression maybeQualifier = JavaMethodContractUtil.findReturnedValue(call);
+      if (maybeQualifier == element) {
+        target = call;
+      }
+    }
 
     return target == null ? null : Pair.create(target, substitutor);
   }

java/java-impl/src/com/intellij/slicer/SliceUtil.java

@@ -18,6 +18,7 @@ package com.intellij.slicer;
 import com.intellij.codeInsight.AnnotationUtil;
 import com.intellij.codeInspection.dataFlow.DfaPsiUtil;
 import com.intellij.codeInspection.dataFlow.DfaUtil;
+import com.intellij.codeInspection.dataFlow.JavaMethodContractUtil;
 import com.intellij.lang.Language;
 import com.intellij.lang.java.JavaLanguage;
 import com.intellij.openapi.progress.ProgressManager;
@@ -149,6 +150,12 @@ class SliceUtil {
       }
     }
     if (expression instanceof PsiMethodCallExpression) { // ctr call can't return value or be container get, so don't use PsiCall here
+      PsiExpression returnedValue = JavaMethodContractUtil.findReturnedValue((PsiMethodCallExpression)expression);
+      if (returnedValue != null) {
+        if (!handToProcessor(returnedValue, processor, parent, parentSubstitutor, indexNesting, syntheticField)) {
+          return false;
+        }
+      }
       PsiMethod method = ((PsiMethodCallExpression)expression).resolveMethod();
       Flow anno = method == null ? null : isMethodFlowAnnotated(method);
       if (anno != null) {
@@ -219,13 +226,7 @@ class SliceUtil {
   }
 
   private static PsiElement simplify(@NotNull PsiElement expression) {
-    if (expression instanceof PsiParenthesizedExpression) {
-      return simplify(((PsiParenthesizedExpression)expression).getExpression());
-    }
-    if (expression instanceof PsiTypeCastExpression) {
-      return simplify(((PsiTypeCastExpression)expression).getOperand());
-    }
-    return expression;
+    return expression instanceof PsiExpression ? PsiUtil.deparenthesizeExpression((PsiExpression)expression) : expression;
   }
 
   private static boolean handToProcessor(@NotNull PsiElement expression,

java/java-tests/testData/codeInsight/slice/backward/Append.java

@@ -0,0 +1,6 @@
+class Test {
+  StringBuilder builder(StringBuilder <flown1111>sb) {
+    StringBuilder foo = <flown111><flown11><flown1>sb.append("foo").append(123);
+    return <caret>foo;
+  }
+}

java/java-tests/testData/codeInsight/slice/backward/InterMethod.java

@@ -10,11 +10,11 @@ class WW {
     }
 
     {
-        x(<flown111111>"zzz");
+        x(<flown1211><flown111111>"zzz");
     }
 
-    String x(String <flown11111>g) {
-        String d = <flown1>foo(<flown1111>g);
+    String x(String <flown121><flown11111>g) {
+        String d = <flown1>foo(<flown12><flown1111>g);
         return <caret>d;
     }
 

java/java-tests/testData/codeInsight/slice/backward/RequireNonNull.java

@@ -0,0 +1,21 @@
+class Test {
+
+  private String <flown1>s;
+  private String s2;
+
+  void test(String <flown112111><flown1111>s, String s2) {
+    this.s = <flown11>requireNonNull(<flown11211><flown111>s);
+    this.s2 = requireNonNull(s2);
+  }
+
+  String foo() {
+    return <caret>s;
+  }
+
+  public static <T> T requireNonNull(T <flown1121>obj) {
+    if (obj == null)
+      throw new NullPointerException();
+    return <flown112>obj;
+  }
+
+}

java/java-tests/testData/codeInsight/slice/forward/Append.java

@@ -0,0 +1,6 @@
+class Test {
+  StringBuilder <flown111111>builder(StringBuilder <caret>sb) {
+    StringBuilder <flown1111>foo = <flown111><flown11><flown1>sb.append("foo").append(1);
+    return <flown11111>foo;
+  }
+}

java/java-tests/testData/codeInsight/slice/forward/RequireNonNull.java

@@ -0,0 +1,16 @@
+class Test {
+  private String <flown111>s;
+  private String s2;
+
+  void test(String <caret>s, String s2) {
+    this.s = <flown11>requireNonNull(<flown1>s);
+    this.s2 = requireNonNull(s2);
+  }
+
+  String <flown11111>foo() {
+    return <flown1111>s;
+  }
+
+  @org.jetbrains.annotations.Contract("null -> fail; !null -> param1")
+  public static native <T> T requireNonNull(T <flown12>obj);
+}

java/java-tests/testSrc/com/intellij/java/slicer/SliceBackwardTest.java

@@ -21,8 +21,6 @@ import com.intellij.openapi.editor.RangeMarker;
 import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.slicer.*;
-import com.intellij.util.containers.IntArrayList;
-import gnu.trove.TIntObjectHashMap;
 
 import java.util.Collection;
 import java.util.Map;
@@ -31,15 +29,13 @@ import java.util.Map;
  * @author cdr
  */
 public class SliceBackwardTest extends SliceTestCase {
-  private final TIntObjectHashMap<IntArrayList> myFlownOffsets = new TIntObjectHashMap<>();
-
   private void doTest() throws Exception {
     configureByFile("/codeInsight/slice/backward/"+getTestName(false)+".java");
     Map<String, RangeMarker> sliceUsageName2Offset = SliceTestUtil.extractSliceOffsetsFromDocument(getEditor().getDocument());
     PsiDocumentManager.getInstance(getProject()).commitAllDocuments();
     PsiElement element = new SliceHandler(true).getExpressionAtCaret(getEditor(), getFile());
     assertNotNull(element);
-    SliceTestUtil.calcRealOffsets(element, sliceUsageName2Offset, myFlownOffsets);
+    SliceTestUtil.Node tree = SliceTestUtil.buildTree(element, sliceUsageName2Offset);
     Collection<HighlightInfo> errors = highlightErrors();
     assertEmpty(errors);
     SliceAnalysisParams params = new SliceAnalysisParams();
@@ -47,7 +43,7 @@ public class SliceBackwardTest extends SliceTestCase {
     params.dataFlowToThis = true;
 
     SliceUsage usage = LanguageSlicing.getProvider(element).createRootUsage(element, params);
-    SliceTestUtil.checkUsages(usage, myFlownOffsets);
+    SliceTestUtil.checkUsages(usage, tree);
   }
 
   public void testSimple() throws Exception { doTest();}
@@ -85,4 +81,6 @@ public class SliceBackwardTest extends SliceTestCase {
   public void testFinalVarAssignedBeforePassingToAnonymous() throws Exception { doTest();}
   public void testLocalVarDeclarationAndAssignment() throws Exception { doTest();}
   public void testSearchOverriddenMethodsInThisClassHierarchy() throws Exception { doTest();}
+  public void testAppend() throws Exception { doTest();}
+  public void testRequireNonNull() throws Exception { doTest();}
 }

java/java-tests/testSrc/com/intellij/java/slicer/SliceForwardTest.java

@@ -21,8 +21,6 @@ import com.intellij.openapi.editor.RangeMarker;
 import com.intellij.psi.PsiDocumentManager;
 import com.intellij.psi.PsiElement;
 import com.intellij.slicer.*;
-import com.intellij.util.containers.IntArrayList;
-import gnu.trove.TIntObjectHashMap;
 
 import java.util.Collection;
 import java.util.Map;
@@ -31,25 +29,25 @@ import java.util.Map;
  * @author cdr
  */
 public class SliceForwardTest extends SliceTestCase {
-  private final TIntObjectHashMap<IntArrayList> myFlownOffsets = new TIntObjectHashMap<>();
-
   private void dotest() throws Exception {
     configureByFile("/codeInsight/slice/forward/"+getTestName(false)+".java");
     Map<String, RangeMarker> sliceUsageName2Offset = SliceTestUtil.extractSliceOffsetsFromDocument(getEditor().getDocument());
     PsiDocumentManager.getInstance(getProject()).commitAllDocuments();
     PsiElement element = new SliceForwardHandler().getExpressionAtCaret(getEditor(), getFile());
     assertNotNull(element);
-    SliceTestUtil.calcRealOffsets(element, sliceUsageName2Offset, myFlownOffsets);
+    SliceTestUtil.Node tree = SliceTestUtil.buildTree(element, sliceUsageName2Offset);
     Collection<HighlightInfo> errors = highlightErrors();
     assertEmpty(errors);
     SliceAnalysisParams params = new SliceAnalysisParams();
     params.scope = new AnalysisScope(getProject());
     params.dataFlowToThis = false;
     SliceUsage usage = LanguageSlicing.getProvider(element).createRootUsage(element, params);
-    SliceTestUtil.checkUsages(usage, myFlownOffsets);
+    SliceTestUtil.checkUsages(usage, tree);
   }
 
   public void testSimple() throws Exception { dotest();}
   public void testInterMethod() throws Exception { dotest();}
   public void testParameters() throws Exception { dotest();}
+  public void testRequireNonNull() throws Exception { dotest();}
+  public void testAppend() throws Exception { dotest();}
 }

platform/lang-impl/testSources/com/intellij/slicer/SliceTestUtil.java

@@ -26,9 +26,7 @@ import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.util.CommonProcessors;
-import com.intellij.util.containers.IntArrayList;
 import gnu.trove.THashMap;
-import gnu.trove.TIntObjectHashMap;
 
 import java.util.*;
 
@@ -39,9 +37,35 @@ public class SliceTestUtil {
   private SliceTestUtil() {
   }
 
-  public static void calcRealOffsets(PsiElement startElement, Map<String, RangeMarker> sliceUsageName2Offset,
-                              final TIntObjectHashMap<IntArrayList> flownOffsets) {
-    fill(sliceUsageName2Offset, "", startElement.getTextOffset(), flownOffsets);
+  public static class Node {
+    public final int myOffset;
+    public final List<Node> myChildren;
+
+    public Node(int offset, List<Node> children) {
+      myOffset = offset;
+      myChildren = children;
+      myChildren.sort(Comparator.comparingInt(n -> n.myOffset));
+    }
+
+    @Override
+    public String toString() {
+      return myOffset + (myChildren.isEmpty() ? "" : " -> " + myChildren);
+    }
+  }
+
+  public static Node buildTree(PsiElement startElement, Map<String, RangeMarker> sliceUsageName2Offset) {
+    return buildNode("", startElement.getTextOffset(), sliceUsageName2Offset);
+  }
+
+  private static Node buildNode(String name, int offset, Map<String, RangeMarker> sliceUsageName2Offset) {
+    List<Node> children = new ArrayList<>();
+    for (int i = 1; i < 9; i++) {
+      String newName = name + i;
+      RangeMarker marker = sliceUsageName2Offset.get(newName);
+      if (marker == null) break;
+      children.add(buildNode(newName, marker.getStartOffset(), sliceUsageName2Offset));
+    }
+    return new Node(offset, children);
   }
 
   public static Map<String, RangeMarker> extractSliceOffsetsFromDocument(final Document document) {
@@ -64,23 +88,6 @@ public class SliceTestUtil {
     return sliceUsageName2Offset;
   }
 
-  private static void fill(Map<String, RangeMarker> sliceUsageName2Offset, String name, int offset,
-                           final TIntObjectHashMap<IntArrayList> flownOffsets) {
-    for (int i=1;i<9;i++) {
-      String newName = name + i;
-      RangeMarker marker = sliceUsageName2Offset.get(newName);
-      if (marker == null) break;
-      IntArrayList offsets = flownOffsets.get(offset);
-      if (offsets == null) {
-        offsets = new IntArrayList();
-        flownOffsets.put(offset, offsets);
-      }
-      int newStartOffset = marker.getStartOffset();
-      offsets.add(newStartOffset);
-      fill(sliceUsageName2Offset, newName, newStartOffset, flownOffsets);
-    }
-  }
-
   private static void extract(final List<Document> documents, final Map<String, RangeMarker> sliceUsageName2Offset, final String name) {
     WriteCommandAction.runWriteCommandAction(null, () -> {
       for (int i = 1; i < 9; i++) {
@@ -107,26 +114,21 @@ public class SliceTestUtil {
     });
   }
 
-  public static void checkUsages(final SliceUsage usage, final TIntObjectHashMap<IntArrayList> flownOffsets) {
+  public static void checkUsages(final SliceUsage usage, final Node tree) {
     final List<SliceUsage> children = new ArrayList<>();
     boolean b = ProgressManager.getInstance().runProcessWithProgressSynchronously(
       () -> usage.processChildren(new CommonProcessors.CollectProcessor<>(children)), "Expanding", true, usage.getElement().getProject());
     assertTrue(b);
     int startOffset = usage.getElement().getTextOffset();
-    IntArrayList list = flownOffsets.get(startOffset);
-    int[] offsets = list == null ? new int[0] : list.toArray();
-    Arrays.sort(offsets);
+    assertEquals(message(startOffset, usage), tree.myOffset, startOffset);
+    List<Node> expectedChildren = tree.myChildren;
 
-    int size = offsets.length;
+    int size = expectedChildren.size();
     assertEquals(message(startOffset, usage), size, children.size());
     children.sort(Comparator.naturalOrder());
 
     for (int i = 0; i < children.size(); i++) {
-      SliceUsage child = children.get(i);
-      int offset = offsets[i];
-      assertEquals(message(offset, child), offset, child.getUsageInfo().getElement().getTextOffset());
-
-      checkUsages(child, flownOffsets);
+      checkUsages(children.get(i), expectedChildren.get(i));
     }
   }
 

plugins/groovy/test/org/jetbrains/plugins/groovy/lang/slicer/GroovySliceTestCase.kt

@@ -19,10 +19,11 @@ import com.intellij.analysis.AnalysisScope
 import com.intellij.codeInsight.daemon.DaemonAnalyzerTestCase
 import com.intellij.psi.PsiDocumentManager
 import com.intellij.psi.PsiManager
-import com.intellij.slicer.*
+import com.intellij.slicer.LanguageSlicing
+import com.intellij.slicer.SliceAnalysisParams
+import com.intellij.slicer.SliceHandler
+import com.intellij.slicer.SliceTestUtil
 import com.intellij.testFramework.UsefulTestCase
-import com.intellij.util.containers.IntArrayList
-import gnu.trove.TIntObjectHashMap
 import org.jetbrains.plugins.groovy.util.TestUtils
 
 abstract class GroovySliceTestCase(private val isDataFlowToThis: Boolean) : DaemonAnalyzerTestCase() {
@@ -51,8 +52,7 @@ abstract class GroovySliceTestCase(private val isDataFlowToThis: Boolean) : Daem
     psiDocumentManager.commitAllDocuments()
 
     val element = SliceHandler(isDataFlowToThis).getExpressionAtCaret(editor, file)!!
-    val flownOffsets = TIntObjectHashMap<IntArrayList>()
-    SliceTestUtil.calcRealOffsets(element, sliceUsageName2Offset, flownOffsets)
+    val tree = SliceTestUtil.buildTree(element, sliceUsageName2Offset)
 
     val errors = highlightErrors()
     UsefulTestCase.assertEmpty(errors)
@@ -62,6 +62,6 @@ abstract class GroovySliceTestCase(private val isDataFlowToThis: Boolean) : Daem
       dataFlowToThis = isDataFlowToThis
     }
     val usage = LanguageSlicing.getProvider(element)!!.createRootUsage(element, params)
-    SliceTestUtil.checkUsages(usage, flownOffsets)
+    SliceTestUtil.checkUsages(usage, tree)
   }
 }