Refactor TreeObserver (#4160)

The `TreeObserver` mechanism is currently only used by a handful of recipes and only via `TreeVisitor#collect()`. This use does however not require the diffing functionality as implemented using `ObjectDiffer`. This turns out to be rather expensive, so this commit makes it an opt-in feature and moves the logic from `TreeVisitor` into the `TreeObserver` itself.

rewrite-core/src/main/java/org/openrewrite/TreeObserver.java

@@ -15,10 +15,17 @@
  */
 package org.openrewrite;
 
+import de.danielbechler.diff.ObjectDiffer;
+import de.danielbechler.diff.ObjectDifferBuilder;
+import de.danielbechler.diff.inclusion.Inclusion;
+import de.danielbechler.diff.inclusion.InclusionResolver;
+import de.danielbechler.diff.node.DiffNode;
 import org.openrewrite.internal.lang.Nullable;
 
+import java.beans.Transient;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Predicate;
 
 @Incubating(since = "7.20.0")
@@ -35,8 +42,37 @@ final class Subscription {
         private final TreeObserver observer;
         private final List<Predicate<Tree>> predicates = new ArrayList<>();
 
+        @Nullable
+        private final ObjectDiffer differ;
+
         public Subscription(TreeObserver observer) {
+            this(observer, false);
+        }
+
+        public Subscription(TreeObserver observer, boolean diffChanges) {
             this.observer = observer;
+            if (diffChanges) {
+                differ = ObjectDifferBuilder.startBuilding()
+                        .inclusion()
+                        .resolveUsing(new InclusionResolver() {
+                            @Override
+                            public Inclusion getInclusion(DiffNode node) {
+                                if (node.getPropertyAnnotation(Transient.class) != null) {
+                                    return Inclusion.EXCLUDED;
+                                }
+                                return Inclusion.DEFAULT;
+                            }
+
+                            @Override
+                            public boolean enablesStrictIncludeMode() {
+                                return false;
+                            }
+                        })
+                        .and()
+                        .build();
+            } else {
+                differ = null;
+            }
         }
 
         public Subscription subscribe(@Nullable Tree tree) {
@@ -87,8 +123,25 @@ public boolean isSubscribed(@Nullable Tree tree) {
             return false;
         }
 
-        public TreeObserver getObserver() {
-            return observer;
+        public <T extends Tree> T treeChanged(Cursor cursor, T after, Tree before) {
+            observer.treeChanged(cursor, after);
+            if (differ != null) {
+                return diff(after, before, cursor);
+            }
+            return after;
+        }
+
+        private <T extends Tree> T diff(T after, Tree before, Cursor cursor) {
+            //noinspection DataFlowIssue
+            DiffNode diff = differ.compare(after, before);
+            AtomicReference<T> t2 = new AtomicReference<>(after);
+            diff.visit((node, visit) -> {
+                if (!node.hasChildren() && node.getPropertyName() != null) {
+                    //noinspection unchecked
+                    t2.set((T) observer.propertyChanged(node.getPropertyName(), cursor, t2.get(), node.canonicalGet(before), node.canonicalGet(t2.get())));
+                }
+            });
+            return t2.get();
         }
     }
 }

rewrite-core/src/main/java/org/openrewrite/TreeVisitor.java

@@ -15,11 +15,6 @@
  */
 package org.openrewrite;
 
-import de.danielbechler.diff.ObjectDiffer;
-import de.danielbechler.diff.ObjectDifferBuilder;
-import de.danielbechler.diff.inclusion.Inclusion;
-import de.danielbechler.diff.inclusion.InclusionResolver;
-import de.danielbechler.diff.node.DiffNode;
 import io.micrometer.core.instrument.DistributionSummary;
 import io.micrometer.core.instrument.Metrics;
 import io.micrometer.core.instrument.Timer;
@@ -30,14 +25,12 @@
 import org.openrewrite.marker.Marker;
 import org.openrewrite.marker.Markers;
 
-import java.beans.Transient;
 import java.lang.reflect.ParameterizedType;
 import java.lang.reflect.Type;
 import java.lang.reflect.TypeVariable;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 
@@ -81,32 +74,6 @@ public static <T extends Tree, P> TreeVisitor<T, P> noop() {
     private int visitCount;
     private final DistributionSummary visitCountSummary = DistributionSummary.builder("rewrite.visitor.visit.method.count").description("Visit methods called per source file visited.").tag("visitor.class", getClass().getName()).register(Metrics.globalRegistry);
 
-    private ObjectDiffer differ;
-
-    private ObjectDiffer getObjectDiffer() {
-        if (differ == null) {
-            differ = ObjectDifferBuilder.startBuilding()
-                    .inclusion()
-                    .resolveUsing(new InclusionResolver() {
-                        @Override
-                        public Inclusion getInclusion(DiffNode node) {
-                            if (node.getPropertyAnnotation(Transient.class) != null) {
-                                return Inclusion.EXCLUDED;
-                            }
-                            return Inclusion.DEFAULT;
-                        }
-
-                        @Override
-                        public boolean enablesStrictIncludeMode() {
-                            return false;
-                        }
-                    })
-                    .and()
-                    .build();
-        }
-        return differ;
-    }
-
     public boolean isAcceptable(SourceFile sourceFile, P p) {
         return true;
     }
@@ -290,16 +257,7 @@ public T visit(@Nullable Tree tree, P p) {
                     ExecutionContext ctx = (ExecutionContext) p;
                     for (TreeObserver.Subscription observer : ctx.getObservers()) {
                         if (observer.isSubscribed(tree)) {
-                            observer.getObserver().treeChanged(getCursor(), t);
-                            DiffNode diff = getObjectDiffer().compare(t, tree);
-                            AtomicReference<T> t2 = new AtomicReference<>(t);
-                            diff.visit((node, visit) -> {
-                                if (!node.hasChildren() && node.getPropertyName() != null) {
-                                    //noinspection unchecked
-                                    t2.set((T) observer.getObserver().propertyChanged(node.getPropertyName(), getCursor(), t2.get(), node.canonicalGet(tree), node.canonicalGet(t2.get())));
-                                }
-                            });
-                            t = t2.get();
+                            t = observer.treeChanged(getCursor(), t, tree);
                         }
                     }
                 }

rewrite-core/src/test/java/org/openrewrite/TreeObserverTest.java

@@ -49,7 +49,7 @@ public Tree propertyChanged(String property, Cursor cursor, Tree newTree, Object
                     }
                     return newTree;
                 }
-            }).subscribeToType(PlainText.class))),
+            }, true).subscribeToType(PlainText.class))),
           text(
             "hello jon",
             "hello jonathan"