Skip to content

Commit

Permalink
Refactor PreservedAnnotationTreeVisitor (#955)
Browse files Browse the repository at this point in the history 
We extract some code from `visitParameterizedType` into a proper
`visitAnnotatedType` method. This will be useful for an upcoming PR that
handles `AnnotatedTypeTree`s in new array expressions. The change
introduces some complexity due to API changes in JDK 21, which we work
around through further use of `MethodHandle`s.
  • Loading branch information
@msridhar
msridhar authored May 14, 2024
1 parent c26ae17 commit e0d5c6b
Showing 1 changed file with 88 additions and 42 deletions.
130 changes: 88 additions & 42 deletions nullaway/src/main/java/com/uber/nullaway/generics/PreservedAnnotationTreeVisitor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

import static com.uber.nullaway.NullabilityUtil.castToNonNull;

import com.google.common.base.Preconditions;
import com.google.errorprone.VisitorState;
import com.google.errorprone.util.ASTHelpers;
import com.sun.source.tree.AnnotatedTypeTree;
Expand All @@ -12,6 +11,7 @@
import com.sun.source.tree.Tree;
import com.sun.source.util.SimpleTreeVisitor;
import com.sun.tools.javac.code.Attribute;
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;
Expand Down Expand Up @@ -44,52 +44,42 @@ public Type visitArrayType(ArrayTypeTree tree, Void p) {

@Override
public Type visitParameterizedType(ParameterizedTypeTree tree, Void p) {
Type.ClassType type = (Type.ClassType) ASTHelpers.getType(tree);
Preconditions.checkNotNull(type);
Type nullableType = GenericsChecks.JSPECIFY_NULLABLE_TYPE_SUPPLIER.get(state);
Type.ClassType baseType = (Type.ClassType) tree.getType().accept(this, null);
List<? extends Tree> typeArguments = tree.getTypeArguments();
List<Type> newTypeArgs = new ArrayList<>();
for (int i = 0; i < typeArguments.size(); i++) {
AnnotatedTypeTree annotatedType = null;
Tree curTypeArg = typeArguments.get(i);
// If the type argument has an annotation, it will either be an AnnotatedTypeTree, or a
// ParameterizedTypeTree in the case of a nested generic type
if (curTypeArg instanceof AnnotatedTypeTree) {
annotatedType = (AnnotatedTypeTree) curTypeArg;
} else if (curTypeArg instanceof ParameterizedTypeTree
&& ((ParameterizedTypeTree) curTypeArg).getType() instanceof AnnotatedTypeTree) {
annotatedType = (AnnotatedTypeTree) ((ParameterizedTypeTree) curTypeArg).getType();
}
List<? extends AnnotationTree> annotations =
annotatedType != null ? annotatedType.getAnnotations() : Collections.emptyList();
boolean hasNullableAnnotation = false;
for (AnnotationTree annotation : annotations) {
if (ASTHelpers.isSameType(
nullableType, ASTHelpers.getType(annotation.getAnnotationType()), state)) {
hasNullableAnnotation = true;
break;
}
}
// construct a TypeMetadata object containing a nullability annotation if needed
com.sun.tools.javac.util.List<Attribute.TypeCompound> nullableAnnotationCompound =
hasNullableAnnotation
? com.sun.tools.javac.util.List.from(
Collections.singletonList(
new Attribute.TypeCompound(
nullableType, com.sun.tools.javac.util.List.nil(), null)))
: com.sun.tools.javac.util.List.nil();
TypeMetadata typeMetadata = TYPE_METADATA_BUILDER.create(nullableAnnotationCompound);
Type currentTypeArgType = curTypeArg.accept(this, null);
Type newTypeArgType =
TYPE_METADATA_BUILDER.cloneTypeWithMetadata(currentTypeArgType, typeMetadata);
newTypeArgs.add(newTypeArgType);
newTypeArgs.add(typeArguments.get(i).accept(this, null));
}
Type.ClassType finalType =
new Type.ClassType(
type.getEnclosingType(), com.sun.tools.javac.util.List.from(newTypeArgs), type.tsym);
Type finalType = TYPE_METADATA_BUILDER.createWithBaseTypeAndTypeArgs(baseType, newTypeArgs);
return finalType;
}

@Override
public Type visitAnnotatedType(AnnotatedTypeTree annotatedType, Void unused) {
List<? extends AnnotationTree> annotations = annotatedType.getAnnotations();
boolean hasNullableAnnotation = false;
Type nullableType = GenericsChecks.JSPECIFY_NULLABLE_TYPE_SUPPLIER.get(state);
for (AnnotationTree annotation : annotations) {
if (ASTHelpers.isSameType(
nullableType, ASTHelpers.getType(annotation.getAnnotationType()), state)) {
hasNullableAnnotation = true;
break;
}
}
// construct a TypeMetadata object containing a nullability annotation if needed
com.sun.tools.javac.util.List<Attribute.TypeCompound> nullableAnnotationCompound =
hasNullableAnnotation
? com.sun.tools.javac.util.List.from(
Collections.singletonList(
new Attribute.TypeCompound(
nullableType, com.sun.tools.javac.util.List.nil(), null)))
: com.sun.tools.javac.util.List.nil();
TypeMetadata typeMetadata = TYPE_METADATA_BUILDER.create(nullableAnnotationCompound);
Type underlyingType = annotatedType.getUnderlyingType().accept(this, null);
Type newType = TYPE_METADATA_BUILDER.cloneTypeWithMetadata(underlyingType, typeMetadata);
return newType;
}

/** By default, just use the type computed by javac */
@Override
protected Type defaultAction(Tree node, Void unused) {
Expand All @@ -104,6 +94,8 @@ 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);
}

/**
Expand All @@ -129,6 +121,15 @@ public TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound>
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());
}
}

/**
Expand All @@ -138,11 +139,14 @@ public Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metadata) {
*/
private static class JDK21TypeMetadataBuilder implements TypeMetadataBuilder {

private static final MethodHandle typeMetadataHandle = createHandle();
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();
Expand All @@ -156,6 +160,32 @@ private static MethodHandle createHandle() {
}
}

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
*
Expand Down Expand Up @@ -183,7 +213,7 @@ public TypeMetadata create(com.sun.tools.javac.util.List<Attribute.TypeCompound>
ListBuffer<Attribute.TypeCompound> b = new ListBuffer<>();
b.appendList(attrs);
try {
return (TypeMetadata) typeMetadataHandle.invoke(b);
return (TypeMetadata) typeMetadataConstructorHandle.invoke(b);
} catch (Throwable e) {
throw new RuntimeException(e);
}
Expand All @@ -209,6 +239,22 @@ public Type cloneTypeWithMetadata(Type typeToBeCloned, TypeMetadata metadata) {
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. */
Expand Down

0 comments on commit e0d5c6b

Please sign in to comment.