Fix bug where ProcessingStep.process(...) was called with too many elements when there had been deferred elements.

RELNOTES=Fix bug where `ProcessingStep.process(...)` was called with too many elements when there had been deferred elements.

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

Author: netdpb

common/src/main/java/com/google/auto/common/BasicAnnotationProcessor.java

@@ -321,7 +321,7 @@ private void process(ImmutableSetMultimap<Class<? extends Annotation>, Element>
     for (ProcessingStep step : steps) {
       ImmutableSetMultimap<Class<? extends Annotation>, Element> stepElements =
           new ImmutableSetMultimap.Builder<Class<? extends Annotation>, Element>()
-              .putAll(indexByAnnotation(elementsDeferredBySteps.get(step)))
+              .putAll(indexByAnnotation(elementsDeferredBySteps.get(step), step.annotations()))
               .putAll(filterKeys(validElements, Predicates.<Object>in(step.annotations())))
               .build();
       if (stepElements.isEmpty()) {
@@ -343,15 +343,14 @@ public ElementName apply(Element element) {
   }
 
   private ImmutableSetMultimap<Class<? extends Annotation>, Element> indexByAnnotation(
-      Set<ElementName> annotatedElements) {
-    ImmutableSet<? extends Class<? extends Annotation>> supportedAnnotationClasses =
-        getSupportedAnnotationClasses();
+      Set<ElementName> annotatedElements,
+      Set<? extends Class<? extends Annotation>> annotationClasses) {
     ImmutableSetMultimap.Builder<Class<? extends Annotation>, Element> deferredElements =
         ImmutableSetMultimap.builder();
     for (ElementName elementName : annotatedElements) {
       Optional<? extends Element> element = elementName.getElement(elements);
       if (element.isPresent()) {
-        findAnnotatedElements(element.get(), supportedAnnotationClasses, deferredElements);
+        findAnnotatedElements(element.get(), annotationClasses, deferredElements);
       }
     }
     return deferredElements.build();
@@ -377,7 +376,7 @@ private ImmutableSetMultimap<Class<? extends Annotation>, Element> indexByAnnota
    */
   private static void findAnnotatedElements(
       Element element,
-      ImmutableSet<? extends Class<? extends Annotation>> annotationClasses,
+      Set<? extends Class<? extends Annotation>> annotationClasses,
       ImmutableSetMultimap.Builder<Class<? extends Annotation>, Element> annotatedElements) {
     for (Element enclosedElement : element.getEnclosedElements()) {
       if (!enclosedElement.getKind().isClass() && !enclosedElement.getKind().isInterface()) {

common/src/test/java/com/google/auto/common/BasicAnnotationProcessorTest.java

@@ -15,23 +15,24 @@
  */
 package com.google.auto.common;
 
+import static com.google.common.collect.Multimaps.transformValues;
 import static com.google.common.truth.Truth.assertAbout;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.testing.compile.JavaSourceSubjectFactory.javaSource;
 import static com.google.testing.compile.JavaSourcesSubjectFactory.javaSources;
-import static java.lang.annotation.ElementType.TYPE;
 import static javax.tools.StandardLocation.SOURCE_OUTPUT;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.ImmutableSetMultimap;
 import com.google.common.collect.SetMultimap;
+import com.google.common.truth.Correspondence;
 import com.google.testing.compile.JavaFileObjects;
 import java.io.IOException;
 import java.io.PrintWriter;
 import java.lang.annotation.Annotation;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
-import java.lang.annotation.Target;
 import java.util.Set;
 import javax.annotation.processing.Filer;
 import javax.lang.model.SourceVersion;
@@ -51,9 +52,11 @@ public class BasicAnnotationProcessorTest {
   /**
    * Rejects elements unless the class generated by {@link GeneratesCode}'s processor is present.
    */
-  public class RequiresGeneratedCodeProcessor extends BasicAnnotationProcessor {
+  private static final class RequiresGeneratedCodeProcessor extends BasicAnnotationProcessor {
 
-    private int rejectedRounds;
+    int rejectedRounds;
+    final ImmutableList.Builder<ImmutableSetMultimap<Class<? extends Annotation>, Element>>
+        processArguments = ImmutableList.builder();
 
     @Override
     public SourceVersion getSupportedSourceVersion() {
@@ -67,6 +70,7 @@ protected Iterable<? extends ProcessingStep> initSteps() {
             @Override
             public Set<Element> process(
                 SetMultimap<Class<? extends Annotation>, Element> elementsByAnnotation) {
+              processArguments.add(ImmutableSetMultimap.copyOf(elementsByAnnotation));
               TypeElement requiredClass =
                   processingEnv.getElementUtils().getTypeElement("test.SomeGeneratedClass");
               if (requiredClass == null) {
@@ -81,8 +85,24 @@ public Set<Element> process(
             public Set<? extends Class<? extends Annotation>> annotations() {
               return ImmutableSet.of(RequiresGeneratedCode.class);
             }
+          },
+          new ProcessingStep() {
+            @Override
+            public Set<? extends Class<? extends Annotation>> annotations() {
+              return ImmutableSet.of(AnAnnotation.class);
+            }
+
+            @Override
+            public Set<? extends Element> process(
+                SetMultimap<Class<? extends Annotation>, Element> elementsByAnnotation) {
+              return ImmutableSet.of();
+            }
           });
     }
+
+    ImmutableList<ImmutableSetMultimap<Class<? extends Annotation>, Element>> processArguments() {
+      return processArguments.build();
+    }
   }
 
   @Retention(RetentionPolicy.SOURCE)
@@ -115,7 +135,6 @@ public Set<? extends Class<? extends Annotation>> annotations() {
     }
   }
 
-  @Target(TYPE)
   public @interface AnAnnotation {}
 
   /** When annotating a type {@code Foo}, generates a class called {@code FooXYZ}. */
@@ -250,11 +269,16 @@ public void properlyDefersProcessing_nestedTypeValidBeforeOuterType() {
 
   @Test
   public void properlyDefersProcessing_rejectsElement() {
-    JavaFileObject classAFileObject = JavaFileObjects.forSourceLines("test.ClassA",
-        "package test;",
-        "",
-        "@" + RequiresGeneratedCode.class.getCanonicalName(),
-        "public class ClassA {}");
+    JavaFileObject classAFileObject =
+        JavaFileObjects.forSourceLines(
+            "test.ClassA",
+            "package test;",
+            "",
+            "@" + RequiresGeneratedCode.class.getCanonicalName(),
+            "public class ClassA {",
+            "  @" + AnAnnotation.class.getCanonicalName(),
+            "  public void method() {}",
+            "}");
     JavaFileObject classBFileObject = JavaFileObjects.forSourceLines("test.ClassB",
         "package test;",
         "",
@@ -270,6 +294,31 @@ public void properlyDefersProcessing_rejectsElement() {
         .generatesFileNamed(
             SOURCE_OUTPUT, "test", "GeneratedByRequiresGeneratedCodeProcessor.java");
     assertThat(requiresGeneratedCodeProcessor.rejectedRounds).isEqualTo(1);
+
+    // Re b/118372780: Assert that the right deferred elements are passed back, and not any enclosed
+    // elements annotated with annotations from a different step.
+    assertThat(requiresGeneratedCodeProcessor.processArguments())
+        .comparingElementsUsing(setMultimapValuesByString())
+        .containsExactly(
+            ImmutableSetMultimap.of(RequiresGeneratedCode.class, "test.ClassA"),
+            ImmutableSetMultimap.of(RequiresGeneratedCode.class, "test.ClassA"))
+        .inOrder();
+  }
+
+  private static <K, V>
+      Correspondence<SetMultimap<K, V>, SetMultimap<K, String>> setMultimapValuesByString() {
+    return new Correspondence<SetMultimap<K, V>, SetMultimap<K, String>>() {
+      @Override
+      public boolean compare(SetMultimap<K, V> actual, SetMultimap<K, String> expected) {
+        return ImmutableSetMultimap.copyOf(transformValues(actual, Object::toString))
+            .equals(expected);
+      }
+
+      @Override
+      public String toString() {
+        return "is equivalent comparing multimap values by `toString()` to";
+      }
+    };
   }
 
   @Test public void reportsMissingType() {