Avoid redundant Map lookups (#722)

while reading through the code i noticed a few redundant map lookups

i've only adjusted the places where it was clear the map never stores null values and where `get` immediately followed a `containsKey`

according to running the JMH benchmarks locally this safes a few cycles and often makes the code clearer imo

build.gradle

@@ -39,7 +39,7 @@ plugins {
   id "net.ltgt.errorprone" version "3.0.1" apply false
   id "com.github.johnrengelman.shadow" version "6.1.0" apply false
   id "com.github.kt3k.coveralls" version "2.12.0" apply false
-  id "me.champeau.jmh" version "0.6.7" apply false
+  id "me.champeau.jmh" version "0.6.8" apply false
   id "com.github.ben-manes.versions" version "0.42.0"
 }
 

nullaway/src/main/java/com/uber/nullaway/CodeAnnotationInfo.java

@@ -292,7 +292,7 @@ public ClassCacheRecord(Symbol.ClassSymbol outermostClassSymbol, boolean isAnnot
     }
 
     public boolean isMethodNullnessAnnotated(Symbol.MethodSymbol methodSymbol) {
-      methodNullnessCache.computeIfAbsent(
+      return methodNullnessCache.computeIfAbsent(
           methodSymbol,
           m -> {
             if (ASTHelpers.hasDirectAnnotationWithSimpleName(
@@ -305,7 +305,6 @@ public boolean isMethodNullnessAnnotated(Symbol.MethodSymbol methodSymbol) {
                   m, NullabilityUtil.NULLMARKED_SIMPLE_NAME);
             }
           });
-      return methodNullnessCache.get(methodSymbol);
     }
   }
 }

nullaway/src/main/java/com/uber/nullaway/NullAway.java

@@ -83,6 +83,7 @@
 import com.sun.source.util.TreePath;
 import com.sun.source.util.Trees;
 import com.sun.tools.javac.code.Symbol;
+import com.sun.tools.javac.code.Symbol.ClassSymbol;
 import com.sun.tools.javac.code.Symbol.VarSymbol;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.processing.JavacProcessingEnvironment;
@@ -1208,12 +1209,10 @@ private boolean fieldInitializedByPreviousInitializer(
       Symbol fieldSymbol, TreePath initTreePath, VisitorState state) {
     TreePath enclosingClassPath = initTreePath.getParentPath();
     ClassTree enclosingClass = (ClassTree) enclosingClassPath.getLeaf();
+    ClassSymbol classSymbol = ASTHelpers.getSymbol(enclosingClass);
     Multimap<Tree, Element> tree2Init =
-        initTree2PrevFieldInit.get(ASTHelpers.getSymbol(enclosingClass));
-    if (tree2Init == null) {
-      tree2Init = computeTree2Init(enclosingClassPath, state);
-      initTree2PrevFieldInit.put(ASTHelpers.getSymbol(enclosingClass), tree2Init);
-    }
+        initTree2PrevFieldInit.computeIfAbsent(
+            classSymbol, sym -> computeTree2Init(enclosingClassPath, state));
     return tree2Init.containsEntry(initTreePath.getLeaf(), fieldSymbol);
   }
 
@@ -2412,11 +2411,9 @@ private static ExpressionTree stripParensAndCasts(ExpressionTree expr) {
    * @return computed nullness for e, if any, else Nullable
    */
   public Nullness getComputedNullness(ExpressionTree e) {
-    if (computedNullnessMap.containsKey(e)) {
-      return computedNullnessMap.get(e);
-    } else {
-      return Nullness.NULLABLE;
-    }
+    // TODO: use Map.getOrDefault after https://github.com/uber/NullAway/issues/723
+    Nullness nullness = computedNullnessMap.get(e);
+    return nullness != null ? nullness : Nullness.NULLABLE;
   }
 
   /**

nullaway/src/main/java/com/uber/nullaway/fixserialization/SerializationService.java

@@ -36,6 +36,7 @@
 import com.uber.nullaway.fixserialization.location.SymbolLocation;
 import com.uber.nullaway.fixserialization.out.ErrorInfo;
 import com.uber.nullaway.fixserialization.out.SuggestedNullableFixInfo;
+import java.util.Map;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import javax.annotation.Nullable;
@@ -68,8 +69,10 @@ public static String escapeSpecialCharacters(String str) {
     // escape existing backslashes
     str = str.replaceAll(Pattern.quote("\\"), Matcher.quoteReplacement("\\\\"));
     // escape special characters
-    for (Character key : escapes.keySet()) {
-      str = str.replaceAll(String.valueOf(key), Matcher.quoteReplacement("\\" + escapes.get(key)));
+    for (Map.Entry<Character, Character> entry : escapes.entrySet()) {
+      str =
+          str.replaceAll(
+              String.valueOf(entry.getKey()), Matcher.quoteReplacement("\\" + entry.getValue()));
     }
     return str;
   }

nullaway/src/main/java/com/uber/nullaway/handlers/InferredJARModelsHandler.java

@@ -342,17 +342,12 @@ private void parseStubStream(InputStream stubxInputStream, String stubxLocation)
   private void cacheAnnotation(String methodSig, Integer argNum, String annotation) {
     // TODO: handle inner classes properly
     String className = methodSig.split(":")[0].replace('$', '.');
-    if (!argAnnotCache.containsKey(className)) {
-      argAnnotCache.put(className, new LinkedHashMap<>());
-    }
-    Map<String, Map<Integer, Set<String>>> cacheForClass = argAnnotCache.get(className);
-    if (!cacheForClass.containsKey(methodSig)) {
-      cacheForClass.put(methodSig, new LinkedHashMap<>());
-    }
-    Map<Integer, Set<String>> cacheForMethod = cacheForClass.get(methodSig);
-    if (!cacheForMethod.containsKey(argNum)) {
-      cacheForMethod.put(argNum, new LinkedHashSet<>());
-    }
-    cacheForMethod.get(argNum).add(annotation);
+    Map<String, Map<Integer, Set<String>>> cacheForClass =
+        argAnnotCache.computeIfAbsent(className, s -> new LinkedHashMap<>());
+    Map<Integer, Set<String>> cacheForMethod =
+        cacheForClass.computeIfAbsent(methodSig, s -> new LinkedHashMap<>());
+    Set<String> cacheForArgument =
+        cacheForMethod.computeIfAbsent(argNum, s -> new LinkedHashSet<>());
+    cacheForArgument.add(annotation);
   }
 }

nullaway/src/main/java/com/uber/nullaway/handlers/LibraryModelsHandler.java

@@ -984,11 +984,8 @@ private <T> NameIndexedMap<T> makeOptimizedLookup(
       Map<Name, Map<MethodRef, T>> nameMapping = new LinkedHashMap<>();
       for (MethodRef ref : refs) {
         Name methodName = names.fromString(ref.methodName);
-        Map<MethodRef, T> mapForName = nameMapping.get(methodName);
-        if (mapForName == null) {
-          mapForName = new LinkedHashMap<>();
-          nameMapping.put(methodName, mapForName);
-        }
+        Map<MethodRef, T> mapForName =
+            nameMapping.computeIfAbsent(methodName, k -> new LinkedHashMap<>());
         mapForName.put(ref, getValForRef.apply(ref));
       }
       return new NameIndexedMap<>(nameMapping);

nullaway/src/main/java/com/uber/nullaway/handlers/StreamNullabilityPropagator.java

@@ -325,10 +325,10 @@ public void onMatchMethodReference(
       MemberReferenceTree tree,
       VisitorState state,
       Symbol.MethodSymbol methodSymbol) {
-    if (mapToFilterMap.containsKey(tree) && ((JCTree.JCMemberReference) tree).kind.isUnbound()) {
+    MaplikeToFilterInstanceRecord callInstanceRecord = mapToFilterMap.get(tree);
+    if (callInstanceRecord != null && ((JCTree.JCMemberReference) tree).kind.isUnbound()) {
       // Unbound method reference, check if we know the corresponding path to be NonNull from the
       // previous filter.
-      MaplikeToFilterInstanceRecord callInstanceRecord = mapToFilterMap.get(tree);
       Tree filterTree = callInstanceRecord.getFilter();
       if (!(filterTree instanceof MethodTree || filterTree instanceof LambdaExpressionTree)) {
         throw new IllegalStateException(
@@ -421,9 +421,9 @@ public NullnessStore.Builder onDataflowInitialStore(
       return nullnessBuilder;
     }
     assert (tree instanceof MethodTree || tree instanceof LambdaExpressionTree);
-    if (mapToFilterMap.containsKey(tree)) {
+    MaplikeToFilterInstanceRecord callInstanceRecord = mapToFilterMap.get(tree);
+    if (callInstanceRecord != null) {
       // Plug Nullness info from filter method into entry to map method.
-      MaplikeToFilterInstanceRecord callInstanceRecord = mapToFilterMap.get(tree);
       Tree filterTree = callInstanceRecord.getFilter();
       assert (filterTree instanceof MethodTree || filterTree instanceof LambdaExpressionTree);
       MaplikeMethodRecord mapMR = callInstanceRecord.getMaplikeMethodRecord();
@@ -459,25 +459,23 @@ public NullnessStore.Builder onDataflowInitialStore(
   @Override
   public void onDataflowVisitReturn(
       ReturnTree tree, NullnessStore thenStore, NullnessStore elseStore) {
-    if (returnToEnclosingMethodOrLambda.containsKey(tree)) {
-      Tree filterTree = returnToEnclosingMethodOrLambda.get(tree);
+    Tree filterTree = returnToEnclosingMethodOrLambda.get(tree);
+    if (filterTree != null) {
       assert (filterTree instanceof MethodTree || filterTree instanceof LambdaExpressionTree);
       ExpressionTree retExpression = tree.getExpression();
       if (canBooleanExpressionEvalToTrue(retExpression)) {
-        if (filterToNSMap.containsKey(filterTree)) {
-          filterToNSMap.put(filterTree, filterToNSMap.get(filterTree).leastUpperBound(thenStore));
-        } else {
-          filterToNSMap.put(filterTree, thenStore);
-        }
+        filterToNSMap.compute(
+            filterTree,
+            (key, value) -> value == null ? thenStore : value.leastUpperBound(thenStore));
       }
     }
   }
 
   @Override
   public void onDataflowVisitLambdaResultExpression(
       ExpressionTree tree, NullnessStore thenStore, NullnessStore elseStore) {
-    if (expressionBodyToFilterLambda.containsKey(tree)) {
-      LambdaExpressionTree filterTree = expressionBodyToFilterLambda.get(tree);
+    LambdaExpressionTree filterTree = expressionBodyToFilterLambda.get(tree);
+    if (filterTree != null) {
       if (canBooleanExpressionEvalToTrue(tree)) {
         filterToNSMap.put(filterTree, thenStore);
       }