public void addTransform(Class<? extends ASTTransformation> transform, ASTNode node) {
    GroovyASTTransformation annotation = transform.getAnnotation(GroovyASTTransformation.class);
    if (annotation != null) {
        Map<Class<? extends ASTTransformation>, Set<ASTNode>> transforms = getTransforms(annotation.phase());
        Set<ASTNode> nodes = transforms.computeIfAbsent(transform, k -> new LinkedHashSet<>());
        nodes.add(node);
    }
}
 

private Map<CompilePhase, Map<Class<? extends ASTTransformation>, Set<ASTNode>>> getTransformInstances() {
    if (transformInstances == null) {
        transformInstances = new EnumMap<>(CompilePhase.class);
        for (CompilePhase phase : CompilePhase.values()) {
            transformInstances.put(phase, new LinkedHashMap<>());
        }
    }
    return transformInstances;
}
 

@Override
public void addTransform(final Class<? extends ASTTransformation> transform, final ASTNode node) {
    throw new UnsupportedOperationException();
}
 

public Map<Class <? extends ASTTransformation>, Set<ASTNode>> getTransforms(CompilePhase phase) {
    return getTransformInstances().get(phase);
}