add array type inference to visitor

uber · Feb 21, 2025 · 461c309 · 461c309
1 parent f7c4a3f
commit 461c309
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 24 deletions.
diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java b/nullaway/src/main/java/com/uber/nullaway/generics/GenericsChecks.java
@@ -435,15 +435,12 @@ public void checkTypeParameterNullnessForAssignability(
         // update inferredTypes cache for assignments
         if (methodSymbol.type instanceof Type.ForAll // generic method call
             && methodInvocationTree.getTypeArguments().isEmpty() // no explicit generic arguments
-            && lhsType != null
-            && !lhsType.getTypeArguments().isEmpty()) { // lhs type has generic
-          if (lhsType != null) {
-            Type returnType = methodSymbol.getReturnType();
-            @Nullable Map<Type, Type> genericNullness =
-                returnType.accept(new InferTypeVisitor(config), lhsType);
-            if (genericNullness != null) {
-              inferredTypes.put(rhsTree, genericNullness);
-            }
+            && lhsType != null) {
+          Type returnType = methodSymbol.getReturnType();
+          @Nullable Map<Type, Type> genericNullness =
+              returnType.accept(new InferTypeVisitor(config), lhsType);
+          if (genericNullness != null) {
+            inferredTypes.put(rhsTree, genericNullness);
           }
         }
       }
@@ -458,20 +455,18 @@ public void checkTypeParameterNullnessForAssignability(
     }
     Type rhsType = getTreeType(rhsTree, config);
 
-    if (rhsType != null && !rhsType.getTypeArguments().isEmpty()) {
-      if (rhsTree instanceof MethodInvocationTree) {
-        // recreate rhsType using inferred types
-        Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol((MethodInvocationTree) rhsTree);
-        if (inferredTypes.containsKey(rhsTree)) {
-          Map<Type, Type> genericNullness = inferredTypes.get(rhsTree);
-          com.sun.tools.javac.util.List<Type> from =
-              com.sun.tools.javac.util.List.from(genericNullness.keySet());
-          com.sun.tools.javac.util.List<Type> to =
-              com.sun.tools.javac.util.List.from(genericNullness.values());
-          rhsType =
-              TypeSubstitutionUtils.subst(
-                  state.getTypes(), methodSymbol.getReturnType(), from, to, config);
-        }
+    if (rhsType != null && rhsTree instanceof MethodInvocationTree) {
+      // recreate rhsType using inferred types
+      Symbol.MethodSymbol methodSymbol = ASTHelpers.getSymbol((MethodInvocationTree) rhsTree);
+      if (inferredTypes.containsKey(rhsTree)) {
+        Map<Type, Type> genericNullness = inferredTypes.get(rhsTree);
+        com.sun.tools.javac.util.List<Type> from =
+            com.sun.tools.javac.util.List.from(genericNullness.keySet());
+        com.sun.tools.javac.util.List<Type> to =
+            com.sun.tools.javac.util.List.from(genericNullness.values());
+        rhsType =
+            TypeSubstitutionUtils.subst(
+                state.getTypes(), methodSymbol.getReturnType(), from, to, config);
       }
     }
 

diff --git a/nullaway/src/main/java/com/uber/nullaway/generics/InferTypeVisitor.java b/nullaway/src/main/java/com/uber/nullaway/generics/InferTypeVisitor.java
@@ -21,7 +21,8 @@ public class InferTypeVisitor extends Types.DefaultTypeVisitor<@Nullable Map<Typ
     Map<Type, Type> genericNullness = new HashMap<>();
     // for each type parameter, call accept with this visitor and add all results to one map
     com.sun.tools.javac.util.List<Type> rhsTypeArguments = rhsType.getTypeArguments();
-    com.sun.tools.javac.util.List<Type> lhsTypeArguments = lhsType.getTypeArguments();
+    com.sun.tools.javac.util.List<Type> lhsTypeArguments =
+        ((Type.ClassType) lhsType).getTypeArguments();
     for (int i = 0; i < rhsTypeArguments.size(); i++) {
       Type rhsTypeArg = rhsTypeArguments.get(i);
       Type lhsTypeArg = lhsTypeArguments.get(i);
@@ -52,6 +53,14 @@ public Map<Type, Type> visitTypeVar(Type.TypeVar rhsType, Type lhsType) { // typ
     return genericNullness;
   }
 
+  @Override
+  public @Nullable Map<Type, Type> visitArrayType(Type.ArrayType rhsType, Type lhsType) {
+    Type rhsComponentType = rhsType.elemtype;
+    Type lhsComponentType = ((Type.ArrayType) lhsType).elemtype;
+    Map<Type, Type> genericNullness = rhsComponentType.accept(this, lhsComponentType);
+    return genericNullness;
+  }
+
   @Override
   public @Nullable Map<Type, Type> visitType(Type t, Type type) {
     return null;