001/*
002 * JDrupes MDoclet
003 * Copyright (C) 2021 Michael N. Lipp
004 * 
005 * This program is free software; you can redistribute it and/or modify it 
006 * under the terms of the GNU Affero General Public License as published by 
007 * the Free Software Foundation; either version 3 of the License, or 
008 * (at your option) any later version.
009 *
010 * This program is distributed in the hope that it will be useful, but 
011 * WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
012 * or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License 
013 * for more details.
014 *
015 * You should have received a copy of the GNU Affero General Public License along 
016 * with this program; if not, see <http://www.gnu.org/licenses/>.
017 */
018
019package org.jdrupes.taglets.plantUml;
020
021import java.io.IOException;
022import java.nio.file.Files;
023import java.nio.file.Paths;
024import java.util.ArrayList;
025import java.util.Collections;
026import java.util.List;
027import java.util.Set;
028
029import javax.lang.model.element.Element;
030import javax.lang.model.element.Name;
031import javax.lang.model.element.PackageElement;
032import javax.lang.model.element.TypeElement;
033import javax.lang.model.util.SimpleElementVisitor8;
034import javax.tools.DocumentationTool;
035import javax.tools.FileObject;
036import javax.tools.JavaFileManager;
037
038import com.sun.source.doctree.CommentTree;
039import com.sun.source.doctree.DocTree;
040import com.sun.source.doctree.ErroneousTree;
041import com.sun.source.doctree.TextTree;
042import com.sun.source.doctree.UnknownBlockTagTree;
043import com.sun.source.util.SimpleDocTreeVisitor;
044
045import jdk.javadoc.doclet.Doclet;
046import jdk.javadoc.doclet.DocletEnvironment;
047import jdk.javadoc.doclet.Taglet;
048import net.sourceforge.plantuml.FileFormat;
049import net.sourceforge.plantuml.FileFormatOption;
050import net.sourceforge.plantuml.SourceStringReader;
051import net.sourceforge.plantuml.preproc.Defines;
052
053/**
054 * A JDK11 doclet that generates UML diagrams from PlantUML 
055 * specifications in the comment.
056 */
057public class PlantUml implements Taglet {
058
059    private DocletEnvironment env;
060    private JavaFileManager fileManager;
061    private List<String> plantConfigData;
062
063    @Override
064    public void init(DocletEnvironment env, Doclet doclet) {
065        this.env = env;
066        fileManager = env.getJavaFileManager();
067    }
068
069    @Override
070    public String getName() {
071        return "plantUml";
072    }
073
074    @Override
075    public Set<Location> getAllowedLocations() {
076        return Set.of(Taglet.Location.OVERVIEW, Taglet.Location.PACKAGE,
077            Taglet.Location.TYPE);
078    }
079
080    @Override
081    public boolean isInlineTag() {
082        return false;
083    }
084
085    @Override
086    public String toString(List<? extends DocTree> tags, Element element) {
087        for (DocTree tagTree : tags) {
088            processTag(tagTree, element);
089        }
090        return "";
091    }
092
093    private void processTag(DocTree tree, Element element) {
094        String plantUmlSource = extractPlantUmlSource(tree);
095        String[] splitSource = plantUmlSource.split("\\s", 2);
096        if (splitSource.length < 2) {
097            throw new IllegalArgumentException("Invalid " + getName()
098                + " tag: Expected filename and PlantUML source");
099        }
100
101        String packageName = extractPackageName(element);
102        FileObject graphicsFile;
103        try {
104            graphicsFile = fileManager.getFileForOutput(
105                DocumentationTool.Location.DOCUMENTATION_OUTPUT, packageName,
106                splitSource[0], null);
107        } catch (IOException e) {
108            throw new RuntimeException(
109                "Error generating output file for " + packageName + "/"
110                    + splitSource[0] + ": " + e.getLocalizedMessage());
111        }
112
113        // render
114        plantUmlSource = "@startuml\n" + splitSource[1].trim() + "\n@enduml";
115        SourceStringReader reader = new SourceStringReader(
116            Defines.createEmpty(), plantUmlSource, plantConfig());
117        try {
118            reader.outputImage(graphicsFile.openOutputStream(),
119                new FileFormatOption(getFileFormat(splitSource[0])));
120        } catch (IOException e) {
121            throw new RuntimeException(
122                "Error generating UML image " + graphicsFile.getName() + ": "
123                    + e.getLocalizedMessage());
124        }
125
126    }
127
128    private String extractPlantUmlSource(DocTree tagTree) {
129        StringBuilder source = new StringBuilder();
130        new SimpleDocTreeVisitor<>() {
131
132            @Override
133            public Object visitUnknownBlockTag(UnknownBlockTagTree node,
134                    Object p) {
135                new SimpleDocTreeVisitor<>() {
136
137                    @Override
138                    public Object visitText(TextTree node, Object p) {
139                        source.append(node.getBody());
140                        return null;
141                    }
142
143                    @Override
144                    public Object visitComment(CommentTree node, Object p) {
145                        String comment = node.getBody();
146                        source
147                            .append(comment.substring(4, comment.length() - 3));
148                        return null;
149                    }
150
151                    @Override
152                    public Object visitErroneous(ErroneousTree node, Object p) {
153                        source.append(node.getBody());
154                        return null;
155                    }
156
157                }.visit(node.getContent(), null);
158                return null;
159            }
160
161            @Override
162            public Object visitText(TextTree node, Object p) {
163                source.append(node.getBody());
164                return null;
165            }
166
167        }.visit(tagTree, null);
168        return source.toString();
169    }
170
171    private String extractPackageName(Element element) {
172        return element.accept(new SimpleElementVisitor8<>() {
173
174            @Override
175            public Name visitPackage(PackageElement e, Object p) {
176                return e.getQualifiedName();
177            }
178
179            @Override
180            public Name visitType(TypeElement e, Object p) {
181                return env.getElementUtils().getPackageOf(e).getQualifiedName();
182            }
183
184        }, null).toString();
185    }
186
187    private FileFormat getFileFormat(String name) {
188        for (FileFormat f : FileFormat.values()) {
189            if (name.toLowerCase().endsWith(f.getFileSuffix())) {
190                return f;
191            }
192        }
193        String msg = "Unsupported file extension: " + name;
194        throw new IllegalArgumentException(msg);
195    }
196
197    private List<String> plantConfig() {
198        if (plantConfigData != null) {
199            return plantConfigData;
200        }
201        plantConfigData = new ArrayList<>();
202        String configFileName
203            = System.getProperty(getClass().getName() + ".config");
204        if (configFileName == null) {
205            return plantConfigData;
206        }
207        try {
208            plantConfigData = Collections.singletonList(
209                new String(Files.readAllBytes(Paths.get(configFileName))));
210        } catch (IOException e) {
211            throw new RuntimeException(
212                "Error loading PlantUML configuration file "
213                    + configFileName + ": " + e.getLocalizedMessage());
214        }
215
216        return plantConfigData;
217    }
218
219}