Have MemoizeExtension recognise @Nullable as a type annotation.

It currently only recognises @Nullable as a method annotation, typically for javax.annotation.Nullable, but does not recognise the TYPE_USE version from the checker framework.

We also annotate the corresponding constructor parameter, though currently only for the TYPE_USE case.

While testing this I discovered that we can pick up two versions of MemoizeExtension because it changed packages. Specifically I saw Maven picking up AutoValue 1.5.3 via compile-testing. The way ServiceLoader works means that we will load and run both versions of MemoizeExtension in this case. So I added a hack to AutoValueProcessor so it will ignore the old version if it sees it.

RELNOTES=MemoizeExtension recognises @Nullable type annotations, not just method annotations.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=218175594

3 files changed, 109 insertions, 18 deletions

diff --git a/value/src/main/java/com/google/auto/value/extension/memoized/processor/MemoizeExtension.java b/value/src/main/java/com/google/auto/value/extension/memoized/processor/MemoizeExtension.java
index ad52f697..8ef2fb65 100644
--- a/value/src/main/java/com/google/auto/value/extension/memoized/processor/MemoizeExtension.java
+++ b/value/src/main/java/com/google/auto/value/extension/memoized/processor/MemoizeExtension.java
@@ -62,6 +62,7 @@ import javax.lang.model.element.ExecutableElement;
 import javax.lang.model.element.Modifier;
 import javax.lang.model.element.TypeElement;
 import javax.lang.model.element.TypeParameterElement;
+import javax.lang.model.type.TypeMirror;
 import javax.lang.model.util.Elements;
 import javax.tools.Diagnostic.Kind;
 
@@ -165,8 +166,7 @@ public final class MemoizeExtension extends AutoValueExtension {
     private MethodSpec constructor() {
       MethodSpec.Builder constructor = constructorBuilder();
       for (Map.Entry<String, ExecutableElement> property : context.properties().entrySet()) {
-        constructor.addParameter(
-            TypeName.get(property.getValue().getReturnType()), property.getKey() + "$");
+        constructor.addParameter(annotatedReturnType(property.getValue()), property.getKey() + "$");
       }
       List<String> namesWithDollars = new ArrayList<String>();
       for (String property : context.properties().keySet()) {
@@ -189,8 +189,7 @@ public final class MemoizeExtension extends AutoValueExtension {
         this.method = method;
         validate();
         cacheField =
-            buildCacheField(
-                TypeName.get(method.getReturnType()), method.getSimpleName().toString());
+            buildCacheField(annotatedReturnType(method), method.getSimpleName().toString());
         fields.add(cacheField);
         override =
             methodBuilder(method.getSimpleName().toString())
@@ -299,14 +298,9 @@ public final class MemoizeExtension extends AutoValueExtension {
         if (method.getReturnType().getKind().isPrimitive()) {
           return new CheckBooleanField();
         }
-        for (AnnotationMirror annotationMirror : method.getAnnotationMirrors()) {
-          if (annotationMirror
-              .getAnnotationType()
-              .asElement()
-              .getSimpleName()
-              .contentEquals("Nullable")) {
-            return new CheckBooleanField();
-          }
+        if (containsNullable(method.getAnnotationMirrors())
+            || containsNullable(method.getReturnType().getAnnotationMirrors())) {
+          return new CheckBooleanField();
         }
         return new NullMeansUninitialized();
       }
@@ -373,4 +367,21 @@ public final class MemoizeExtension extends AutoValueExtension {
     }
     return Optional.of(AnnotationSpec.builder(LAZY_INIT).build());
   }
+
+  /** True if one of the given annotations is {@code @Nullable} in any package. */
+  private static boolean containsNullable(List<? extends AnnotationMirror> annotations) {
+    return annotations.stream()
+        .map(a -> a.getAnnotationType().asElement().getSimpleName())
+        .anyMatch(n -> n.contentEquals("Nullable"));
+  }
+
+  /** The return type of the given method, including type annotations. */
+  private static TypeName annotatedReturnType(ExecutableElement method) {
+    TypeMirror returnType = method.getReturnType();
+    List<AnnotationSpec> annotations =
+        returnType.getAnnotationMirrors().stream()
+            .map(AnnotationSpec::get)
+            .collect(toList());
+    return TypeName.get(returnType).annotated(annotations);
+  }
 }
diff --git a/value/src/main/java/com/google/auto/value/processor/AutoValueProcessor.java b/value/src/main/java/com/google/auto/value/processor/AutoValueProcessor.java
index 594645b3..cc555e9c 100644
--- a/value/src/main/java/com/google/auto/value/processor/AutoValueProcessor.java
+++ b/value/src/main/java/com/google/auto/value/processor/AutoValueProcessor.java
@@ -30,6 +30,7 @@ import com.google.common.collect.ImmutableBiMap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableListMultimap;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
 import java.lang.annotation.Annotation;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -67,6 +68,13 @@ import net.ltgt.gradle.incap.IncrementalAnnotationProcessorType;
 public class AutoValueProcessor extends AutoValueOrOneOfProcessor {
   private static final String OMIT_IDENTIFIERS_OPTION = "com.google.auto.value.OmitIdentifiers";
 
+  // We moved MemoizeExtension to a different package, which had an unexpected effect:
+  // now if an old version of AutoValue is in the class path, ServiceLoader can pick up both the
+  // old and the new versions of MemoizeExtension. So we exclude the old version if we see it.
+  // The new version will be bundled with this processor so we should always find it.
+  private static final String OLD_MEMOIZE_EXTENSION =
+      "com.google.auto.value.extension.memoized.MemoizeExtension";
+
   public AutoValueProcessor() {
     this(AutoValueProcessor.class.getClassLoader());
   }
@@ -97,8 +105,11 @@ public class AutoValueProcessor extends AutoValueOrOneOfProcessor {
 
     if (extensions == null) {
       try {
-        extensions =
-            ImmutableList.copyOf(ServiceLoader.load(AutoValueExtension.class, loaderForExtensions));
+        ServiceLoader<AutoValueExtension> serviceLoader =
+            ServiceLoader.load(AutoValueExtension.class, loaderForExtensions);
+        extensions = ImmutableList.copyOf(
+            Iterables.filter(
+                serviceLoader, ext -> !ext.getClass().getName().equals(OLD_MEMOIZE_EXTENSION)));
         // ServiceLoader.load returns a lazily-evaluated Iterable, so evaluate it eagerly now
         // to discover any exceptions.
       } catch (Throwable t) {
diff --git a/value/src/test/java/com/google/auto/value/extension/memoized/MemoizedTest.java b/value/src/test/java/com/google/auto/value/extension/memoized/MemoizedTest.java
index 50331561..ff312048 100644
--- a/value/src/test/java/com/google/auto/value/extension/memoized/MemoizedTest.java
+++ b/value/src/test/java/com/google/auto/value/extension/memoized/MemoizedTest.java
@@ -20,7 +20,9 @@ import static org.junit.Assert.fail;
 
 import com.google.auto.value.AutoValue;
 import com.google.common.collect.ImmutableList;
-import javax.annotation.Nullable;
+import java.lang.reflect.AnnotatedType;
+import java.lang.reflect.Constructor;
+import java.lang.reflect.Method;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -55,11 +57,16 @@ public class MemoizedTest {
     private int notNullableCount;
     private int nullableCount;
     private int returnsNullCount;
+    private int nullableWithTypeAnnotationCount;
+    private int returnsNullWithTypeAnnotationCount;
     private int notNullableButReturnsNullCount;
     private int throwsExceptionCount;
 
+    @javax.annotation.Nullable
     abstract String string();
 
+    abstract @org.checkerframework.checker.nullness.qual.Nullable String stringWithTypeAnnotation();
+
     abstract HashCodeAndToStringCounter counter();
 
     @Memoized
@@ -74,20 +81,35 @@ public class MemoizedTest {
     }
 
     @Memoized
-    @Nullable
+    @javax.annotation.Nullable
     String nullable() {
       nullableCount++;
       return "nullable derived " + string() + " " + nullableCount;
     }
 
     @Memoized
-    @Nullable
+    @javax.annotation.Nullable
     String returnsNull() {
       returnsNullCount++;
       return null;
     }
 
     @Memoized
+    @org.checkerframework.checker.nullness.qual.Nullable
+    String nullableWithTypeAnnotation() {
+      nullableWithTypeAnnotationCount++;
+      return "nullable derived " + stringWithTypeAnnotation() + " "
+          + nullableWithTypeAnnotationCount;
+    }
+
+    @Memoized
+    @org.checkerframework.checker.nullness.qual.Nullable
+    String returnsNullWithTypeAnnotation() {
+      returnsNullWithTypeAnnotationCount++;
+      return null;
+    }
+
+    @Memoized
     String notNullableButReturnsNull() {
       notNullableButReturnsNullCount++;
       return null;
@@ -140,7 +162,8 @@ public class MemoizedTest {
 
   @Before
   public void setUp() {
-    value = new AutoValue_MemoizedTest_Value("string", new HashCodeAndToStringCounter());
+    value = new AutoValue_MemoizedTest_Value(
+        "string", "stringWithTypeAnnotation", new HashCodeAndToStringCounter());
     listValue = new AutoValue_MemoizedTest_ListValue<Integer, String>(0, "hello");
   }
 
@@ -176,6 +199,14 @@ public class MemoizedTest {
   }
 
   @Test
+  public void nullableWithTypeAnnotation() {
+    assertThat(value.nullableWithTypeAnnotation())
+        .isEqualTo("nullable derived stringWithTypeAnnotation 1");
+    assertThat(value.nullableWithTypeAnnotation()).isSameAs(value.nullableWithTypeAnnotation());
+    assertThat(value.nullableWithTypeAnnotationCount).isEqualTo(1);
+  }
+
+  @Test
   public void returnsNull() {
     assertThat(value.returnsNull()).isNull();
     assertThat(value.returnsNull()).isNull();
@@ -183,6 +214,13 @@ public class MemoizedTest {
   }
 
   @Test
+  public void returnsNullWithTypeAnnotation() {
+    assertThat(value.returnsNullWithTypeAnnotation()).isNull();
+    assertThat(value.returnsNullWithTypeAnnotation()).isNull();
+    assertThat(value.returnsNullWithTypeAnnotationCount).isEqualTo(1);
+  }
+
+  @Test
   public void notNullableButReturnsNull() {
     try {
       value.notNullableButReturnsNull();
@@ -231,4 +269,35 @@ public class MemoizedTest {
     assertThat(value.getNative0()).isFalse();
     assertThat(value.getMemoizedNative0()).isFalse();
   }
+
+  @Test
+  public void nullableHasAnnotation() throws ReflectiveOperationException {
+    Method nullable = AutoValue_MemoizedTest_Value.class.getDeclaredMethod("nullable");
+    assertThat(nullable.isAnnotationPresent(javax.annotation.Nullable.class)).isTrue();
+  }
+
+  @Test
+  public void nullableWithTypeAnnotationHasAnnotation() throws ReflectiveOperationException {
+    Method nullable =
+        AutoValue_MemoizedTest_Value.class.getDeclaredMethod("nullableWithTypeAnnotation");
+    AnnotatedType returnType = nullable.getAnnotatedReturnType();
+    assertThat(returnType.isAnnotationPresent(
+                   org.checkerframework.checker.nullness.qual.Nullable.class))
+        .isTrue();
+  }
+
+  @Test
+  public void nullableConstructorParameter() throws ReflectiveOperationException {
+    // Constructor parameters are potentially:
+    // [0] @javax.annotation.Nullable String string,
+    // [1] @org.checkerframework.checker.nullness.qual.Nullable String stringWithTypeAnnotation,
+    // [2] HashCodeAndToStringCounter counter
+    // We don't currently copy @javax.annotation.Nullable because it is not a TYPE_USE annotation.
+    Constructor<?> constructor = AutoValue_MemoizedTest_Value.class.getDeclaredConstructor(
+        String.class, String.class, HashCodeAndToStringCounter.class);
+    AnnotatedType paramType = constructor.getAnnotatedParameterTypes()[1];
+    assertThat(paramType.isAnnotationPresent(
+                   org.checkerframework.checker.nullness.qual.Nullable.class))
+        .isTrue();
+  }
 }