Optimise proto file resolution in archives to only need a single pass

src/main/java/io/github/ascopes/protobufmavenplugin/generate/SourceCodeGenerator.java

@@ -23,7 +23,7 @@
 import io.github.ascopes.protobufmavenplugin.execute.ArgLineBuilder;
 import io.github.ascopes.protobufmavenplugin.execute.CommandLineExecutor;
 import io.github.ascopes.protobufmavenplugin.source.ProtoFileListing;
-import io.github.ascopes.protobufmavenplugin.source.ProtoListingResolver;
+import io.github.ascopes.protobufmavenplugin.source.ProtoSourceResolver;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -50,15 +50,15 @@ public final class SourceCodeGenerator {
   private final MavenDependencyPathResolver mavenDependencyPathResolver;
   private final ProtocResolver protocResolver;
   private final PluginResolver pluginResolver;
-  private final ProtoListingResolver protoListingResolver;
+  private final ProtoSourceResolver protoListingResolver;
   private final CommandLineExecutor commandLineExecutor;
 
   @Inject
   public SourceCodeGenerator(
       MavenDependencyPathResolver mavenDependencyPathResolver,
       ProtocResolver protocResolver,
       PluginResolver pluginResolver,
-      ProtoListingResolver protoListingResolver,
+      ProtoSourceResolver protoListingResolver,
       CommandLineExecutor commandLineExecutor
   ) {
     this.mavenDependencyPathResolver = mavenDependencyPathResolver;

src/main/java/io/github/ascopes/protobufmavenplugin/source/ProtoArchiveExtractor.java

@@ -15,21 +15,19 @@
  */
 package io.github.ascopes.protobufmavenplugin.source;
 
-import static io.github.ascopes.protobufmavenplugin.source.ProtoFilePredicates.isProtoFile;
-
 import io.github.ascopes.protobufmavenplugin.system.FileUtils;
 import java.io.IOException;
-import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import java.nio.file.SimpleFileVisitor;
-import java.nio.file.attribute.BasicFileAttributes;
 import java.nio.file.spi.FileSystemProvider;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
+import java.util.ArrayList;
 import java.util.Base64;
+import java.util.Collection;
 import java.util.Map;
 import java.util.Optional;
+import java.util.stream.Collectors;
 import javax.inject.Inject;
 import javax.inject.Named;
 import org.apache.maven.project.MavenProject;
@@ -47,7 +45,7 @@ public final class ProtoArchiveExtractor {
   private static final Logger log = LoggerFactory.getLogger(ProtoArchiveExtractor.class);
 
   private final FileSystemProvider jarFileSystemProvider;
-  private final Path extractionRoot;
+  private final Path extractionBaseDir;
 
   @Inject
   public ProtoArchiveExtractor(MavenProject mavenProject) {
@@ -57,38 +55,57 @@ public ProtoArchiveExtractor(MavenProject mavenProject) {
         .findFirst()
         .orElseThrow();
 
-    extractionRoot = Path.of(mavenProject.getBuild().getDirectory())
+    extractionBaseDir = Path.of(mavenProject.getBuild().getDirectory())
         .resolve("protobuf-maven-plugin")
         .resolve("extracted");
   }
 
-  public Optional<Path> tryExtractProtoFiles(Path zipPath) throws IOException {
-    var extractionTarget = extractionRoot.resolve(generateUniqueName(zipPath));
-
+  public Optional<ProtoFileListing> extractProtoFiles(Path zipPath) throws IOException {
     try (var vfs = jarFileSystemProvider.newFileSystem(zipPath, Map.of())) {
       var vfsRoot = vfs.getRootDirectories().iterator().next();
+      var sourceFiles = findProtoFilesInArchive(vfsRoot);
+
+      if (sourceFiles.isEmpty()) {
+        return Optional.empty();
+      }
+
+      var extractionRoot = extractionBaseDir.resolve(generateUniqueName(zipPath));
+      Files.createDirectories(extractionRoot);
+
+      var targetFiles = new ArrayList<Path>();
 
-      Files.walkFileTree(vfsRoot, new SimpleFileVisitor<Path>() {
-        @Override
-        public FileVisitResult visitFile(
-            Path sourceFile,
-            BasicFileAttributes attrs
-        ) throws IOException {
-          if (isProtoFile(sourceFile)) {
-            var targetFile = changeRelativePath(extractionTarget, vfsRoot, sourceFile);
-            log.trace("Extracting {} to {}", sourceFile.toUri(), targetFile);
-            Files.createDirectories(targetFile.getParent());
-            Files.copy(sourceFile, targetFile);
-          }
-
-          return FileVisitResult.CONTINUE;
-        }
-      });
+      for (var sourceFile : sourceFiles) {
+        var targetFile = changeRelativePath(extractionRoot, vfsRoot, sourceFile);
+        log.trace("Copying {} to {}", sourceFile.toUri(), targetFile);
+
+        // We have to do this on each iteration to ensure the directory hierarchy exists.
+        Files.createDirectories(targetFile.getParent());
+        Files.copy(sourceFile, targetFile);
+        targetFiles.add(targetFile);
+      }
+
+      var listing = ImmutableProtoFileListing
+          .builder()
+          .originalRoot(zipPath)
+          .protoFilesRoot(extractionRoot)
+          .protoFiles(targetFiles)
+          .build();
+
+      return Optional.of(listing);
     }
+  }
 
-    // Will this create odd behaviour on reruns if archives change? Do we care? If it becomes
-    // a problem, we can use a boolean flag instead.
-    return Optional.of(extractionTarget).filter(Files::isDirectory);
+  private Collection<Path> findProtoFilesInArchive(Path archiveRootPath) throws IOException {
+    try (var stream = Files.walk(archiveRootPath)) {
+      return stream
+          .filter(ProtoFilePredicates::isProtoFile)
+          .peek(protoFile -> log.trace(
+              "Found proto file {} in archive {}",
+              protoFile.toUri(),
+              archiveRootPath
+          ))
+          .collect(Collectors.toUnmodifiableList());
+    }
   }
 
   private String generateUniqueName(Path archive) {

src/main/java/io/github/ascopes/protobufmavenplugin/source/ProtoListingResolver.java → src/main/java/io/github/ascopes/protobufmavenplugin/source/ProtoSourceResolver.java

@@ -47,15 +47,15 @@
  * @author Ashley Scopes
  */
 @Named
-public final class ProtoListingResolver implements AutoCloseable {
+public final class ProtoSourceResolver implements AutoCloseable {
 
   private static final Logger log = LoggerFactory.getLogger(ProtoArchiveExtractor.class);
 
   private final ProtoArchiveExtractor protoArchiveExtractor;
   private final ExecutorService executorService;
 
   @Inject
-  public ProtoListingResolver(ProtoArchiveExtractor protoArchiveExtractor) {
+  public ProtoSourceResolver(ProtoArchiveExtractor protoArchiveExtractor) {
     var concurrency = Runtime.getRuntime().availableProcessors() * 4;
 
     this.protoArchiveExtractor = protoArchiveExtractor;
@@ -103,29 +103,15 @@ public Collection<ProtoFileListing> createProtoFileListings(
     return results;
   }
 
-  public Optional<ProtoFileListing> createProtoFileListing(Path originalPath) throws IOException {
-    Path rootPath;
-
-    if (Files.isRegularFile(originalPath)) {
-      // Treat it as a ZIP archive. We might want to support other formats in the future but
-      // for now, let's leave it alone. If we get no result from this call then we already know
-      // that there are no proto files to extract, so we can skip the lookup.
-
-      var extractionResult = protoArchiveExtractor.tryExtractProtoFiles(originalPath);
-
-      if (extractionResult.isPresent()) {
-        rootPath = extractionResult.get();
-      } else {
-        return Optional.empty();
-      }
-    } else {
-      rootPath = originalPath;
+  public Optional<ProtoFileListing> createProtoFileListing(Path path) throws IOException {
+    if (Files.isRegularFile(path)) {
+      return protoArchiveExtractor.extractProtoFiles(path);
     }
 
-    try (var stream = Files.walk(rootPath)) {
+    try (var stream = Files.walk(path)) {
       var protoFiles = stream
           .filter(ProtoFilePredicates::isProtoFile)
-          .peek(protoFile -> log.trace("Found proto file in root {}: {}", rootPath, protoFile))
+          .peek(protoFile -> log.trace("Found proto file in root {}: {}", path, protoFile))
           .collect(Collectors.toUnmodifiableSet());
 
       if (protoFiles.isEmpty()) {
@@ -135,8 +121,8 @@ public Optional<ProtoFileListing> createProtoFileListing(Path originalPath) thro
       var listing = ImmutableProtoFileListing
           .builder()
           .addAllProtoFiles(protoFiles)
-          .protoFilesRoot(rootPath)
-          .originalRoot(originalPath)
+          .protoFilesRoot(path)
+          .originalRoot(path)
           .build();
 
       return Optional.of(listing);