Check for strict deps exemptions at the top level

--
MOS_MIGRATED_REVID=138223170

1 files changed, 55 insertions, 55 deletions

diff --git a/src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/dependency/StrictJavaDepsPlugin.java b/src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/dependency/StrictJavaDepsPlugin.java
index 5e0b354284..1c08f2b3b3 100644
--- a/src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/dependency/StrictJavaDepsPlugin.java
+++ b/src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/dependency/StrictJavaDepsPlugin.java
@@ -14,6 +14,7 @@
 
 package com.google.devtools.build.buildjar.javac.plugins.dependency;
 
+import static com.google.common.collect.Iterables.getOnlyElement;
 import static com.google.devtools.build.buildjar.javac.plugins.dependency.DependencyModule.StrictJavaDeps.ERROR;
 import static com.google.devtools.build.buildjar.javac.plugins.dependency.ImplicitDependencyExtractor.getPlatformJars;
 
@@ -32,6 +33,7 @@ import com.sun.tools.javac.comp.AttrContext;
 import com.sun.tools.javac.comp.Env;
 import com.sun.tools.javac.main.JavaCompiler;
 import com.sun.tools.javac.tree.JCTree;
+import com.sun.tools.javac.tree.TreeInfo;
 import com.sun.tools.javac.tree.TreeScanner;
 import com.sun.tools.javac.util.Context;
 import com.sun.tools.javac.util.Log;
@@ -136,23 +138,29 @@ public final class StrictJavaDepsPlugin extends BlazeJavaCompilerPlugin {
    */
   @Override
   public void postAttribute(Env<AttrContext> env) {
-    // We want to generate warnings/errors as if we were javac, and in order to
-    // use the internal log properly, we need to set its current source file
-    // information. The useSource call does just that, and is a common pattern
-    // from JavaCompiler: set source to current file and save the previous
-    // value, do work and generate warnings, reset source.
-    JavaFileObject prev = log.useSource(
-        env.enclClass.sym.sourcefile != null
-            ? env.enclClass.sym.sourcefile
-            : env.toplevel.sourcefile);
-    if (trees.add(env.tree)) {
-      checkingTreeScanner.scan(env.tree);
-    }
-    if (toplevels.add(env.toplevel)) {
-      checkingTreeScanner.scan(env.toplevel.getImports());
-      dependencyModule.addPackage(env.toplevel.packge.toString());
+    JavaFileObject previousSource =
+        log.useSource(
+            env.enclClass.sym.sourcefile != null
+                ? env.enclClass.sym.sourcefile
+                : env.toplevel.sourcefile);
+    boolean previousExemption = checkingTreeScanner.isStrictDepsExempt;
+    try {
+      ProcessorDependencyMode mode = isAnnotationProcessorExempt(env.toplevel);
+      if (mode == ProcessorDependencyMode.EXEMPT_NORECORD) {
+        return;
+      }
+      checkingTreeScanner.isStrictDepsExempt |= mode == ProcessorDependencyMode.EXEMPT_RECORD;
+      if (trees.add(env.tree)) {
+        checkingTreeScanner.scan(env.tree);
+      }
+      if (toplevels.add(env.toplevel)) {
+        checkingTreeScanner.scan(env.toplevel.getImports());
+        dependencyModule.addPackage(env.toplevel.packge.toString());
+      }
+    } finally {
+      checkingTreeScanner.isStrictDepsExempt = previousExemption;
+      log.useSource(previousSource);
     }
-    log.useSource(prev);
   }
 
   @Override
@@ -214,9 +222,6 @@ public final class StrictJavaDepsPlugin extends BlazeJavaCompilerPlugin {
     /** The set of jars on the compilation bootclasspath. */
     private final Set<String> platformJars;
 
-    /** The set of generators we exempt from this testing. */
-    private final Set<String> exemptGenerators;
-
     /** Was the node being visited generated by an exempt annotation processor? */
     private boolean isStrictDepsExempt = false;
 
@@ -233,7 +238,6 @@ public final class StrictJavaDepsPlugin extends BlazeJavaCompilerPlugin {
       this.directDependenciesMap = dependencyModule.getExplicitDependenciesMap();
       this.platformJars = platformJars;
       this.fileManager = fileManager;
-      this.exemptGenerators = dependencyModule.getExemptGenerators();
     }
 
     Set<ClassSymbol> getSeenClasses() {
@@ -355,49 +359,45 @@ public final class StrictJavaDepsPlugin extends BlazeJavaCompilerPlugin {
       }
       scan(tree.body);
     }
+  }
 
-    private static final String DAGGER_PROCESSOR_PREFIX = "dagger.";
+  private static final String DAGGER_PROCESSOR_PREFIX = "dagger.";
 
-    enum ProcessorDependencyMode {
-      DEFAULT,
-      EXEMPT_RECORD,
-      EXEMPT_NORECORD;
-    }
+  enum ProcessorDependencyMode {
+    DEFAULT,
+    EXEMPT_RECORD,
+    EXEMPT_NORECORD;
+  }
 
-    public ProcessorDependencyMode isAnnotationProcessorExempt(JCTree.JCClassDecl tree) {
-      if (tree.sym == null) {
-        return ProcessorDependencyMode.DEFAULT;
-      }
-      Generated generated = tree.sym.getAnnotation(Generated.class);
-      if (generated == null) {
-        return ProcessorDependencyMode.DEFAULT;
-      }
-      for (String value : generated.value()) {
-        // Relax strict deps for dagger-generated code (b/17979436).
-        if (value.startsWith(DAGGER_PROCESSOR_PREFIX)) {
-          return ProcessorDependencyMode.EXEMPT_NORECORD;
-        }
-        if (exemptGenerators.contains(value)) {
-          return ProcessorDependencyMode.EXEMPT_RECORD;
-        }
-      }
+  /**
+   * Returns true if the compilation unit contains a single top-level class generated by an exempt
+   * annotation processor (according to its {@link @Generated} annotation).
+   *
+   * <p>Annotation processors are expected to never generate more than one top level class, as
+   * required by the style guide.
+   */
+  public ProcessorDependencyMode isAnnotationProcessorExempt(JCTree.JCCompilationUnit unit) {
+    if (unit.getTypeDecls().size() != 1) {
       return ProcessorDependencyMode.DEFAULT;
     }
-
-    @Override
-    public void visitClassDef(JCTree.JCClassDecl tree) {
-      ProcessorDependencyMode mode = isAnnotationProcessorExempt(tree);
-      if (mode == ProcessorDependencyMode.EXEMPT_NORECORD) {
-        return;
+    Symbol sym = TreeInfo.symbolFor(getOnlyElement(unit.getTypeDecls()));
+    if (sym == null) {
+      return ProcessorDependencyMode.DEFAULT;
+    }
+    Generated generated = sym.getAnnotation(Generated.class);
+    if (generated == null) {
+      return ProcessorDependencyMode.DEFAULT;
+    }
+    for (String value : generated.value()) {
+      // Relax strict deps for dagger-generated code (b/17979436).
+      if (value.startsWith(DAGGER_PROCESSOR_PREFIX)) {
+        return ProcessorDependencyMode.EXEMPT_NORECORD;
       }
-      boolean previous = isStrictDepsExempt;
-      try {
-        isStrictDepsExempt |= mode == ProcessorDependencyMode.EXEMPT_RECORD;
-        super.visitClassDef(tree);
-      } finally {
-        isStrictDepsExempt = previous;
+      if (dependencyModule.getExemptGenerators().contains(value)) {
+        return ProcessorDependencyMode.EXEMPT_RECORD;
       }
     }
+    return ProcessorDependencyMode.DEFAULT;
   }
 
   /**