From 11da6cee326e078bf523f8ced366fba21336968f Mon Sep 17 00:00:00 2001
From: Tagir Valeev <Tagir.Valeev@jetbrains.com>
Date: Mon, 25 May 2020 12:34:35 +0700
Subject: [PATCH] JetBrains NPE handling: support overridden methods

GitOrigin-RevId: 7ee47707879532d31872cc8befb8eaa76dfd767c
---
 .../DataflowExceptionAnalysisProviderTest.java    | 15 +++++++++++++++
 ...BrainsNotNullInstrumentationExceptionInfo.java | 15 ++++++++++++---
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/java/java-tests/testSrc/com/intellij/java/slicer/DataflowExceptionAnalysisProviderTest.java b/java/java-tests/testSrc/com/intellij/java/slicer/DataflowExceptionAnalysisProviderTest.java
index 152a54863e9c..589038b641ec 100644
--- a/java/java-tests/testSrc/com/intellij/java/slicer/DataflowExceptionAnalysisProviderTest.java
+++ b/java/java-tests/testSrc/com/intellij/java/slicer/DataflowExceptionAnalysisProviderTest.java
@@ -330,6 +330,21 @@ public class DataflowExceptionAnalysisProviderTest extends LightJavaCodeInsightT
            "}");
   }
   
+  public void testNpeJetBrainsOverride() {
+    doTest("Exception in thread \"main\" java.lang.IllegalArgumentException: Argument for @NotNull parameter 's' of MainTest$XImpl.foo must not be null\n" +
+           "\tat MainTest$XImpl.$$$reportNull$$$0(MainTest.java)\n" +
+           "\tat MainTest$XImpl.foo(MainTest.java)",
+           "Find why 's1' could be null",
+           "import org.jetbrains.annotations.NotNull;\n" +
+           "\n" +
+           "public class MainTest {\n" +
+           "    static void test(X x, String s, String s1) { x.foo(s, s1); }\n" +
+           "    interface X { void foo(String s, String t);}\n" +
+           "    static class XImpl implements X { @Override public void foo(String t, @NotNull String s) {}}" +
+           "    public static void main(String[] args) { test(new XImpl(), \"\", null); }" +
+           "}");
+  }
+  
   public void testArrayCopySource() {
     doTest("java.lang.ArrayIndexOutOfBoundsException: arraycopy: source index -1 out of bounds for int[10]\n" +
            "\tat java.base/java.lang.System.arraycopy(Native Method)",
diff --git a/java/openapi/src/com/intellij/execution/filters/JetBrainsNotNullInstrumentationExceptionInfo.java b/java/openapi/src/com/intellij/execution/filters/JetBrainsNotNullInstrumentationExceptionInfo.java
index 92ed4b0c0f76..28363b9c332e 100644
--- a/java/openapi/src/com/intellij/execution/filters/JetBrainsNotNullInstrumentationExceptionInfo.java
+++ b/java/openapi/src/com/intellij/execution/filters/JetBrainsNotNullInstrumentationExceptionInfo.java
@@ -3,6 +3,7 @@ package com.intellij.execution.filters;
 
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.*;
+import com.intellij.psi.util.ClassUtil;
 import com.intellij.util.ObjectUtils;
 import org.jetbrains.annotations.NotNull;
 
@@ -18,6 +19,7 @@ public class JetBrainsNotNullInstrumentationExceptionInfo extends ExceptionInfo
   private final String myParameterName;
   private final String myClassName;
   private final String myMethodName;
+  private final String myFullClassName;
   private final int myWantLines;
 
   private JetBrainsNotNullInstrumentationExceptionInfo(int offset,
@@ -29,6 +31,7 @@ public class JetBrainsNotNullInstrumentationExceptionInfo extends ExceptionInfo
                                                        int wantLines) {
     super(offset, exceptionClassName, exceptionMessage);
     myParameterName = parameterName;
+    myFullClassName = className;
     myClassName = StringUtil.getShortName(className, '/');
     myMethodName = methodName;
     myWantLines = wantLines;
@@ -53,7 +56,13 @@ public class JetBrainsNotNullInstrumentationExceptionInfo extends ExceptionInfo
     if (method == null) return null;
     PsiClass psiClass = method.getContainingClass();
     if (psiClass == null) return null;
-    if (!myClassName.equals(psiClass.getName())) return null;
+    if (!myClassName.equals(psiClass.getName())) {
+      PsiClass aClass = ClassUtil.findPsiClass(method.getManager(), myFullClassName.replace('/', '.'), null, true);
+      if (aClass == null || !aClass.isInheritor(psiClass, true)) return null;
+      PsiMethod subClassMethod = aClass.findMethodBySignature(method, false);
+      if (subClassMethod == null) return null;
+      method = subClassMethod;
+    }
     PsiParameter[] parameters = method.getParameterList().getParameters();
     for (int i = 0; i < parameters.length; i++) {
       if (parameters[i].getName().equals(myParameterName)) {
@@ -73,13 +82,13 @@ public class JetBrainsNotNullInstrumentationExceptionInfo extends ExceptionInfo
       case 2:
         if (line.contains(myClassName+".$$$reportNull$$$0")) {
           return new JetBrainsNotNullInstrumentationExceptionInfo(getClassNameOffset(), getExceptionClassName(), getExceptionMessage(), 
-                                                                  myParameterName, myClassName, myMethodName, 1); 
+                                                                  myParameterName, myFullClassName, myMethodName, 1); 
         }
         break;
       case 1:
         if (line.contains(myClassName+"."+myMethodName)) {
           return new JetBrainsNotNullInstrumentationExceptionInfo(getClassNameOffset(), getExceptionClassName(), getExceptionMessage(),
-                                                                  myParameterName, myClassName, myMethodName, 0);
+                                                                  myParameterName, myFullClassName, myMethodName, 0);
         }
     }
     return null;