com.slimgears.apt.util.ElementVisitors Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of apt-utils Show documentation
General purpose utils / module: apt-utils
There is a newer version: 0.7.58
package com.slimgears.apt.util;

import javax.lang.model.element.*;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.util.AbstractElementVisitor8;
import java.util.*;
import java.util.function.BiFunction;
import java.util.function.BinaryOperator;

import static com.slimgears.apt.util.VisitorUtils.takeFirstNonNull;

public class ElementVisitors {
    public static  Builder builder() {
        return builder(takeFirstNonNull());
    }

    public static  Builder builder(BinaryOperator resultCombiner) {
        return new Builder<>(resultCombiner);
    }

    public static Collection allReferencedTypes(Element e) {
        List types = new ArrayList<>();
        e.accept(ElementVisitors.builder()
                .on(TypeElement.class, type -> { types.addAll(TypeVisitors.findAllOf(DeclaredType.class, type.asType())); })
                .on(VariableElement.class, var -> { types.addAll(TypeVisitors.findAllOf(DeclaredType.class, var.asType())); })
                .on(ExecutableElement.class, ex -> { types.addAll(TypeVisitors.findAllOf(DeclaredType.class, ex.getReturnType())); })
                .build(), null);
        return new HashSet<>(types);
    }

    public static  List findAllOf(Class cls, Element e) {
        return VisitorUtils.findAllOf(cls, e, ElementVisitors::builder);
    }

    public static class Builder extends VisitorUtils.AbstractBuilder, R, P, Builder> {
        public Builder(BinaryOperator combiner) {
            super(combiner);
        }

        @Override
        protected Builder self() {
            return this;
        }

        @Override
        public ElementVisitor build() {
            return new AbstractElementVisitor8() {
                private Map visitedElements = new HashMap<>();

                @Override
                public R visitPackage(PackageElement e, P param) {
                    return visitElement(PackageElement.class, e, param);
                }

                @Override
                public R visitType(TypeElement e, P param) {
                    return visitElement(TypeElement.class, e, param);
                }

                @Override
                public R visitVariable(VariableElement e, P param) {
                    return aggregate(
                            visitElement(VariableElement.class, e, param));
                }

                @Override
                public R visitExecutable(ExecutableElement e, P param) {
                    return aggregate(
                            visitElement(ExecutableElement.class, e, param),
                            e.getParameters().stream().map(pe -> pe.accept(this, param)));
                }

                @Override
                public R visitTypeParameter(TypeParameterElement e, P param) {
                    return aggregate(
                            visitElement(TypeParameterElement.class, e, param),
                            e.getGenericElement().accept(this, param));

                }

                private  R visitElement(Class cls, E element, P param) {
                    if (!filter.test(element, param)) {
                        return null;
                    }

                    BiFunction func = listenerOf(cls);

                    return visitedElements.computeIfAbsent(
                            element,
                            _element -> aggregate(
                            aggregate(onAny.apply(element, param), func.apply(element, param)),
                            element.getEnclosedElements().stream().map(el -> el.accept(this, param))));
                }
            };
        }

        @Override
        protected R accept(Element element, ElementVisitor visitor, P param) {
            return element.accept(visitor, param);
        }
    }
}