Add support for invoking AOT in the Gradle plugin

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/dsl/SpringBootExtension.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2021 the original author or authors.
+ * Copyright 2012-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,18 +17,26 @@
 package org.springframework.boot.gradle.dsl;
 
 import java.io.File;
+import java.util.List;
 
 import org.gradle.api.Action;
 import org.gradle.api.Project;
+import org.gradle.api.artifacts.Configuration;
+import org.gradle.api.artifacts.ConfigurationContainer;
+import org.gradle.api.file.FileCollection;
 import org.gradle.api.plugins.BasePlugin;
 import org.gradle.api.plugins.JavaPlugin;
 import org.gradle.api.plugins.JavaPluginExtension;
 import org.gradle.api.provider.Property;
 import org.gradle.api.tasks.SourceSet;
+import org.gradle.api.tasks.SourceSetContainer;
 import org.gradle.api.tasks.TaskContainer;
 import org.gradle.api.tasks.TaskProvider;
 import org.gradle.jvm.tasks.Jar;
 
+import org.springframework.boot.gradle.plugin.ResolveMainClassName;
+import org.springframework.boot.gradle.plugin.SpringBootPlugin;
+import org.springframework.boot.gradle.tasks.aot.GenerateAotSources;
 import org.springframework.boot.gradle.tasks.buildinfo.BuildInfo;
 import org.springframework.boot.gradle.tasks.buildinfo.BuildInfoProperties;
 
@@ -129,4 +137,39 @@ private Jar findArtifactTask() {
 		return (Jar) this.project.getTasks().findByName("bootJar");
 	}
 
+	public void aot() {
+		this.project.getPlugins().withType(JavaPlugin.class).all((plugin) -> {
+			JavaPluginExtension javaPluginExtension = this.project.getExtensions().getByType(JavaPluginExtension.class);
+			SourceSetContainer sourceSets = javaPluginExtension.getSourceSets();
+			SourceSet main = sourceSets.getByName(SourceSet.MAIN_SOURCE_SET_NAME);
+			FileCollection runtimeClasspath = main.getRuntimeClasspath();
+			TaskProvider<ResolveMainClassName> resolveMainClassName = this.project.getTasks()
+					.named(SpringBootPlugin.RESOLVE_MAIN_CLASS_NAME_TASK_NAME, ResolveMainClassName.class);
+			SourceSet aotSourceSet = sourceSets.create("aot", (aot) -> {
+				aot.getJava().setSrcDirs(List.of("build/generated/aotSources"));
+				aot.getResources().setSrcDirs(List.of("build/generated/aotResources"));
+				aot.setCompileClasspath(aot.getCompileClasspath().plus(main.getOutput()));
+				main.setRuntimeClasspath(runtimeClasspath.plus(aot.getOutput()));
+				ConfigurationContainer configurations = this.project.getConfigurations();
+				Configuration aotImplementation = configurations.getByName(aot.getImplementationConfigurationName());
+				aotImplementation.extendsFrom(configurations.getByName(main.getImplementationConfigurationName()));
+				aotImplementation.extendsFrom(configurations.getByName(main.getRuntimeOnlyConfigurationName()));
+			});
+			TaskProvider<GenerateAotSources> generateAotSources = this.project.getTasks().register("generateAotSources",
+					GenerateAotSources.class, (task) -> {
+						task.getApplicationClass()
+								.set(resolveMainClassName.flatMap((thing) -> thing.readMainClassName()));
+						task.setClasspath(aotSourceSet.getCompileClasspath());
+						task.getSourcesDir().set(aotSourceSet.getJava().getSrcDirs().iterator().next());
+						task.getResourcesDir().set(aotSourceSet.getResources().getSrcDirs().iterator().next());
+					});
+			this.project.getTasks().getByName(aotSourceSet.getCompileJavaTaskName(), (compile) -> {
+				compile.dependsOn(generateAotSources);
+			});
+			this.project.getTasks().getByName(aotSourceSet.getProcessResourcesTaskName(), (processResources) -> {
+				processResources.dependsOn(generateAotSources);
+			});
+		});
+	}
+
 }

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/plugin/ResolveMainClassName.java

@@ -144,7 +144,7 @@ private String findMainClass(File file) {
 		}
 	}
 
-	Provider<String> readMainClassName() {
+	public Provider<String> readMainClassName() {
 		return this.outputFile.map(new ClassNameReader());
 	}
 

spring-boot-project/spring-boot-tools/spring-boot-gradle-plugin/src/main/java/org/springframework/boot/gradle/tasks/aot/GenerateAotSources.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2012-2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.gradle.tasks.aot;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.gradle.api.file.DirectoryProperty;
+import org.gradle.api.provider.Property;
+import org.gradle.api.tasks.Input;
+import org.gradle.api.tasks.JavaExec;
+import org.gradle.api.tasks.OutputDirectory;
+import org.gradle.api.tasks.TaskAction;
+
+/**
+ * Custom {@link JavaExec} task for generating sources ahead of time.
+ *
+ * @author Andy Wilkinson
+ * @since 3.0
+ */
+public class GenerateAotSources extends JavaExec {
+
+	private final Property<String> applicationClass;
+
+	private final DirectoryProperty sourcesDir;
+
+	private final DirectoryProperty resourcesDir;
+
+	public GenerateAotSources() {
+		this.applicationClass = getProject().getObjects().property(String.class);
+		this.sourcesDir = getProject().getObjects().directoryProperty();
+		this.resourcesDir = getProject().getObjects().directoryProperty();
+		getMainClass().set("org.springframework.boot.AotProcessor");
+	}
+
+	@Input
+	public Property<String> getApplicationClass() {
+		return this.applicationClass;
+	}
+
+	@OutputDirectory
+	public DirectoryProperty getSourcesDir() {
+		return this.sourcesDir;
+	}
+
+	@OutputDirectory
+	public DirectoryProperty getResourcesDir() {
+		return this.resourcesDir;
+	}
+
+	@Override
+	@TaskAction
+	public void exec() {
+		List<String> args = new ArrayList<>();
+		args.add(this.applicationClass.get());
+		args.add(this.sourcesDir.getAsFile().get().getAbsolutePath());
+		args.add(this.resourcesDir.getAsFile().get().getAbsolutePath());
+		args.addAll(super.getArgs());
+		this.setArgs(args);
+		super.exec();
+	}
+
+}