Observers for AST nodes in OpenRewrite

OpenRewrite is at its core a programmatic means of rewriting source code. It is recombinant by nature, but at its core is always a visitor-based refactoring technology that offers the full expressiveness of a programming language to mutate the code. Always looking for inspiration from other language engineers on techniques that might enhance the expressive power of the framework, I recently encountered Federico Tomassetti's blog post on the TreeObserver milestone feature of JavaParser 3.0.

Seeing the potential usefulness of the concept, I created a riff on this concept in OpenRewrite, and in some cases expanded on the idea from Federico's post. I hope this second look at the feature is as useful to the JavaParser community as the implementation will be to the OpenRewrite community. What follows is a parallel blog post documenting the feature as it was introduced in OpenRewrite 7.20.0.

What observers on AST nodes could be used for?

OpenRewrite recipes are mostly visitor-based, and visitors themselves are a form of event-driven programming. Event-driven systems are essential to operating with source code, because most recipes need to act on a particular structure of code that itself could exist arbitrarily deeply inside a source file. Trying to act on a top-level source file imperatively is impractical in many cases to express the inherently recursive structure of code.

A TreeVisitor is an event-driven response to a single source file. A TreeVisitor operating over many source files can produce changes in one or more of these source files. A TreeObserver is another form of event-driven response, but oriented toward a response to a specific node or cross-sectionally to any node that matches some criteria.

A TreeObserver could, for example, allow us to respond in some way to every modification of a class field across any number of source files. It's kind of like Aspect Oriented Programming, but for AST visiting en masse across a whole codebase.

The TreeObserver

We introduced a new interface named TreeObserver. A TreeObserver can subscribe to changes to matching AST elements.

public interface TreeObserver {
  Tree propertyChanged(
    String property,
    Cursor cursor, 
    Tree newTree,
    Object oldValue,
    Object newValue
  );
}

The propertyChanged event handler gives you access to property changes to an AST element. These property changes are limited to leaf nodes in the tree, i.e. things like string and numeric properties of an AST element and not subtrees.

Notice that the propertyChanged method also returns a Tree. This permits you to further modify the tree as part of a cross-cutting observer, e.g. to update data flow edges, formatting, or type information.

Given that a TreeObserver matches on AST elements by a predicate match, they logically are not associated with a particular AST element. This is a key difference between the OpenRewrite and JavaParser implementations today. The natural place in OpenRewrite to register TreeObserver is the ExecutionContext that gets passed to recipe runs.

PlainText pt = ...;

ExecutionContext ctx = new InMemoryExecutionContext();
TreeObserver.Subscription sub = new TreeObserver.Subscription(
        (property, cursor, newTree, oldValue, newValue) -> {
            System.out.println("Property " + property + " changed from " +
                    oldValue + " to " + newValue);
            return newTree;
        });

// subscribe to ALL changes
sub.subscribe(t -> true);

// subscribe to changes to AST elements of type PlainText
sub.subscribeToType(PlainText.class);

// subscribe to changes to a particular AST instance
sub.subscribe(pt);

ctx.addObserver(sub);

Implementing support for observers

In OpenRewrite, implementing support for observers only touched one piece of code in TreeVisitor. If the context object being passed to a visitor is an ExecutionContext and an AST element changes, each registered TreeObserver is checked to see if the changing AST element is subscribed and the observer called if so.

Conclusions

Many thanks to Federico and the rest of theJavaParser for the interesting concept. Looking forward to seeing how this will be used in practice!