Initial handling of constructor diamond operators

nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java

@@ -28,6 +28,7 @@
 import com.sun.source.tree.Tree;
 import com.sun.source.tree.VariableTree;
 import com.sun.source.util.TreePath;
+import com.sun.source.util.TreePathScanner;
 import com.sun.source.util.TreeScanner;
 import com.sun.tools.javac.code.Attribute;
 import com.sun.tools.javac.code.Symbol;
@@ -36,6 +37,7 @@
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.Types;
 import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.tree.TreeInfo;
 import com.sun.tools.javac.util.Name;
 import com.sun.tools.javac.util.Names;
 import com.uber.nullaway.CodeAnnotationInfo;
@@ -434,14 +436,26 @@ private void reportInvalidOverridingMethodParamTypeError(
       }
       return result;
     }
-    if (tree instanceof NewClassTree
-        && ((NewClassTree) tree).getIdentifier() instanceof ParameterizedTypeTree paramTypedTree) {
-      if (paramTypedTree.getTypeArguments().isEmpty()) {
-        // diamond operator, which we do not yet support; for now, return null
-        // TODO: support diamond operators
-        return null;
+    if (tree instanceof NewClassTree newClassTree) {
+      if (TreeInfo.isDiamond((JCTree) newClassTree)) {
+        if (newClassTree.getClassBody() != null) {
+          // Keep existing behavior for diamond anonymous classes, which are not yet fully
+          // supported.
+          return null;
+        }
+        // For constructor calls using diamond operator, infer from assignment context.
+        // TODO handle diamond constructor calls passed to generic methods
+        // https://github.com/uber/NullAway/issues/1470
+        Type fromAssignmentContext = getDiamondTypeFromContext(newClassTree, state);
+        if (fromAssignmentContext != null) {
+          return fromAssignmentContext;
+        }
       }
-      return typeWithPreservedAnnotations(paramTypedTree);
+      if (newClassTree.getIdentifier() instanceof ParameterizedTypeTree paramTypedTree
+          && !paramTypedTree.getTypeArguments().isEmpty()) {
+        return typeWithPreservedAnnotations(paramTypedTree);
+      }
+      return ASTHelpers.getType(tree);
     } else if (tree instanceof NewArrayTree
         && ((NewArrayTree) tree).getType() instanceof AnnotatedTypeTree) {
       return typeWithPreservedAnnotations(tree);
@@ -522,6 +536,149 @@ private void reportInvalidOverridingMethodParamTypeError(
     }
   }
 
+  /**
+   * Gets the type of a constructor call using a diamond operator from its assignment context, if
+   * available.
+   */
+  private @Nullable Type getDiamondTypeFromContext(NewClassTree tree, VisitorState state) {
+    TreePath treePath = findPathToSubtree(state.getPath(), tree);
+    if (treePath == null) {
+      return null;
+    }
+    return getDiamondTypeFromParentContext(tree, state, castToNonNull(treePath.getParentPath()));
+  }
+
+  /**
+   * Computes the assignment-context type for an inferred constructor call, given a path to its
+   * parent context.
+   */
+  private @Nullable Type getDiamondTypeFromParentContext(
+      NewClassTree tree, VisitorState state, TreePath parentPath) {
+    Tree parent = parentPath.getLeaf();
+    while (parent instanceof ParenthesizedTree) {
+      parentPath = parentPath.getParentPath();
+      if (parentPath == null) {
+        return null;
+      }
+      parent = parentPath.getLeaf();
+    }
+    if (parent instanceof VariableTree || parent instanceof AssignmentTree) {
+      return getTreeType(parent, state);
+    }
+    if (parent instanceof ReturnTree) {
+      TreePath enclosingMethodOrLambda =
+          NullabilityUtil.findEnclosingMethodOrLambdaOrInitializer(parentPath);
+      if (enclosingMethodOrLambda != null
+          && enclosingMethodOrLambda.getLeaf() instanceof MethodTree enclosingMethod) {
+        Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(enclosingMethod);
+        if (methodSymbol != null) {
+          return methodSymbol.getReturnType();
+        }
+      }
+      return null;
+    }
+    if (parent instanceof MethodInvocationTree parentInvocation) {
+      if (isGenericCallNeedingInference(parentInvocation)) {
+        // TODO support full integration of diamond constructor calls with generic method inference
+        // for now, just give up and return null
+        return null;
+      }
+      Type methodType = ASTHelpers.getType(parentInvocation.getMethodSelect());
+      if (methodType == null) {
+        return null;
+      }
+      return getFormalParameterTypeForArgument(parentInvocation, methodType.asMethodType(), tree);
+    }
+    if (parent instanceof NewClassTree parentConstructorCall) {
+      // get the type returned by the parent constructor call
+      Type parentClassType = getTreeType(parentConstructorCall, state.withPath(parentPath));
+      if (parentClassType != null) {
+        Symbol parentCtorSymbol = ASTHelpers.getSymbol(parentConstructorCall);
+        // get the proper type for the constructor, as a member of the type returned by the
+        // constructor
+        Type parentCtorType =
+            TypeSubstitutionUtils.memberType(
+                state.getTypes(), parentClassType, parentCtorSymbol, config);
+        return getFormalParameterTypeForArgument(
+            parentConstructorCall, parentCtorType.asMethodType(), tree);
+      }
+    }
+    return null;
+  }
+
+  /** Returns the inferred/declared formal parameter type corresponding to {@code argumentTree}. */
+  private @Nullable Type getFormalParameterTypeForArgument(
+      Tree invocationTree, Type.MethodType invocationType, Tree argumentTree) {
+    AtomicReference<@Nullable Type> formalParamTypeRef = new AtomicReference<>();
+    new InvocationArguments(invocationTree, invocationType)
+        .forEach(
+            (arg, pos, formalParamType, unused) -> {
+              if (ASTHelpers.stripParentheses(arg) == argumentTree) {
+                formalParamTypeRef.set(formalParamType);
+              }
+            });
+    return formalParamTypeRef.get();
+  }
+
+  /** Finds the path to {@code target} within {@code rootPath}, or null when not found. */
+  private static @Nullable TreePath findPathToSubtree(TreePath rootPath, Tree target) {
+    if (rootPath.getLeaf() == target) {
+      return rootPath;
+    }
+    return new TreePathScanner<@Nullable TreePath, @Nullable TreePath>() {
+      @Override
+      public @Nullable TreePath scan(Tree tree, @Nullable TreePath prevPath) {
+        if (tree == target) {
+          TreePath currentPath = getCurrentPath();
+          if (currentPath != null && currentPath.getLeaf() == tree) {
+            return currentPath;
+          }
+          // When overriding scan(), getCurrentPath() can still point at the parent.
+          return new TreePath(currentPath, tree);
+        }
+        return super.scan(tree, null);
+      }
+
+      @Override
+      public @Nullable TreePath reduce(@Nullable TreePath r1, @Nullable TreePath r2) {
+        // we should only find the target once, so at most one of r1 and r2 should be non-null
+        Verify.verify(r1 == null || r2 == null);
+        return r1 != null ? r1 : r2;
+      }
+    }.scan(rootPath, null);
+  }
+
+  /**
+   * Creates a state whose path points to {@code subtree}, when that subtree is reachable.
+   *
+   * @param state the original state with a path to some subtree of the AST
+   * @param subtree the subtree to find within the original state's path
+   * @return a state with the same information as the original, but with a path to {@code subtree}
+   *     if it is reachable from the original state's path, or the original state if {@code subtree}
+   *     is not reachable
+   */
+  private static VisitorState withPathToSubtree(VisitorState state, Tree subtree) {
+    TreePath subtreePath = findPathToSubtree(state.getPath(), subtree);
+    return subtreePath == null ? state : state.withPath(subtreePath);
+  }
+
+  /**
+   * Returns true when javac inferred class type arguments for a constructor call, i.e. there are
+   * instantiated type arguments at the type level, but no explicit non-diamond source type args.
+   */
+  private static boolean hasInferredClassTypeArguments(NewClassTree newClassTree) {
+    if (newClassTree.getClassBody() != null) {
+      // we still need to properly handle anonymous classes
+      return false;
+    }
+    if (!TreeInfo.isDiamond((JCTree) newClassTree)) {
+      // explicit class type arguments in source
+      return false;
+    }
+    Type newClassType = ASTHelpers.getType(newClassTree);
+    return newClassType != null && !newClassType.getTypeArguments().isEmpty();
+  }
+
   /**
    * Gets the inferred type of lambda parameter, if the lambda was passed to a generic method and
    * its type was inferred previously
@@ -606,7 +763,7 @@ public void checkTypeParameterNullnessForAssignability(Tree tree, VisitorState s
         && isAssignmentToField(tree)) {
       maybeStoreLambdaTypeFromTarget(lambdaExpressionTree, lhsType);
     }
-    Type rhsType = getTreeType(rhsTree, state);
+    Type rhsType = getTreeType(rhsTree, withPathToSubtree(state, rhsTree));
     if (rhsType != null) {
       if (isGenericCallNeedingInference(rhsTree)) {
         rhsType =
@@ -1124,6 +1281,7 @@ private Type updateTypeWithNullness(
   private static boolean isGenericCallNeedingInference(ExpressionTree argument) {
     // For now, we only support calls to generic methods.
     // TODO also support calls to generic constructors that use the diamond operator
+    // https://github.com/uber/NullAway/issues/1470
     if (argument instanceof MethodInvocationTree methodInvocation) {
       Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol(methodInvocation);
       // true for generic method calls with no explicit type arguments
@@ -1154,7 +1312,7 @@ public void checkTypeParameterNullnessForFunctionReturnType(
       // bail out of any checking involving raw types for now
       return;
     }
-    Type returnExpressionType = getTreeType(retExpr, state);
+    Type returnExpressionType = getTreeType(retExpr, withPathToSubtree(state, retExpr));
     if (returnExpressionType != null) {
       if (isGenericCallNeedingInference(retExpr)) {
         returnExpressionType =
@@ -1304,7 +1462,21 @@ public void compareGenericTypeParameterNullabilityForCall(
       return;
     }
     Type invokedMethodType = methodSymbol.type;
-    Type enclosingType = getEnclosingTypeForCallExpression(methodSymbol, tree, null, state, false);
+    Type enclosingType = null;
+    if (tree instanceof NewClassTree newClassTree) {
+      if (hasInferredClassTypeArguments(newClassTree)) {
+        TreePath currentPath = state.getPath();
+        if (currentPath != null && ASTHelpers.stripParentheses(currentPath.getLeaf()) == tree) {
+          TreePath parentPath = currentPath.getParentPath();
+          if (parentPath != null) {
+            enclosingType = getDiamondTypeFromParentContext(newClassTree, state, parentPath);
+          }
+        }
+      }
+    }
+    if (enclosingType == null) {
+      enclosingType = getEnclosingTypeForCallExpression(methodSymbol, tree, null, state, false);
+    }
     if (enclosingType != null) {
       invokedMethodType =
           TypeSubstitutionUtils.memberType(state.getTypes(), enclosingType, methodSymbol, config);
@@ -1335,7 +1507,8 @@ public void compareGenericTypeParameterNullabilityForCall(
               if (inferredPolyType != null) {
                 actualParameterType = inferredPolyType;
               } else {
-                actualParameterType = getTreeType(currentActualParam, state);
+                actualParameterType =
+                    getTreeType(currentActualParam, withPathToSubtree(state, currentActualParam));
               }
               if (actualParameterType != null) {
                 if (isGenericCallNeedingInference(currentActualParam)) {
@@ -1670,16 +1843,11 @@ private InvocationAndContext getInvocationAndContextForInference(
         }
         // the generic invocation is either a regular parameter to the parent call, or the
         // receiver expression
-        AtomicReference<@Nullable Type> formalParamTypeRef = new AtomicReference<>();
-        Type type = ASTHelpers.getSymbol(parentInvocation).type;
-        new InvocationArguments(parentInvocation, type.asMethodType())
-            .forEach(
-                (arg, pos, formalParamType, unused) -> {
-                  if (ASTHelpers.stripParentheses(arg) == invocation) {
-                    formalParamTypeRef.set(formalParamType);
-                  }
-                });
-        Type formalParamType = formalParamTypeRef.get();
+        Type formalParamType =
+            getFormalParameterTypeForArgument(
+                parentInvocation,
+                ASTHelpers.getSymbol(parentInvocation).type.asMethodType(),
+                invocation);
         if (formalParamType == null) {
           // this can happen if the invocation is the receiver expression of the call, e.g.,
           // id(x).foo() (note that foo() need not be generic)
@@ -1830,14 +1998,14 @@ public Nullness getGenericParameterNullnessAtInvocation(
                   false,
                   calledFromDataflow);
         } else {
-          enclosingType = getTreeType(receiver, state);
+          enclosingType = getTreeType(receiver, withPathToSubtree(state, receiver));
         }
       }
     } else {
       Verify.verify(tree instanceof NewClassTree);
       // for a constructor invocation, the type from the invocation itself is the "enclosing type"
       // for the purposes of determining type arguments
-      enclosingType = getTreeType(tree, state);
+      enclosingType = getTreeType(tree, withPathToSubtree(state, tree));
     }
     return enclosingType;
   }