Implement radix tree based JavaTypeCache

rewrite-benchmarks/build.gradle.kts

@@ -13,6 +13,7 @@ dependencies {
     jmh(project(":rewrite-maven"))
     jmh("org.rocksdb:rocksdbjni:latest.release")
     jmh("org.openjdk.jmh:jmh-core:latest.release")
+    jmh("org.openjdk.jol:jol-core:latest.release")
     jmh("io.github.fastfilter:fastfilter:latest.release")
 
     // Nebula doesn't like having jmhAnnotationProcessor without jmh so we just add it twice.

rewrite-benchmarks/src/jmh/java/org/openrewrite/benchmarks/java/JavaCompilationUnitState.java

@@ -15,38 +15,53 @@
  */
 package org.openrewrite.benchmarks.java;
 
+import org.jspecify.annotations.Nullable;
 import org.openjdk.jmh.annotations.*;
 import org.openjdk.jmh.infra.Blackhole;
+import org.openjdk.jol.info.GraphLayout;
 import org.openrewrite.InMemoryExecutionContext;
 import org.openrewrite.LargeSourceSet;
 import org.openrewrite.SourceFile;
 import org.openrewrite.internal.InMemoryLargeSourceSet;
 import org.openrewrite.java.JavaParser;
+import org.openrewrite.java.internal.AdaptiveRadixJavaTypeCache;
+import org.openrewrite.java.internal.JavaTypeCache;
 
 import java.net.URISyntaxException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 
 @State(Scope.Benchmark)
 public class JavaCompilationUnitState {
+    JavaParser.Builder<? extends JavaParser, ?> javaParser;
     List<SourceFile> sourceFiles;
+    List<Path> inputs;
+    JavaTypeCache snappyTypeCache;
+    AdaptiveRadixJavaTypeCache radixMapTypeCache;
+    MapJavaTypeCache typeCache;
+
+    public static void main(String[] args) throws URISyntaxException {
+        new JavaCompilationUnitState().setup();
+    }
 
     @Setup(Level.Trial)
     public void setup() throws URISyntaxException {
         Path rewriteRoot = Paths.get(ChangeTypeBenchmark.class.getResource("./")
                 .toURI()).resolve("../../../../../../../../").normalize();
 
-        List<Path> inputs = Arrays.asList(
+        inputs = Arrays.asList(
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/internal/lang/Nullable.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/internal/lang/NullUtils.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/internal/MetricsHelper.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/internal/ListUtils.java"),
-                rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/internal/StringUtils.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/internal/PropertyPlaceholderHelper.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/internal/RecipeIntrospectionUtils.java"),
+                rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/internal/StringUtils.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/Tree.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/ExecutionContext.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/InMemoryExecutionContext.java"),
@@ -61,21 +76,45 @@ public void setup() throws URISyntaxException {
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/Result.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/SourceFile.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/Recipe.java"),
+                rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/ScanningRecipe.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/Validated.java"),
+                rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/ValidationException.java"),
+                rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/TreeVisitor.java"),
+                rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/TreeObserver.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/config/DeclarativeRecipe.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/config/ResourceLoader.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/config/YamlResourceLoader.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/config/ClasspathScanningLoader.java"),
                 rewriteRoot.resolve("rewrite-core/src/main/java/org/openrewrite/config/RecipeIntrospectionException.java")
         );
 
-        sourceFiles = JavaParser.fromJavaVersion()
-                .classpath("jsr305", "classgraph", "jackson-annotations", "micrometer-core", "slf4j-api",
-                        "org.openrewrite.jgit")
+        javaParser = JavaParser.fromJavaVersion()
+                .classpath("jsr305", "classgraph", "jackson-annotations", "micrometer-core",
+                        "jgit", "jspecify", "lombok", "annotations");
 //                .logCompilationWarningsAndErrors(true)
-                .build()
-                .parse(inputs, null, new InMemoryExecutionContext(Throwable::printStackTrace))
+
+        typeCache = new MapJavaTypeCache();
+        JavaParser parser = javaParser.typeCache(typeCache).build();
+        sourceFiles = parser
+                .parse(inputs, null, new InMemoryExecutionContext())
                 .collect(Collectors.toList());
+
+        radixMapTypeCache = new AdaptiveRadixJavaTypeCache();
+        for (Map.Entry<String, Object> entry : typeCache.map().entrySet()) {
+            radixMapTypeCache.put(entry.getKey(), entry.getValue());
+        }
+
+        snappyTypeCache = new JavaTypeCache();
+        for (Map.Entry<String, Object> entry : typeCache.map().entrySet()) {
+            snappyTypeCache.put(entry.getKey(), entry.getValue());
+        }
+    }
+
+    void printMemory() {
+        long retainedSize = GraphLayout.parseInstance(radixMapTypeCache).totalSize();
+        System.out.printf("Retained AdaptiveRadixTree size: %10d bytes\n", retainedSize);
+        retainedSize = GraphLayout.parseInstance(snappyTypeCache).totalSize();
+        System.out.printf("Retained Snappy size:            %10d bytes\n", retainedSize);
     }
 
     @TearDown(Level.Trial)
@@ -90,4 +129,37 @@ public LargeSourceSet getSourceSet() {
     public List<SourceFile> getSourceFiles() {
         return sourceFiles;
     }
+
+    static class MapJavaTypeCache extends JavaTypeCache {
+
+        Map<String, Object> typeCache = new HashMap<>();
+
+        public <T> @Nullable T get(String signature) {
+            //noinspection unchecked
+            return (T) typeCache.get(signature);
+        }
+
+        public void put(String signature, Object o) {
+            typeCache.put(signature, o);
+        }
+
+        public Map<String, Object> map() {
+            return typeCache;
+        }
+
+        public void clear() {
+            typeCache.clear();
+        }
+
+        public int size() {
+            return typeCache.size();
+        }
+
+        @Override
+        public MapJavaTypeCache clone() {
+            MapJavaTypeCache clone = (MapJavaTypeCache) super.clone();
+            clone.typeCache = new HashMap<>(this.typeCache);
+            return clone;
+        }
+    }
 }

rewrite-benchmarks/src/jmh/java/org/openrewrite/benchmarks/java/JavaParserBenchmark.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2022 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.benchmarks.java;
+
+import org.openjdk.jmh.annotations.*;
+import org.openjdk.jmh.infra.Blackhole;
+import org.openjdk.jmh.profile.GCProfiler;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+import org.openrewrite.InMemoryExecutionContext;
+import org.openrewrite.java.JavaParser;
+import org.openrewrite.java.internal.AdaptiveRadixJavaTypeCache;
+import org.openrewrite.java.internal.JavaTypeCache;
+
+import java.net.URISyntaxException;
+import java.util.concurrent.TimeUnit;
+
+@Fork(1)
+@Measurement(iterations = 2)
+@Warmup(iterations = 2)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Threads(4)
+public class JavaParserBenchmark {
+
+    @Benchmark
+    public void snappy(JavaCompilationUnitState state, Blackhole bh) {
+        JavaTypeCache typeCache = new JavaTypeCache();
+        JavaParser parser = state.javaParser.typeCache(typeCache).build();
+        parser
+                .parse(state.inputs, null, new InMemoryExecutionContext())
+                .forEach(bh::consume);
+    }
+
+    @Benchmark
+    public void adaptiveRadix(JavaCompilationUnitState state, Blackhole bh) {
+        AdaptiveRadixJavaTypeCache typeCache = new AdaptiveRadixJavaTypeCache();
+        JavaParser parser = state.javaParser.typeCache(typeCache).build();
+        parser
+                .parse(state.inputs, null, new InMemoryExecutionContext())
+                .forEach(bh::consume);
+    }
+
+    public static void main(String[] args) throws RunnerException, URISyntaxException {
+        Options opt = new OptionsBuilder()
+                .include(JavaParserBenchmark.class.getSimpleName())
+                .addProfiler(GCProfiler.class)
+                .shouldFailOnError(true)
+                .build();
+        new Runner(opt).run();
+        JavaCompilationUnitState state = new JavaCompilationUnitState();
+        state.setup();
+        state.printMemory();
+    }
+}

rewrite-benchmarks/src/jmh/java/org/openrewrite/benchmarks/java/JavaTypeCacheBenchmark.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2022 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.benchmarks.java;
+
+import org.openjdk.jmh.annotations.*;
+import org.openjdk.jmh.infra.Blackhole;
+import org.openjdk.jmh.profile.GCProfiler;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.Options;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+import org.openrewrite.java.internal.JavaTypeCache;
+import org.openrewrite.java.internal.AdaptiveRadixJavaTypeCache;
+
+import java.net.URISyntaxException;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+@Fork(1)
+@Measurement(iterations = 2)
+@Warmup(iterations = 2)
+@BenchmarkMode(Mode.Throughput)
+@OutputTimeUnit(TimeUnit.SECONDS)
+@Threads(4)
+public class JavaTypeCacheBenchmark {
+
+    @Benchmark
+    public void writeSnappy(JavaCompilationUnitState state, Blackhole bh) {
+        JavaTypeCache typeCache = new JavaTypeCache();
+        for (Map.Entry<String, Object> entry : state.typeCache.map().entrySet()) {
+            typeCache.put(entry.getKey(), entry.getValue());
+        }
+    }
+
+    @Benchmark
+    public void writeAdaptiveRadix(JavaCompilationUnitState state, Blackhole bh) {
+        AdaptiveRadixJavaTypeCache typeCache = new AdaptiveRadixJavaTypeCache();
+        for (Map.Entry<String, Object> entry : state.typeCache.map().entrySet()) {
+            typeCache.put(entry.getKey(), entry.getValue());
+        }
+    }
+
+    @Benchmark
+    public void readSnappy(JavaCompilationUnitState state, Blackhole bh) {
+        for (Map.Entry<String, Object> entry : state.typeCache.map().entrySet()) {
+            bh.consume(state.snappyTypeCache.get(entry.getKey()));
+        }
+    }
+
+    @Benchmark
+    public void readAdaptiveRadix(JavaCompilationUnitState state, Blackhole bh) {
+        for (Map.Entry<String, Object> entry : state.typeCache.map().entrySet()) {
+            bh.consume(state.radixMapTypeCache.get(entry.getKey()));
+        }
+    }
+
+    public static void main(String[] args) throws RunnerException, URISyntaxException {
+        Options opt = new OptionsBuilder()
+                .include(JavaTypeCacheBenchmark.class.getSimpleName())
+                .addProfiler(GCProfiler.class)
+                .shouldFailOnError(true)
+                .build();
+        new Runner(opt).run();
+        JavaCompilationUnitState state = new JavaCompilationUnitState();
+        state.setup();
+        state.printMemory();
+    }
+}