From cd631709b042199e57cca618b09442df1a70799f Mon Sep 17 00:00:00 2001
From: Josiah Noel <32279667+SentryMan@users.noreply.github.com>
Date: Fri, 21 Jun 2024 16:51:37 -0400
Subject: [PATCH] Support Alternate Loading of Generated SPIs (#25)

* support loading generated SPIs

* Update pom.xml

* add validator generated
---
 .../avaje/spi/internal/ServiceProcessor.java  | 81 ++++++++++++-------
 pom.xml                                       |  2 +-
 2 files changed, 53 insertions(+), 30 deletions(-)

diff --git a/avaje-spi-core/src/main/java/io/avaje/spi/internal/ServiceProcessor.java b/avaje-spi-core/src/main/java/io/avaje/spi/internal/ServiceProcessor.java
index 113675c..7207377 100644
--- a/avaje-spi-core/src/main/java/io/avaje/spi/internal/ServiceProcessor.java
+++ b/avaje-spi-core/src/main/java/io/avaje/spi/internal/ServiceProcessor.java
@@ -2,7 +2,6 @@
 
 import static io.avaje.spi.internal.APContext.*;
 import static java.util.stream.Collectors.joining;
-import static java.util.stream.Collectors.toMap;
 
 import java.io.BufferedReader;
 import java.io.FileNotFoundException;
@@ -18,6 +17,8 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -62,6 +63,7 @@
   "io.avaje.http.api.Controller",
   "io.avaje.recordbuilder.Generated",
   "io.avaje.prism.GenerateAPContext",
+  "io.avaje.validation.spi.Generated",
   "javax.annotation.processing.Generated",
   "javax.annotation.processing.SupportedAnnotationTypes",
   "javax.annotation.processing.SupportedOptions",
@@ -83,44 +85,70 @@ public SourceVersion getSupportedSourceVersion() {
 
   private ModuleElement moduleElement;
 
-  private boolean writtenLocator;
-
   private Path servicesDirectory;
 
+  private Path generatedSpisDir;
+
   @Override
   public synchronized void init(ProcessingEnvironment env) {
     super.init(env);
     this.elements = env.getElementUtils();
     this.types = env.getTypeUtils();
     APContext.init(env);
+
+    final var filer = env.getFiler();
+    try {
+      final var uri =
+          filer
+              .createResource(
+                  StandardLocation.CLASS_OUTPUT, "", "META-INF/services/spi-service-locator")
+              .toUri();
+      this.servicesDirectory = Path.of(uri).getParent();
+      this.generatedSpisDir =
+          Path.of(
+              URI.create(
+                  uri.toString()
+                      .replace(
+                          "META-INF/services/spi-service-locator", "META-INF/generated-services")));
+    } catch (IOException e) {
+      logError("Failed to write service locator file");
+    }
   }
 
   @Override
   public boolean process(Set<? extends TypeElement> tes, RoundEnvironment roundEnv) {
 
-    if (!writtenLocator) {
-      try {
-        this.servicesDirectory =
-            Path.of(
-                    processingEnv
-                        .getFiler()
-                        .createResource(
-                            StandardLocation.CLASS_OUTPUT,
-                            "",
-                            "META-INF/services/spi-service-locator")
-                        .toUri())
-                .getParent();
-      } catch (IOException e) {
-        logError("Failed to write service locator file");
-      }
-      writtenLocator = true;
-    }
     final var annotated =
         Optional.ofNullable(typeElement(ServiceProviderPrism.PRISM_TYPE))
             .map(roundEnv::getElementsAnnotatedWith)
             .orElseGet(Set::of);
 
     // discover services from the current compilation sources
+    processSpis(annotated);
+
+    findModule(tes, roundEnv);
+    if (roundEnv.processingOver()) {
+      //load generated service files into main services
+      var generatedSpis = loadMetaInfServices(generatedSpisDir);
+      generatedSpis.forEach(
+          (key, value) ->
+              services.merge(
+                  key,
+                  value,
+                  (oldValue, newValue) -> {
+                    if (oldValue == null) {
+                      oldValue = new HashSet<>();
+                    }
+                    oldValue.addAll(newValue);
+                    return oldValue;
+                  }));
+      write();
+      validateModule();
+    }
+    return false;
+  }
+
+  private void processSpis(final Collection<? extends Element> annotated) {
     for (final var type : ElementFilter.typesIn(annotated)) {
       validate(type);
 
@@ -135,12 +163,6 @@ public boolean process(Set<? extends TypeElement> tes, RoundEnvironment roundEnv
         v.add(elements.getBinaryName(type).toString());
       }
     }
-    findModule(tes, roundEnv);
-    if (roundEnv.processingOver()) {
-      write();
-      validateModule();
-    }
-    return false;
   }
 
   private void validate(final TypeElement type) {
@@ -163,7 +185,8 @@ private void validate(final TypeElement type) {
 
   private void write() {
     // Read the existing service files
-    var allServices = loadMetaInfServices();
+    var allServices = loadMetaInfServices(servicesDirectory);
+    allServices.putAll(services);
     // Write the service files
     for (final Map.Entry<String, Set<String>> e : allServices.entrySet()) {
       final String contract = e.getKey();
@@ -190,8 +213,8 @@ private void write() {
     services.putAll(allServices);
   }
 
-  private Map<String, Set<String>> loadMetaInfServices() {
-    var allServices = new ConcurrentHashMap<>(services);
+  private Map<String, Set<String>> loadMetaInfServices(Path servicesDirectory) {
+    var allServices = new HashMap<String, Set<String>>();
 
     // Read the existing service files
     try (var servicePaths = Files.walk(servicesDirectory, 1).skip(1)) {
diff --git a/pom.xml b/pom.xml
index 5d99ed5..e087157 100644
--- a/pom.xml
+++ b/pom.xml
@@ -3,7 +3,7 @@
 	<parent>
 		<groupId>org.avaje</groupId>
 		<artifactId>java11-oss</artifactId>
-		<version>4.0</version>
+		<version>4.3</version>
 	</parent>
 	<groupId>io.avaje</groupId>
 	<artifactId>avaje-spi-parent</artifactId>