1package com.github.javaparser.generator;
2
3import com.github.javaparser.ast.CompilationUnit;
4import com.github.javaparser.ast.body.ClassOrInterfaceDeclaration;
5import com.github.javaparser.ast.body.MethodDeclaration;
6import com.github.javaparser.ast.expr.MarkerAnnotationExpr;
7import com.github.javaparser.ast.expr.Name;
8import com.github.javaparser.metamodel.BaseNodeMetaModel;
9import com.github.javaparser.metamodel.JavaParserMetaModel;
10import com.github.javaparser.utils.Log;
11import com.github.javaparser.utils.SourceRoot;
12
13import java.util.Optional;
14
15import static com.github.javaparser.ast.Modifier.PUBLIC;
16
17/**
18 * Makes it easier to generate visitor classes.
19 * It will create missing visit methods on the fly,
20 * and will ask you to fill in the bodies of the visit methods.
21 */
22public abstract class VisitorGenerator extends Generator {
23    private final String pkg;
24    private final String visitorClassName;
25    private final String returnType;
26    private final String argumentType;
27    private final boolean createMissingVisitMethods;
28
29    protected VisitorGenerator(SourceRoot sourceRoot, String pkg, String visitorClassName, String returnType, String argumentType, boolean createMissingVisitMethods) {
30        super(sourceRoot);
31        this.pkg = pkg;
32        this.visitorClassName = visitorClassName;
33        this.returnType = returnType;
34        this.argumentType = argumentType;
35        this.createMissingVisitMethods = createMissingVisitMethods;
36    }
37
38    public final void generate() throws Exception {
39        Log.info("Running %s", getClass().getSimpleName());
40
41        final CompilationUnit compilationUnit = sourceRoot.tryToParse(pkg, visitorClassName + ".java").getResult().get();
42
43        Optional<ClassOrInterfaceDeclaration> visitorClassOptional = compilationUnit.getClassByName(visitorClassName);
44        if (!visitorClassOptional.isPresent()) {
45            visitorClassOptional = compilationUnit.getInterfaceByName(visitorClassName);
46        }
47        final ClassOrInterfaceDeclaration visitorClass = visitorClassOptional.get();
48
49        JavaParserMetaModel.getNodeMetaModels().stream()
50                .filter((baseNodeMetaModel) -> !baseNodeMetaModel.isAbstract())
51                .forEach(node -> generateVisitMethodForNode(node, visitorClass, compilationUnit));
52        after();
53    }
54
55    protected void after() throws Exception {
56
57    }
58
59    private void generateVisitMethodForNode(BaseNodeMetaModel node, ClassOrInterfaceDeclaration visitorClass, CompilationUnit compilationUnit) {
60        final Optional<MethodDeclaration> existingVisitMethod = visitorClass.getMethods().stream()
61                .filter(m -> m.getNameAsString().equals("visit"))
62                .filter(m -> m.getParameter(0).getType().toString().equals(node.getTypeName()))
63                .findFirst();
64
65        if (existingVisitMethod.isPresent()) {
66            generateVisitMethodBody(node, existingVisitMethod.get(), compilationUnit);
67        } else if (createMissingVisitMethods) {
68            MethodDeclaration newVisitMethod = visitorClass.addMethod("visit")
69                    .addParameter(node.getTypeNameGenerified(), "n")
70                    .addParameter(argumentType, "arg")
71                    .setType(returnType);
72            if (!visitorClass.isInterface()) {
73                newVisitMethod
74                        .addAnnotation(new MarkerAnnotationExpr(new Name("Override")))
75                        .addModifier(PUBLIC);
76            }
77            generateVisitMethodBody(node, newVisitMethod, compilationUnit);
78        }
79    }
80
81    protected abstract void generateVisitMethodBody(BaseNodeMetaModel node, MethodDeclaration visitMethod, CompilationUnit compilationUnit);
82}
83