JSpecify: preserve explicit nullability annotations on type variables when performing substitutions (#1143)

Fixes #1091 

Consider the following code:
```java
abstract class Test {
  abstract <V> @nullable V foo(Function<@nullable V, @nullable V> f);
  void testPositive(Function<String, String> f) {
     this.<String>foo(f); // error
  }
} 
```
The call to `foo` should not type check. Since the type of its parameter
`f` is `Function<@nullable V, @nullable V>`, with explicit `@Nullable`
annotations on the type variables, any `Function` passed to `foo` must
have `@Nullable` type arguments. In typechecking this code, NullAway
previously substituted the type arguments for the type variables in
`foo` just using built-in `javac` routines. But, this would yield a
formal parameter type `Function<String, String>`, as the `javac` routine
would not retain the explicit type arguments in the right places. So we
would miss reporting an error. This PR fixes the substitutions and
re-introduces the annotations on type variables, so we get the type
`Function<@nullable String, @nullable String>` for the formal parameter
at the call, and report an error correctly. Substitutions were broken in
other cases as well; substituting `@Nullable V` for `@Nullable V` (where
`V` is a type variable) yielded just `V`, which led to false positives
(like #1091).

The main logic changes are in `TypeSubstitutionUtils`. We add a new
`RestoreNullnessAnnotationsVisitor` and use it to restore nullability
annotations from type variables after performing a substitution.

We also extract the `TypeMetadataBuilder` logic to a top-level source
file, and add new methods as needed for this PR. Some of this could have
been split into a separate PR but it's a bit of a pain to extract it
now.

Author: msridhar

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

@@ -107,7 +107,7 @@ private static NullnessStore lambdaInitialStore(
     // annotations in case of generic type arguments.  Only used in JSpecify mode.
     List<Type> overridenMethodParamTypeList =
         TypeSubstitutionUtils.memberType(
-                types, castToNonNull(ASTHelpers.getType(code)), fiMethodSymbol)
+                types, castToNonNull(ASTHelpers.getType(code)), fiMethodSymbol, config)
             .getParameterTypes();
     // If fiArgumentPositionNullness[i] == null, parameter position i is unannotated
     Nullness[] fiArgumentPositionNullness = new Nullness[fiMethodParameters.size()];

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

@@ -634,13 +634,14 @@ public static void compareGenericTypeParameterNullabilityForCall(
       }
       if (enclosingType != null) {
         invokedMethodType =
-            TypeSubstitutionUtils.memberType(state.getTypes(), enclosingType, methodSymbol);
+            TypeSubstitutionUtils.memberType(state.getTypes(), enclosingType, methodSymbol, config);
       }
     }
     // substitute type arguments for generic methods
     if (tree instanceof MethodInvocationTree && methodSymbol.type instanceof Type.ForAll) {
       invokedMethodType =
-          substituteTypeArgsInGenericMethodType((MethodInvocationTree) tree, methodSymbol, state);
+          substituteTypeArgsInGenericMethodType(
+              (MethodInvocationTree) tree, methodSymbol, state, config);
     }
     List<Type> formalParamTypes = invokedMethodType.getParameterTypes();
     int n = formalParamTypes.size();
@@ -706,7 +707,7 @@ public static void checkTypeParameterNullnessForMethodOverriding(
     // method's class
     Type methodWithTypeParams =
         TypeSubstitutionUtils.memberType(
-            state.getTypes(), overridingMethod.owner.type, overriddenMethod);
+            state.getTypes(), overridingMethod.owner.type, overriddenMethod, analysis.getConfig());
 
     checkTypeParameterNullnessForOverridingMethodReturnType(
         tree, methodWithTypeParams, analysis, state);
@@ -786,7 +787,7 @@ public static Nullness getGenericMethodReturnTypeNullness(
       return Nullness.NONNULL;
     }
     Type overriddenMethodType =
-        TypeSubstitutionUtils.memberType(state.getTypes(), enclosingType, method);
+        TypeSubstitutionUtils.memberType(state.getTypes(), enclosingType, method, config);
     verify(
         overriddenMethodType instanceof ExecutableType,
         "expected ExecutableType but instead got %s",
@@ -835,7 +836,8 @@ public static Nullness getGenericReturnNullnessAtInvocation(
     if (!invokedMethodSymbol.getTypeParameters().isEmpty()) {
       // Substitute type arguments inside the return type
       Type substitutedReturnType =
-          substituteTypeArgsInGenericMethodType(tree, invokedMethodSymbol, state).getReturnType();
+          substituteTypeArgsInGenericMethodType(tree, invokedMethodSymbol, state, config)
+              .getReturnType();
       // If this condition evaluates to false, we fall through to the subsequent logic, to handle
       // type variables declared on the enclosing class
       if (substitutedReturnType != null
@@ -875,20 +877,22 @@ private static com.sun.tools.javac.util.List<Type> convertTreesToTypes(
    * @param methodInvocationTree the method invocation tree
    * @param methodSymbol symbol for the invoked generic method
    * @param state the visitor state
+   * @param config the NullAway config
    * @return the substituted method type for the generic method
    */
   private static Type substituteTypeArgsInGenericMethodType(
       MethodInvocationTree methodInvocationTree,
       Symbol.MethodSymbol methodSymbol,
-      VisitorState state) {
+      VisitorState state,
+      Config config) {
 
     List<? extends Tree> typeArgumentTrees = methodInvocationTree.getTypeArguments();
     com.sun.tools.javac.util.List<Type> explicitTypeArgs = convertTreesToTypes(typeArgumentTrees);
 
     Type.ForAll forAllType = (Type.ForAll) methodSymbol.type;
     Type.MethodType underlyingMethodType = (Type.MethodType) forAllType.qtype;
     return TypeSubstitutionUtils.subst(
-        state.getTypes(), underlyingMethodType, forAllType.tvars, explicitTypeArgs);
+        state.getTypes(), underlyingMethodType, forAllType.tvars, explicitTypeArgs, config);
   }
 
   /**
@@ -936,7 +940,7 @@ public static Nullness getGenericParameterNullnessAtInvocation(
       // Substitute the argument types within the MethodType
       // NOTE: if explicitTypeArgs is empty, this is a noop
       List<Type> substitutedParamTypes =
-          substituteTypeArgsInGenericMethodType(tree, invokedMethodSymbol, state)
+          substituteTypeArgsInGenericMethodType(tree, invokedMethodSymbol, state, config)
               .getParameterTypes();
       // If this condition evaluates to false, we fall through to the subsequent logic, to handle
       // type variables declared on the enclosing class
@@ -1022,7 +1026,8 @@ public static Nullness getGenericMethodParameterNullness(
       // @Nullable annotation is handled elsewhere)
       return Nullness.NONNULL;
     }
-    Type methodType = TypeSubstitutionUtils.memberType(state.getTypes(), enclosingType, method);
+    Type methodType =
+        TypeSubstitutionUtils.memberType(state.getTypes(), enclosingType, method, config);
     Type paramType = methodType.getParameterTypes().get(parameterIndex);
     return getTypeNullness(paramType, config);
   }

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

@@ -1,6 +1,7 @@
 package com.uber.nullaway.generics;
 
 import static com.uber.nullaway.NullabilityUtil.castToNonNull;
+import static com.uber.nullaway.generics.TypeMetadataBuilder.TYPE_METADATA_BUILDER;
 
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.AnnotatedTypeTree;
@@ -14,12 +15,8 @@
 import com.sun.tools.javac.code.Symbol;
 import com.sun.tools.javac.code.Type;
 import com.sun.tools.javac.code.TypeMetadata;
-import com.sun.tools.javac.util.ListBuffer;
 import com.uber.nullaway.Config;
 import com.uber.nullaway.Nullness;
-import java.lang.invoke.MethodHandle;
-import java.lang.invoke.MethodHandles;
-import java.lang.invoke.MethodType;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -58,7 +55,8 @@ public Type visitParameterizedType(ParameterizedTypeTree tree, Void p) {
     for (int i = 0; i < typeArguments.size(); i++) {
       newTypeArgs.add(typeArguments.get(i).accept(this, null));
     }
-    Type finalType = TYPE_METADATA_BUILDER.createWithBaseTypeAndTypeArgs(baseType, newTypeArgs);
+    Type finalType =
+        TYPE_METADATA_BUILDER.createClassType(baseType, baseType.getEnclosingType(), newTypeArgs);
     return finalType;
   }
 
@@ -97,181 +95,4 @@ public Type visitAnnotatedType(AnnotatedTypeTree annotatedType, Void unused) {
   protected Type defaultAction(Tree node, Void unused) {
     return castToNonNull(ASTHelpers.getType(node));
   }
-
-  /**
-   * Abstracts over the different APIs for building {@link TypeMetadata} objects in different JDK
-   * versions.
-   */
-  private interface TypeMetadataBuilder {
-    TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound> attrs);
-
-    Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metaData);
-
-    Type createWithBaseTypeAndTypeArgs(Type baseType, List<Type> typeArgs);
-  }
-
-  /**
-   * Provides implementations for methods under TypeMetadataBuilder compatible with JDK 17 and
-   * earlier versions.
-   */
-  private static class JDK17AndEarlierTypeMetadataBuilder implements TypeMetadataBuilder {
-
-    @Override
-    public TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound> attrs) {
-      return new TypeMetadata(new TypeMetadata.Annotations(attrs));
-    }
-
-    /**
-     * Clones the given type with the specified Metadata for getting the right nullability
-     * annotations.
-     *
-     * @param typeToBeCloned The Type we want to clone with the required Nullability Metadata
-     * @param metadata The required Nullability metadata which is lost from the type
-     * @return Type after it has been cloned by applying the required Nullability metadata
-     */
-    @Override
-    public Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metadata) {
-      return typeToBeCloned.cloneWithMetadata(metadata);
-    }
-
-    @Override
-    public Type createWithBaseTypeAndTypeArgs(Type baseType, List<Type> typeArgs) {
-      return new Type.ClassType(
-          baseType.getEnclosingType(),
-          com.sun.tools.javac.util.List.from(typeArgs),
-          baseType.tsym,
-          baseType.getMetadata());
-    }
-  }
-
-  /**
-   * Provides implementations for methods under TypeMetadataBuilder compatible with the updates made
-   * to the library methods for Jdk 21. The implementation calls the logic specific to JDK 21
-   * indirectly using MethodHandles since we still need the code to compile on earlier versions.
-   */
-  private static class JDK21TypeMetadataBuilder implements TypeMetadataBuilder {
-
-    private static final MethodHandle typeMetadataConstructorHandle = createHandle();
-    private static final MethodHandle addMetadataHandle =
-        createVirtualMethodHandle(Type.class, TypeMetadata.class, Type.class, "addMetadata");
-    private static final MethodHandle dropMetadataHandle =
-        createVirtualMethodHandle(Type.class, Class.class, Type.class, "dropMetadata");
-    private static final MethodHandle getMetadataHandler = createGetMetadataHandle();
-    private static final MethodHandle classTypeConstructorHandle =
-        createClassTypeConstructorHandle();
-
-    private static MethodHandle createHandle() {
-      MethodHandles.Lookup lookup = MethodHandles.lookup();
-      MethodType mt = MethodType.methodType(void.class, com.sun.tools.javac.util.ListBuffer.class);
-      try {
-        return lookup.findConstructor(TypeMetadata.Annotations.class, mt);
-      } catch (NoSuchMethodException e) {
-        throw new RuntimeException(e);
-      } catch (IllegalAccessException e) {
-        throw new RuntimeException(e);
-      }
-    }
-
-    private static MethodHandle createGetMetadataHandle() {
-      MethodHandles.Lookup lookup = MethodHandles.lookup();
-      MethodType mt = MethodType.methodType(com.sun.tools.javac.util.List.class);
-      try {
-        return lookup.findVirtual(Type.class, "getMetadata", mt);
-      } catch (NoSuchMethodException | IllegalAccessException e) {
-        throw new RuntimeException(e);
-      }
-    }
-
-    private static MethodHandle createClassTypeConstructorHandle() {
-      try {
-        MethodHandles.Lookup lookup = MethodHandles.lookup();
-        MethodType methodType =
-            MethodType.methodType(
-                void.class, // return type for a constructor is void
-                Type.class,
-                com.sun.tools.javac.util.List.class,
-                Symbol.TypeSymbol.class,
-                com.sun.tools.javac.util.List.class);
-        return lookup.findConstructor(Type.ClassType.class, methodType);
-      } catch (NoSuchMethodException | IllegalAccessException e) {
-        throw new RuntimeException(e);
-      }
-    }
-
-    /**
-     * Used to get a MethodHandle for a virtual method from the specified class
-     *
-     * @param retTypeClass Class to indicate the return type of the desired method
-     * @param paramTypeClass Class to indicate the parameter type of the desired method
-     * @param refClass Class within which the desired method is contained
-     * @param methodName Name of the desired method
-     * @return The appropriate MethodHandle for the virtual method
-     */
-    private static MethodHandle createVirtualMethodHandle(
-        Class<?> retTypeClass, Class<?> paramTypeClass, Class<?> refClass, String methodName) {
-      MethodHandles.Lookup lookup = MethodHandles.lookup();
-      MethodType mt = MethodType.methodType(retTypeClass, paramTypeClass);
-      try {
-        return lookup.findVirtual(refClass, methodName, mt);
-      } catch (NoSuchMethodException e) {
-        throw new RuntimeException(e);
-      } catch (IllegalAccessException e) {
-        throw new RuntimeException(e);
-      }
-    }
-
-    @Override
-    public TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound> attrs) {
-      ListBuffer<Attribute.TypeCompound> b = new ListBuffer<>();
-      b.appendList(attrs);
-      try {
-        return (TypeMetadata) typeMetadataConstructorHandle.invoke(b);
-      } catch (Throwable e) {
-        throw new RuntimeException(e);
-      }
-    }
-
-    /**
-     * Calls dropMetadata and addMetadata using MethodHandles for JDK 21, which removed the previous
-     * cloneWithMetadata method.
-     *
-     * @param typeToBeCloned The Type we want to clone with the required Nullability metadata
-     * @param metadata The required Nullability metadata
-     * @return Cloned Type with the necessary Nullability metadata
-     */
-    @Override
-    public Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metadata) {
-      try {
-        // In JDK 21 addMetadata works if there is no metadata associated with the type, so we
-        // create a copy without the existing metadata first and then add it
-        Type clonedTypeWithoutMetadata =
-            (Type) dropMetadataHandle.invoke(typeToBeCloned, metadata.getClass());
-        return (Type) addMetadataHandle.invoke(clonedTypeWithoutMetadata, metadata);
-      } catch (Throwable e) {
-        throw new RuntimeException(e);
-      }
-    }
-
-    @Override
-    public Type createWithBaseTypeAndTypeArgs(Type baseType, List<Type> typeArgs) {
-      try {
-        com.sun.tools.javac.util.List<TypeMetadata> metadata =
-            (com.sun.tools.javac.util.List<TypeMetadata>) getMetadataHandler.invoke(baseType);
-        return (Type)
-            classTypeConstructorHandle.invoke(
-                baseType.getEnclosingType(),
-                com.sun.tools.javac.util.List.from(typeArgs),
-                baseType.tsym,
-                metadata);
-      } catch (Throwable e) {
-        throw new RuntimeException(e);
-      }
-    }
-  }
-
-  /** The TypeMetadataBuilder to be used for the current JDK version. */
-  private static final TypeMetadataBuilder TYPE_METADATA_BUILDER =
-      Runtime.version().feature() >= 21
-          ? new JDK21TypeMetadataBuilder()
-          : new JDK17AndEarlierTypeMetadataBuilder();
 }