1 files changed, 92 insertions, 39 deletions

diff --git a/src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/processing/AnnotationProcessingModule.java b/src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/processing/AnnotationProcessingModule.java
index e9d4c1832a..65a831f1e5 100644
--- a/src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/processing/AnnotationProcessingModule.java
+++ b/src/java_tools/buildjar/java/com/google/devtools/build/buildjar/javac/plugins/processing/AnnotationProcessingModule.java
@@ -15,50 +15,108 @@
 package com.google.devtools.build.buildjar.javac.plugins.processing;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 import com.google.devtools.build.buildjar.javac.plugins.BlazeJavaCompilerPlugin;
+import com.google.devtools.build.buildjar.proto.JavaCompilation.CompilationUnit;
+import com.google.devtools.build.buildjar.proto.JavaCompilation.Manifest;
 
+import java.io.IOException;
+import java.io.OutputStream;
+import java.nio.file.Files;
 import java.nio.file.Path;
-import java.util.HashSet;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * A module for information about the compilation's annotation processing.
  */
 public class AnnotationProcessingModule {
-  
+
   /**
    * A builder for {@link AnnotationProcessingModule}s.
    */
   public static class Builder {
     private Path sourceGenDir;
-    private Path classDir;
+    private Path manifestProto;
+    private ImmutableSet.Builder<Path> sourceRoots = ImmutableSet.builder();
 
     private Builder() {}
 
     public AnnotationProcessingModule build() {
-      return new AnnotationProcessingModule(sourceGenDir, classDir);
+      return new AnnotationProcessingModule(
+          sourceGenDir, manifestProto, validateSourceRoots(sourceRoots.build()));
+    }
+
+    /**
+     * Verify that source roots do not contain other source roots.
+     *
+     * <p>If one source root is an ancestor of another, the source path to
+     * use in the manifest will be ambiguous.
+     */
+    private ImmutableSet<Path> validateSourceRoots(ImmutableSet<Path> roots) {
+      // It's sad that this is quadratic, but the number of source roots
+      // should be <= 2.
+      for (Path a : roots) {
+        for (Path b : roots) {
+          if (a.equals(b) || b.getNameCount() == 0) {
+            continue;
+          }
+          if (a.startsWith(b)) {
+            throw new IllegalArgumentException(
+                String.format("Source root %s is a parent of %s", b, a));
+          }
+        }
+      }
+      return roots;
     }
 
     public void setSourceGenDir(Path sourceGenDir) {
-      this.sourceGenDir = sourceGenDir.toAbsolutePath();
+      this.sourceGenDir = sourceGenDir;
+    }
+
+    public void setManifestProtoPath(Path manifestProto) {
+      this.manifestProto = manifestProto.toAbsolutePath();
     }
 
-    public void setClassDir(Path classDir) {
-      this.classDir = classDir.toAbsolutePath();
+    public void addAllSourceRoots(Set<String> sourceRoots) {
+      for (String root : sourceRoots) {
+        this.sourceRoots.add(Paths.get(root));
+      }
     }
   }
 
   private final boolean enabled;
   private final Path sourceGenDir;
-  private final Path classDir;
+  private final Path manifestProto;
+  private final ImmutableSet<Path> sourceRoots;
+
+  public boolean isGenerated(Path path) {
+    return path.startsWith(sourceGenDir);
+  }
 
-  public Path sourceGenDir() {
-    return sourceGenDir;
+  public Path stripSourceRoot(Path path) {
+    if (path.startsWith(sourceGenDir)) {
+      return sourceGenDir.relativize(path);
+    }
+    for (Path sourceRoot : sourceRoots) {
+      if (path.startsWith(sourceRoot)) {
+        return sourceRoot.relativize(path);
+      }
+    }
+    return path;
   }
 
-  private AnnotationProcessingModule(Path sourceGenDir, Path classDir) {
+  private AnnotationProcessingModule(
+      Path sourceGenDir, Path manifestProto, ImmutableSet<Path> sourceRoots) {
     this.sourceGenDir = sourceGenDir;
-    this.classDir = classDir;
-    this.enabled = sourceGenDir != null && classDir != null;
+    this.manifestProto = manifestProto;
+    this.sourceRoots = sourceRoots;
+    this.enabled = sourceGenDir != null && manifestProto != null;
   }
 
   public static Builder builder() {
@@ -71,38 +129,33 @@ public class AnnotationProcessingModule {
     }
   }
 
-  /**
-   * The set of prefixes of generated class files.
-   */
-  private final HashSet<String> pathPrefixes = new HashSet<>();
+  private final Map<String, CompilationUnit> units = new HashMap<>();
 
-  /**
-   * Record the prefix of a group of generated class files.
-   *
-   * <p>Prefixes are used to handle generated inner classes. Since
-   * e.g. j/c/g/Foo.class and j/c/g/Foo$Inner.class both correspond to the
-   * same generated source file, only the prefix "j/c/g/Foo" is recorded.
-   */
-  void recordPrefix(String pathPrefix) {
-    pathPrefixes.add(pathPrefix);
+  public void recordUnit(CompilationUnit unit) {
+    units.put(unit.getPath(), unit);
   }
 
-  /** Returns true if the given path is to a generated source file. */
-  public boolean isGeneratedSource(Path sourcePath) {
-    return sourcePath.toAbsolutePath().startsWith(sourceGenDir);
+  private Manifest buildManifestProto() {
+    Manifest.Builder builder = Manifest.newBuilder();
+
+    List<String> keys = new ArrayList<>(units.keySet());
+    Collections.sort(keys);
+    for (String key : keys) {
+      CompilationUnit unit = units.get(key);
+      builder.addCompilationUnit(unit);
+    }
+
+    return builder.build();
   }
 
-  /** Returns true if the given path is to a generated class file. */
-  public boolean isGeneratedClass(Path path) {
-    if (!path.getFileName().toString().endsWith(".class")) {
-      return false;
+  public void emitManifestProto() throws IOException {
+    if (!enabled) {
+      return;
     }
-    String prefix = classDir.relativize(path.toAbsolutePath()).toString();
-    prefix = prefix.substring(0, prefix.length() - ".class".length());
-    int idx = prefix.lastIndexOf('$');
-    if (idx > 0) {
-      prefix = prefix.substring(0, idx);
+    try (OutputStream out = Files.newOutputStream(manifestProto)) {
+      buildManifestProto().writeTo(out);
+    } catch (IOException ex) {
+      throw new IOException("Cannot write manifest to " + manifestProto, ex);
     }
-    return pathPrefixes.contains(prefix);
   }
 }