Optimize library model lookups

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

@@ -22,15 +22,12 @@
 
 package com.uber.nullaway;
 
-import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSetMultimap;
-import com.google.errorprone.util.ASTHelpers;
 import com.sun.tools.javac.code.Symbol;
-import com.sun.tools.javac.code.Types;
 import java.util.Objects;
-import javax.annotation.Nullable;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 /** Provides models for library routines for the null checker. */
 public interface LibraryModels {
@@ -94,92 +91,77 @@ public interface LibraryModels {
    */
   final class MethodRef {
 
-    public final String clazz;
-    public final String methodSig;
-
-    private MethodRef(String clazz, String methodSig) {
-      this.clazz = clazz;
-      this.methodSig = methodSig;
-    }
-
+    public final String enclosingClass;
     /**
-     * @param clazz containing class
-     * @param methodSig method signature in the appropriate format (see class docs)
-     * @return corresponding {@link MethodRef}
+     * we store the method name separately to enable fast comparison with MethodSymbols. See {@link
+     * com.uber.nullaway.handlers.LibraryModelsHandler.OptimizedLibraryModels}
      */
-    public static MethodRef methodRef(Class<?> clazz, String methodSig) {
-      return methodRef(clazz.getName(), methodSig);
+    public final String methodName;
+
+    public final String fullMethodSig;
+
+    private MethodRef(String enclosingClass, String methodName, String fullMethodSig) {
+      this.enclosingClass = enclosingClass;
+      this.methodName = methodName;
+      this.fullMethodSig = fullMethodSig;
     }
 
+    private static final Pattern METHOD_SIG_PATTERN = Pattern.compile("^(<.*>)?(\\w+)(\\(.*\\))$");
+
     /**
-     * @param clazz containing class
-     * @param methodSig method signature in the appropriate format (see class docs)
+     * @param enclosingClass containing class
+     * @param methodSignature method signature in the appropriate format (see class docs)
      * @return corresponding {@link MethodRef}
      */
-    public static MethodRef methodRef(String clazz, String methodSig) {
-      Preconditions.checkArgument(
-          isValidMethodSig(methodSig), methodSig + " is not a valid method signature");
-      return new MethodRef(clazz, methodSig);
-    }
-
-    private static boolean isValidMethodSig(String methodSig) {
-      // some basic checking to make sure it's not just a method name
-      return methodSig.contains("(") && methodSig.contains(")");
+    public static MethodRef methodRef(String enclosingClass, String methodSignature) {
+      Matcher matcher = METHOD_SIG_PATTERN.matcher(methodSignature);
+      if (matcher.find()) {
+        String methodName = matcher.group(2);
+        if (methodName.equals(enclosingClass.substring(enclosingClass.lastIndexOf('.') + 1))) {
+          // constructor
+          methodName = "<init>";
+        }
+        return new MethodRef(enclosingClass, methodName, methodSignature);
+      } else {
+        throw new IllegalArgumentException("malformed method signature " + methodSignature);
+      }
     }
 
     public static MethodRef fromSymbol(Symbol.MethodSymbol symbol) {
-      return methodRef(symbol.owner.getQualifiedName().toString(), symbol.toString());
+      String methodStr = symbol.toString();
+
+      return new MethodRef(
+          symbol.owner.getQualifiedName().toString(), symbol.name.toString(), methodStr);
     }
 
     @Override
-    public boolean equals(Object obj) {
-      if (obj instanceof MethodRef) {
-        MethodRef other = (MethodRef) obj;
-        return clazz.equals(other.clazz) && methodSig.equals(other.methodSig);
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
       }
-      return false;
+      MethodRef methodRef = (MethodRef) o;
+      return Objects.equals(enclosingClass, methodRef.enclosingClass)
+          && Objects.equals(fullMethodSig, methodRef.fullMethodSig);
     }
 
     @Override
     public int hashCode() {
-      return Objects.hash(clazz, methodSig);
+      return Objects.hash(enclosingClass, fullMethodSig);
     }
 
     @Override
     public String toString() {
-      return "MethodRef{" + "clazz='" + clazz + '\'' + ", methodSig='" + methodSig + '\'' + '}';
-    }
-  }
-
-  /** utility methods for dealing with library models */
-  final class LibraryModelUtil {
-
-    private LibraryModelUtil() {}
-
-    public static boolean hasNullableReturn(
-        LibraryModels models, Symbol.MethodSymbol symbol, Types types) {
-      // need to check if symbol is in the model or if it overrides a method in the model
-      return methodInSet(symbol, types, models.nullableReturns()) != null;
-    }
-
-    public static boolean hasNonNullReturn(
-        LibraryModels models, Symbol.MethodSymbol symbol, Types types) {
-      // need to check if symbol is in the model or if it overrides a method in the model
-      return methodInSet(symbol, types, models.nonNullReturns()) != null;
-    }
-
-    @Nullable
-    private static Symbol.MethodSymbol methodInSet(
-        Symbol.MethodSymbol symbol, Types types, ImmutableCollection<MethodRef> memberNames) {
-      if (memberNames.contains(MethodRef.fromSymbol(symbol))) {
-        return symbol;
-      }
-      for (Symbol.MethodSymbol superSymbol : ASTHelpers.findSuperMethods(symbol, types)) {
-        if (memberNames.contains(MethodRef.fromSymbol(superSymbol))) {
-          return superSymbol;
-        }
-      }
-      return null;
+      return "MethodRef{"
+          + "enclosingClass='"
+          + enclosingClass
+          + '\''
+          + ", fullMethodSig='"
+          + fullMethodSig
+          + '\''
+          + '}';
     }
   }
 }

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

@@ -536,7 +536,7 @@ private Description checkParamOverriding(
     if (NullabilityUtil.isUnannotated(overriddenMethod, config)) {
       nullableParamsOfOverriden =
           handler.onUnannotatedInvocationGetExplicitlyNullablePositions(
-              overriddenMethod, ImmutableSet.of());
+              state.context, overriddenMethod, ImmutableSet.of());
     } else {
       ImmutableSet.Builder<Integer> builder = ImmutableSet.builder();
       for (int i = startParam; i < superParamSymbols.size(); i++) {

nullaway/src/main/java/com/uber/nullaway/dataflow/AccessPathNullnessPropagation.java

@@ -208,7 +208,7 @@ private NullnessStore lambdaInitialStore(
     NullnessStore environmentMapping =
         Objects.requireNonNull(
             environmentNullness.getEnvironmentMapping(underlyingAST.getLambdaTree()),
-            "no environment stored for lambda " + underlyingAST.getLambdaTree());
+            "no environment stored for lambda");
     NullnessStore.Builder result = environmentMapping.toBuilder();
     LambdaExpressionTree code = underlyingAST.getLambdaTree();
     // need to check annotation for i'th parameter of functional interface declaration
@@ -217,7 +217,7 @@ private NullnessStore lambdaInitialStore(
         fiMethodSymbol.getParameters();
     ImmutableSet<Integer> nullableParamsFromHandler =
         handler.onUnannotatedInvocationGetExplicitlyNullablePositions(
-            fiMethodSymbol, ImmutableSet.of());
+            context, fiMethodSymbol, ImmutableSet.of());
 
     for (int i = 0; i < parameters.size(); i++) {
       LocalVariableNode param = parameters.get(i);
@@ -858,7 +858,7 @@ public TransferResult<Nullness, NullnessStore> visitMethodInvocation(
         node, callee, node.getArguments(), types, values(input), thenUpdates, bothUpdates);
     NullnessHint nullnessHint =
         handler.onDataflowVisitMethodInvocation(
-            node, types, values(input), thenUpdates, elseUpdates, bothUpdates);
+            node, types, context, values(input), thenUpdates, elseUpdates, bothUpdates);
     Nullness nullness = returnValueNullness(node, input, nullnessHint);
     if (booleanReturnType(node)) {
       ResultingStore thenStore = updateStore(input.getThenStore(), thenUpdates, bothUpdates);

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

@@ -28,6 +28,7 @@
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.dataflow.AccessPath;
@@ -64,6 +65,7 @@ public void onMatchTopLevelClass(
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
       Types types,
+      Context context,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,

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

@@ -33,6 +33,7 @@
 import com.sun.source.tree.ReturnTree;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
 import com.uber.nullaway.dataflow.NullnessStore;
@@ -101,7 +102,9 @@ public void onMatchReturn(NullAway analysis, ReturnTree tree, VisitorState state
 
   @Override
   public ImmutableSet<Integer> onUnannotatedInvocationGetExplicitlyNullablePositions(
-      Symbol.MethodSymbol methodSymbol, ImmutableSet<Integer> explicitlyNullablePositions) {
+      Context context,
+      Symbol.MethodSymbol methodSymbol,
+      ImmutableSet<Integer> explicitlyNullablePositions) {
     // NoOp
     return explicitlyNullablePositions;
   }
@@ -143,6 +146,7 @@ public NullnessStore.Builder onDataflowInitialStore(
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
       Types types,
+      Context context,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,

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

@@ -34,6 +34,7 @@
 import com.sun.source.tree.ReturnTree;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
 import com.uber.nullaway.dataflow.NullnessStore;
@@ -115,11 +116,13 @@ public void onMatchReturn(NullAway analysis, ReturnTree tree, VisitorState state
 
   @Override
   public ImmutableSet<Integer> onUnannotatedInvocationGetExplicitlyNullablePositions(
-      Symbol.MethodSymbol methodSymbol, ImmutableSet<Integer> explicitlyNullablePositions) {
+      Context context,
+      Symbol.MethodSymbol methodSymbol,
+      ImmutableSet<Integer> explicitlyNullablePositions) {
     for (Handler h : handlers) {
       explicitlyNullablePositions =
           h.onUnannotatedInvocationGetExplicitlyNullablePositions(
-              methodSymbol, explicitlyNullablePositions);
+              context, methodSymbol, explicitlyNullablePositions);
     }
     return explicitlyNullablePositions;
   }
@@ -174,6 +177,7 @@ public NullnessStore.Builder onDataflowInitialStore(
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
       Types types,
+      Context context,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,
@@ -182,7 +186,7 @@ public NullnessHint onDataflowVisitMethodInvocation(
     for (Handler h : handlers) {
       NullnessHint n =
           h.onDataflowVisitMethodInvocation(
-              node, types, inputs, thenUpdates, elseUpdates, bothUpdates);
+              node, types, context, inputs, thenUpdates, elseUpdates, bothUpdates);
       nullnessHint = nullnessHint.merge(n);
     }
     return nullnessHint;

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

@@ -29,6 +29,7 @@
 import com.sun.source.tree.Tree;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.dataflow.AccessPath;
@@ -90,6 +91,7 @@ public void onMatchTopLevelClass(
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
       Types types,
+      Context context,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,

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

@@ -33,6 +33,7 @@
 import com.sun.source.tree.ReturnTree;
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.Nullness;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
@@ -143,7 +144,9 @@ void onMatchMethodReference(
    * @return Updated parameter nullability computed by this handler.
    */
   ImmutableSet<Integer> onUnannotatedInvocationGetExplicitlyNullablePositions(
-      Symbol.MethodSymbol methodSymbol, ImmutableSet<Integer> explicitlyNullablePositions);
+      Context context,
+      Symbol.MethodSymbol methodSymbol,
+      ImmutableSet<Integer> explicitlyNullablePositions);
 
   /**
    * Called when NullAway encounters an unannotated method and asks if return value is explicitly
@@ -229,6 +232,7 @@ NullnessStore.Builder onDataflowInitialStore(
   NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
       Types types,
+      Context context,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,

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

@@ -33,6 +33,7 @@
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.Types;
+import com.sun.tools.javac.util.Context;
 import com.uber.nullaway.Config;
 import com.uber.nullaway.NullAway;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
@@ -163,6 +164,7 @@ public ImmutableSet<Integer> onUnannotatedInvocationGetNonNullPositions(
   public NullnessHint onDataflowVisitMethodInvocation(
       MethodInvocationNode node,
       Types types,
+      Context context,
       AccessPathNullnessPropagation.SubNodeValues inputs,
       AccessPathNullnessPropagation.Updates thenUpdates,
       AccessPathNullnessPropagation.Updates elseUpdates,