Add ability to cache ASTs in-memory when using the Gradle daemon

plugin/src/main/java/org/openrewrite/gradle/AbstractRewriteTask.java

@@ -49,6 +49,8 @@ public abstract class AbstractRewriteTask extends DefaultTask implements Rewrite
     private List<Project> projects;
     private RewriteExtension extension;
     private RewriteReflectiveFacade rewrite;
+    private static final Map<File, byte[]> astCache = new HashMap<>();
+    protected boolean useAstCache;
 
     AbstractRewriteTask setConfiguration(Configuration configuration) {
         this.configuration = configuration;
@@ -146,52 +148,64 @@ protected InMemoryExecutionContext executionContext() {
     }
 
     protected ResultsContainer listResults() {
-        try {
-            Path baseDir = getProject().getRootProject().getRootDir().toPath();
-            Environment env = environment();
-            Set<String> activeRecipes = getActiveRecipes();
-            Set<String> activeStyles = getActiveStyles();
-            getLog().lifecycle(String.format("Using active recipe(s) %s", activeRecipes));
-            getLog().lifecycle(String.format("Using active styles(s) %s", activeStyles));
-            if (activeRecipes.isEmpty()) {
-                return new ResultsContainer(baseDir, emptyList());
-            }
-            List<NamedStyles> styles = env.activateStyles(activeStyles);
-            File checkstyleConfig = extension.getCheckstyleConfigFile();
-            if (checkstyleConfig != null && checkstyleConfig.exists()) {
-                NamedStyles checkstyle = getRewrite().loadCheckstyleConfig(checkstyleConfig.toPath(), extension.getCheckstyleProperties());
-                styles.add(checkstyle);
-            }
+        Path baseDir = getProject().getRootProject().getRootDir().toPath();
+        Environment env = environment();
+        Set<String> activeRecipes = getActiveRecipes();
+        Set<String> activeStyles = getActiveStyles();
+        getLog().lifecycle(String.format("Using active recipe(s) %s", activeRecipes));
+        getLog().lifecycle(String.format("Using active styles(s) %s", activeStyles));
+        if (activeRecipes.isEmpty()) {
+            return new ResultsContainer(baseDir, emptyList());
+        }
+        List<NamedStyles> styles = env.activateStyles(activeStyles);
+        File checkstyleConfig = extension.getCheckstyleConfigFile();
+        if (checkstyleConfig != null && checkstyleConfig.exists()) {
+            NamedStyles checkstyle = getRewrite().loadCheckstyleConfig(checkstyleConfig.toPath(), extension.getCheckstyleProperties());
+            styles.add(checkstyle);
+        }
 
-            Recipe recipe = env.activateRecipes(activeRecipes);
-
-            getLog().lifecycle("Validating active recipes");
-            Collection<Validated> validated = recipe.validateAll();
-            List<Validated.Invalid> failedValidations = validated.stream().map(Validated::failures)
-                    .flatMap(Collection::stream).collect(toList());
-            if (!failedValidations.isEmpty()) {
-                failedValidations.forEach(failedValidation -> getLog().error(
-                        "Recipe validation error in " + failedValidation.getProperty() + ": " +
-                                failedValidation.getMessage(), failedValidation.getException()));
-                if (getExtension().getFailOnInvalidActiveRecipes()) {
-                    throw new RuntimeException("Recipe validation errors detected as part of one or more activeRecipe(s). Please check error logs.");
-                } else {
-                    getLog().error("Recipe validation errors detected as part of one or more activeRecipe(s). Execution will continue regardless.");
-                }
+        Recipe recipe = env.activateRecipes(activeRecipes);
+
+        getLog().lifecycle("Validating active recipes");
+        Collection<Validated> validated = recipe.validateAll();
+        List<Validated.Invalid> failedValidations = validated.stream().map(Validated::failures)
+                .flatMap(Collection::stream).collect(toList());
+        if (!failedValidations.isEmpty()) {
+            failedValidations.forEach(failedValidation -> getLog().error(
+                    "Recipe validation error in " + failedValidation.getProperty() + ": " +
+                            failedValidation.getMessage(), failedValidation.getException()));
+            if (getExtension().getFailOnInvalidActiveRecipes()) {
+                throw new RuntimeException("Recipe validation errors detected as part of one or more activeRecipe(s). Please check error logs.");
+            } else {
+                getLog().error("Recipe validation errors detected as part of one or more activeRecipe(s). Execution will continue regardless.");
             }
+        }
 
+        List<SourceFile> sourceFiles;
+        if (useAstCache && astCache.containsKey(getProject().getRootProject().getRootDir())) {
+            getLog().lifecycle("Using cached in-memory ASTs...");
+            sourceFiles = getRewrite().toSourceFile(astCache.get(getProject().getRootProject().getRootDir()));
+        } else {
             InMemoryExecutionContext ctx = executionContext();
-            List<SourceFile> sourceFiles = projects.stream()
+            sourceFiles = projects.stream()
                     .flatMap(p -> parse(p, styles, ctx).stream())
                     .collect(toList());
+            if (useAstCache) {
+                astCache.put(getProject().getRootProject().getRootDir(), getRewrite().toBytes(sourceFiles));
+            }
+        }
+        getLog().lifecycle("Running recipe(s)...");
+        List<Result> results = recipe.run(sourceFiles);
 
-            getLog().lifecycle("Running recipe(s)...");
-            List<Result> results = recipe.run(sourceFiles);
+        return new ResultsContainer(baseDir, results);
+    }
 
-            return new ResultsContainer(baseDir, results);
-        } finally {
-            rewrite.shutdown();
-        }
+    protected void clearAstCache() {
+        astCache.clear();
+    }
+
+    protected void shutdownRewrite() {
+        rewrite.shutdown();
     }
 
     protected List<SourceFile> parse(Project subproject, List<NamedStyles> styles, InMemoryExecutionContext ctx) {

plugin/src/main/java/org/openrewrite/gradle/RewriteClearCacheTask.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright 2020 the original author or authors.
+ * <p>
+ * 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
+ * <p>
+ * https://www.apache.org/licenses/LICENSE-2.0
+ * <p>
+ * 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.openrewrite.gradle;
+
+import org.gradle.api.logging.Logger;
+import org.gradle.api.logging.Logging;
+import org.gradle.api.tasks.TaskAction;
+import org.openrewrite.gradle.RewriteReflectiveFacade.Environment;
+import org.openrewrite.gradle.RewriteReflectiveFacade.RecipeDescriptor;
+
+import javax.inject.Inject;
+import java.util.Collection;
+
+public class RewriteClearCacheTask extends AbstractRewriteTask {
+    private static final Logger log = Logging.getLogger(RewriteClearCacheTask.class);
+
+    @Override
+    protected Logger getLog() {
+        return log;
+    }
+
+    @Inject
+    public RewriteClearCacheTask() {
+        setGroup("rewrite");
+        setDescription("Clear in-memory AST cache");
+    }
+
+    @TaskAction
+    public void run() {
+        Environment env = environment();
+        Collection<RecipeDescriptor> availableRecipeDescriptors = env.listRecipeDescriptors();
+
+        log.quiet("Clearing AST Cache...");
+        clearAstCache();
+
+    }
+
+}

plugin/src/main/java/org/openrewrite/gradle/RewriteDryRunTask.java

@@ -18,8 +18,10 @@
 import org.gradle.api.logging.Logger;
 import org.gradle.api.logging.Logging;
 import org.gradle.api.specs.Specs;
+import org.gradle.api.tasks.Input;
 import org.gradle.api.tasks.Internal;
 import org.gradle.api.tasks.TaskAction;
+import org.gradle.api.tasks.options.Option;
 import org.openrewrite.gradle.RewriteReflectiveFacade.Result;
 
 import javax.inject.Inject;
@@ -46,67 +48,80 @@ public RewriteDryRunTask() {
         getOutputs().upToDateWhen(Specs.SATISFIES_NONE);
     }
 
+    @Option(description = "Cache the AST results in-memory when using the Gradle daemon.", option = "useAstCache")
+    public void setUseAstCache(boolean useAstCache) {
+        this.useAstCache = useAstCache;
+    }
+
+    @Input
+    public boolean isUseAstCache() {
+        return useAstCache;
+    }
+
     @Override
     protected Logger getLog() {
         return log;
     }
 
     @TaskAction
     public void run() {
-        ResultsContainer results = listResults();
-
-        if (results.isNotEmpty()) {
-            for (Result result : results.generated) {
-                assert result.getAfter() != null;
-                getLog().warn("These recipes would generate new file {}:", result.getAfter().getSourcePath());
-                logRecipesThatMadeChanges(result);
-            }
-            for (Result result : results.deleted) {
-                assert result.getBefore() != null;
-                getLog().warn("These recipes would delete file {}:", result.getBefore().getSourcePath());
-                logRecipesThatMadeChanges(result);
-            }
-            for (Result result : results.moved) {
-                assert result.getBefore() != null;
-                assert result.getAfter() != null;
-                getLog().warn("These recipes would move file from {} to {}:", result.getBefore().getSourcePath(), result.getAfter().getSourcePath());
-                logRecipesThatMadeChanges(result);
-            }
-            for (Result result : results.refactoredInPlace) {
-                assert result.getBefore() != null;
-                getLog().warn("These recipes would make results to {}:", result.getBefore().getSourcePath());
-                logRecipesThatMadeChanges(result);
-            }
+        try {
+            ResultsContainer results = listResults();
 
-            Path patchFile = getReportPath();
-            //noinspection ResultOfMethodCallIgnored
-            patchFile.getParent().toFile().mkdirs();
-            try (BufferedWriter writer = Files.newBufferedWriter(patchFile)) {
-                Stream.concat(
-                        Stream.concat(results.generated.stream(), results.deleted.stream()),
-                        Stream.concat(results.moved.stream(), results.refactoredInPlace.stream())
-                )
-                        .map(Result::diff)
-                        .forEach(diff -> {
-                            try {
-                                writer.write(diff + "\n");
-                            } catch (IOException e) {
-                                throw new RuntimeException(e);
-                            }
-                        });
+            if (results.isNotEmpty()) {
+                for (Result result : results.generated) {
+                    assert result.getAfter() != null;
+                    getLog().warn("These recipes would generate new file {}:", result.getAfter().getSourcePath());
+                    logRecipesThatMadeChanges(result);
+                }
+                for (Result result : results.deleted) {
+                    assert result.getBefore() != null;
+                    getLog().warn("These recipes would delete file {}:", result.getBefore().getSourcePath());
+                    logRecipesThatMadeChanges(result);
+                }
+                for (Result result : results.moved) {
+                    assert result.getBefore() != null;
+                    assert result.getAfter() != null;
+                    getLog().warn("These recipes would move file from {} to {}:", result.getBefore().getSourcePath(), result.getAfter().getSourcePath());
+                    logRecipesThatMadeChanges(result);
+                }
+                for (Result result : results.refactoredInPlace) {
+                    assert result.getBefore() != null;
+                    getLog().warn("These recipes would make results to {}:", result.getBefore().getSourcePath());
+                    logRecipesThatMadeChanges(result);
+                }
 
-            } catch (Exception e) {
-                throw new RuntimeException("Unable to generate rewrite result file.", e);
-            }
-            getLog().warn("Report available:");
-            getLog().warn(indent(1, patchFile.normalize().toString()));
-            getLog().warn("Run 'gradle rewriteRun' to apply the recipes.");
+                Path patchFile = getReportPath();
+                //noinspection ResultOfMethodCallIgnored
+                patchFile.getParent().toFile().mkdirs();
+                try (BufferedWriter writer = Files.newBufferedWriter(patchFile)) {
+                    Stream.concat(
+                                    Stream.concat(results.generated.stream(), results.deleted.stream()),
+                                    Stream.concat(results.moved.stream(), results.refactoredInPlace.stream())
+                            )
+                            .map(Result::diff)
+                            .forEach(diff -> {
+                                try {
+                                    writer.write(diff + "\n");
+                                } catch (IOException e) {
+                                    throw new RuntimeException(e);
+                                }
+                            });
+                } catch (Exception e) {
+                    throw new RuntimeException("Unable to generate rewrite result file.", e);
+                }
+                getLog().warn("Report available:");
+                getLog().warn(indent(1, patchFile.normalize().toString()));
+                getLog().warn("Run 'gradle rewriteRun' to apply the recipes.");
 
-            if (getExtension().getFailOnDryRunResults()) {
-                throw new RuntimeException("Applying recipes would make changes. See logs for more details.");
+                if (getExtension().getFailOnDryRunResults()) {
+                    throw new RuntimeException("Applying recipes would make changes. See logs for more details.");
+                }
+            } else {
+                getLog().lifecycle("Applying recipes would make no changes. No report generated.");
             }
-        } else {
-            getLog().lifecycle("Applying recipes would make no changes. No report generated.");
+        } finally {
+            shutdownRewrite();
         }
     }
 }

plugin/src/main/java/org/openrewrite/gradle/RewritePlugin.java

@@ -24,9 +24,7 @@
 import org.gradle.api.plugins.quality.CheckstylePlugin;
 
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
 
 /**
  * Adds the RewriteExtension to the current project and registers tasks per-sourceSet.
@@ -74,6 +72,10 @@ public void apply(Project rootProject) {
                 .setConfiguration(rewriteConf)
                 .setExtension(extension)
                 .setProjects(projects);
+        Task rewriteCleanCache = rootProject.getTasks().create("rewriteClearCache", RewriteClearCacheTask.class)
+                .setConfiguration(rewriteConf)
+                .setExtension(extension)
+                .setProjects(projects);
 
         rootProject.allprojects(project -> {
 

plugin/src/main/java/org/openrewrite/gradle/RewriteReflectiveFacade.java

@@ -938,15 +938,34 @@ public XmlParser xmlParser() {
         }
     }
 
-    public void shutdown() {
+    byte[] toBytes(List<SourceFile> sourceFiles) {
+        List<Object> trees = sourceFiles.stream().map(s -> s.real).collect(Collectors.toList());
         try {
-            Class<?> c = getClassLoader().loadClass("org.openrewrite.java.tree.J");
-            Method javaClearCaches = c.getMethod("clearCaches");
-            javaClearCaches.invoke(null);
-
-            getClassLoader().loadClass("org.openrewrite.shaded.jgit.api.Git").getMethod("shutdown").invoke(null);
+            Class<?> serializerClass = getClassLoader().loadClass("org.openrewrite.TreeSerializer");
+            Object serializer = serializerClass.getDeclaredConstructor().newInstance();
+            return (byte[]) serializerClass.getMethod("write", Iterable.class).invoke(serializer, trees);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
 
+    List<SourceFile> toSourceFile(byte[] trees) {
+        try {
+            Class<?> serializerClass = getClassLoader().loadClass("org.openrewrite.TreeSerializer");
+            Object serializer = serializerClass.getDeclaredConstructor().newInstance();
+            List<Object> sources = (List<Object>) serializerClass.getMethod("readList", byte[].class).invoke(serializer, trees);
+            return sources.stream().map(SourceFile::new).collect(Collectors.toList());
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
 
+    public void shutdown() {
+        try {
+            getClassLoader().loadClass("org.openrewrite.java.tree.J").getMethod("clearCaches").invoke(null);
+            getClassLoader().loadClass("org.openrewrite.shaded.jgit.api.Git").getMethod("shutdown").invoke(null);
+            getClassLoader().loadClass("org.openrewrite.scheduling.ForkJoinScheduler").getMethod("shutdown").invoke(null);
+            classLoader = null;
         } catch (Exception e) {
             throw new RuntimeException(e);
         }