From f139ef7215ca77caa343ec4e6afac20189b82fab Mon Sep 17 00:00:00 2001
From: Manu Sridharan <msridhar@gmail.com>
Date: Wed, 24 Apr 2024 14:29:16 -0700
Subject: [PATCH] removeIf

---
 .../handlers/SynchronousCallbackHandler.java  | 15 ++++++++-----
 .../com/uber/nullaway/SyncLambdasTests.java   | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/nullaway/src/main/java/com/uber/nullaway/handlers/SynchronousCallbackHandler.java b/nullaway/src/main/java/com/uber/nullaway/handlers/SynchronousCallbackHandler.java
index e93ff295cc..00a2e544f6 100644
--- a/nullaway/src/main/java/com/uber/nullaway/handlers/SynchronousCallbackHandler.java
+++ b/nullaway/src/main/java/com/uber/nullaway/handlers/SynchronousCallbackHandler.java
@@ -12,7 +12,7 @@
 import com.sun.source.tree.Tree;
 import com.sun.source.util.TreePath;
 import com.sun.tools.javac.code.Symbol;
-import com.uber.nullaway.LibraryModels;
+import com.uber.nullaway.LibraryModels.MethodRef;
 import com.uber.nullaway.dataflow.AccessPath;
 import java.util.function.Predicate;
 
@@ -20,14 +20,19 @@ public class SynchronousCallbackHandler extends BaseNoOpHandler {
 
   // TODO this should work on subtypes of the methods as well, like java.util.HashMap.  Use a
   // Matcher?
-  private static final ImmutableMap<String, ImmutableMap<LibraryModels.MethodRef, Integer>>
+  private static final ImmutableMap<String, ImmutableMap<MethodRef, Integer>>
       METHOD_NAME_TO_SIG_AND_PARAM_INDEX =
           ImmutableMap.of(
               "forEach",
               ImmutableMap.of(
-                  LibraryModels.MethodRef.methodRef(
+                  MethodRef.methodRef(
                       "java.util.Map",
                       "forEach(java.util.function.BiConsumer<? super K,? super V>)"),
+                  0),
+              "removeIf",
+              ImmutableMap.of(
+                  MethodRef.methodRef(
+                      "java.util.Collection", "removeIf(java.util.function.Predicate<? super E>)"),
                   0));
 
   @Override
@@ -46,9 +51,9 @@ public Predicate<AccessPath> getAccessPathPredForSavedContext(TreePath path, Vis
       }
       String invokedMethodName = symbol.getSimpleName().toString();
       if (METHOD_NAME_TO_SIG_AND_PARAM_INDEX.containsKey(invokedMethodName)) {
-        ImmutableMap<LibraryModels.MethodRef, Integer> entriesForMethodName =
+        ImmutableMap<MethodRef, Integer> entriesForMethodName =
             METHOD_NAME_TO_SIG_AND_PARAM_INDEX.get(invokedMethodName);
-        for (LibraryModels.MethodRef methodRef : entriesForMethodName.keySet()) {
+        for (MethodRef methodRef : entriesForMethodName.keySet()) {
           if (symbol.toString().equals(methodRef.fullMethodSig)
               && ASTHelpers.isSubtype(
                   symbol.owner.type, state.getTypeFromString(methodRef.enclosingClass), state)) {
diff --git a/nullaway/src/test/java/com/uber/nullaway/SyncLambdasTests.java b/nullaway/src/test/java/com/uber/nullaway/SyncLambdasTests.java
index 83e91bcff3..0a63ed604b 100644
--- a/nullaway/src/test/java/com/uber/nullaway/SyncLambdasTests.java
+++ b/nullaway/src/test/java/com/uber/nullaway/SyncLambdasTests.java
@@ -86,4 +86,26 @@ public void otherForEach() {
             "}")
         .doTest();
   }
+
+  @Test
+  public void removeIf() {
+    defaultCompilationHelper
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import java.util.List;",
+            "import java.util.ArrayList;",
+            "import org.jspecify.annotations.Nullable;",
+            "public class Test {",
+            "    private @Nullable Object f;",
+            "    public void test1() {",
+            "        if (this.f == null) {",
+            "            throw new IllegalArgumentException();",
+            "        }",
+            "        List<Object> l = new ArrayList<>();",
+            "        l.removeIf(v -> this.f.toString().equals(v.toString()));",
+            "    }",
+            "}")
+        .doTest();
+  }
 }