8234689: facilitate writing additional custom attributes in a class file

Reviewed-by: jlahoda

src/jdk.compiler/share/classes/com/sun/tools/javac/jvm/ClassWriter.java

@@ -30,6 +30,7 @@ import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.Set;
 import java.util.LinkedHashSet;
+import java.util.function.Function;
 
 import javax.tools.JavaFileManager;
 import javax.tools.FileObject;
@@ -113,6 +114,8 @@ public class ClassWriter extends ClassFile {
      */
     public boolean multiModuleMode;
 
+    private List<Function<Symbol, Integer>> extraAttributeHooks = List.nil();
+
     /** The initial sizes of the data and constant pool buffers.
      *  Sizes are increased when buffers get full.
      */
@@ -121,7 +124,7 @@ public class ClassWriter extends ClassFile {
 
     /** An output buffer for member info.
      */
-    ByteBuffer databuf = new ByteBuffer(DATA_BUF_SIZE);
+    public ByteBuffer databuf = new ByteBuffer(DATA_BUF_SIZE);
 
     /** An output buffer for the constant pool.
      */
@@ -188,6 +191,10 @@ public class ClassWriter extends ClassFile {
         }
     }
 
+    public void addExtraAttributes(Function<Symbol, Integer> addExtraAttributes) {
+        extraAttributeHooks = extraAttributeHooks.prepend(addExtraAttributes);
+    }
+
 /******************************************************************
  * Diagnostics: dump generated class names and modifiers
  ******************************************************************/
@@ -276,7 +283,7 @@ public class ClassWriter extends ClassFile {
     /** Write header for an attribute to data buffer and return
      *  position past attribute length index.
      */
-    int writeAttr(Name attrName) {
+    public int writeAttr(Name attrName) {
         int index = poolWriter.putName(attrName);
         databuf.appendChar(index);
         databuf.appendInt(0);
@@ -285,7 +292,7 @@ public class ClassWriter extends ClassFile {
 
     /** Fill in attribute length.
      */
-    void endAttr(int index) {
+    public void endAttr(int index) {
         putInt(databuf, index - 4, databuf.length - index);
     }
 
@@ -942,6 +949,7 @@ public class ClassWriter extends ClassFile {
             acount++;
         }
         acount += writeMemberAttrs(v, false);
+        acount += writeExtraAttributes(v);
         endAttrs(acountIdx, acount);
     }
 
@@ -988,6 +996,7 @@ public class ClassWriter extends ClassFile {
         acount += writeMemberAttrs(m, false);
         if (!m.isLambdaMethod())
             acount += writeParameterAttrs(m.params);
+        acount += writeExtraAttributes(m);
         endAttrs(acountIdx, acount);
     }
 
@@ -1605,6 +1614,7 @@ public class ClassWriter extends ClassFile {
             acount += writeFlagAttrs(c.owner.flags() & ~DEPRECATED);
         }
         acount += writeExtraClassAttributes(c);
+        acount += writeExtraAttributes(c);
 
         poolbuf.appendInt(JAVA_MAGIC);
         if (preview.isEnabled()) {
@@ -1643,14 +1653,26 @@ public class ClassWriter extends ClassFile {
         poolWriter.reset(); // to save space
 
         out.write(databuf.elems, 0, databuf.length);
-     }
+    }
 
-    /**Allows subclasses to write additional class attributes
+     /**Allows subclasses to write additional class attributes
+      *
+      * @return the number of attributes written
+      */
+    protected int writeExtraClassAttributes(ClassSymbol c) {
+        return 0;
+    }
+
+    /**Allows friends to write additional attributes
      *
      * @return the number of attributes written
      */
-    protected int writeExtraClassAttributes(ClassSymbol c) {
-        return 0;
+    protected int writeExtraAttributes(Symbol sym) {
+        int i = 0;
+        for (Function<Symbol, Integer> hook : extraAttributeHooks) {
+            i += hook.apply(sym);
+        }
+        return i;
     }
 
     int adjustFlags(final long flags) {

test/langtools/tools/javac/classwriter/ExtraAttributes.java

@@ -0,0 +1,155 @@
+/*
+ * Copyright (c) 2019, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+/*
+ * @test
+ * @bug 8234689
+ * @summary facilitate writing additional custom attributes in a class file
+ * @library /tools/lib
+ * @modules jdk.compiler/com.sun.tools.javac.api
+ *          jdk.compiler/com.sun.tools.javac.code
+ *          jdk.compiler/com.sun.tools.javac.jvm
+ *          jdk.compiler/com.sun.tools.javac.main
+ *          jdk.compiler/com.sun.tools.javac.util
+ *          jdk.jdeps/com.sun.tools.javap
+ * @build toolbox.JarTask toolbox.JavacTask toolbox.JavapTask toolbox.ToolBox
+ * @run main ExtraAttributes
+ */
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+
+import com.sun.source.util.JavacTask;
+import com.sun.source.util.Plugin;
+
+import com.sun.tools.javac.api.BasicJavacTask;
+import com.sun.tools.javac.code.Symbol;
+import com.sun.tools.javac.jvm.ClassWriter;
+import com.sun.tools.javac.util.Context;
+import com.sun.tools.javac.util.Name;
+import com.sun.tools.javac.util.Names;
+
+import toolbox.JarTask;
+import toolbox.JavapTask;
+import toolbox.Task;
+import toolbox.ToolBox;
+
+
+public class ExtraAttributes implements Plugin {
+    public static void main(String... args) throws Exception {
+        new ExtraAttributes().run();
+    }
+
+    void run() throws Exception {
+        ToolBox tb = new ToolBox();
+        Path pluginClasses = Path.of("plugin-classes");
+        tb.writeFile(pluginClasses.resolve("META-INF").resolve("services").resolve(Plugin.class.getName()),
+                ExtraAttributes.class.getName() + "\n");
+        Files.copy(Path.of(ToolBox.testClasses).resolve("ExtraAttributes.class"),
+                pluginClasses.resolve("ExtraAttributes.class"));
+
+        Path pluginJar = Path.of("plugin.jar");
+        new JarTask(tb, pluginJar)
+                .baseDir(pluginClasses)
+                .files(".")
+                .run();
+
+        Path src = Path.of("src");
+            tb.writeJavaFiles(src,
+                    "public class HelloWorld {\n"
+                    + "    public static String message = \"Hello World!\";\n"
+                    + "    public static void main(String... args) {\n"
+                    + "        System.out.println(message);\n"
+                    + "    }\n"
+                    + "}\n");
+
+        List<String> stdout = new toolbox.JavacTask(tb)
+                .classpath(pluginJar)
+                .options("-Xplugin:ExtraAttributes")
+                .outdir(Files.createDirectories(Path.of("classes")))
+                .files(tb.findJavaFiles(src))
+                .run()
+                .writeAll()
+                .getOutputLines(Task.OutputKind.STDOUT);
+
+        // cannot rely on order of output, so sort it
+        stdout.sort(CharSequence::compare);
+
+        tb.checkEqual(stdout,
+                List.of(
+                        "Add attributes for <clinit>()",
+                        "Add attributes for HelloWorld",
+                        "Add attributes for HelloWorld()",
+                        "Add attributes for main(java.lang.String...)",
+                        "Add attributes for message"
+                ));
+
+        List<String> lines = new JavapTask(tb)
+                .options("-p",
+                        "-v",
+                        Path.of("classes").resolve("HelloWorld.class").toString())
+                .run()
+                .getOutputLines(Task.OutputKind.DIRECT);
+
+        long attrs = lines.stream()
+                .filter(s -> s.contains("testAttr:"))
+                .count();
+        if (attrs != 5) {
+            throw new Exception("expected attributes not found; expected: 5; found: " + attrs);
+        }
+    }
+
+    // Plugin impl...
+
+    private ClassWriter classWriter;
+    private Names names;
+
+    @Override
+    public String getName() { return "ExtraAttributes"; }
+
+    @Override
+    public void init(JavacTask task, String... args) {
+        Context c = ((BasicJavacTask) task).getContext();
+        classWriter = ClassWriter.instance(c);
+        names = Names.instance(c);
+
+        // register callback
+        classWriter.addExtraAttributes(this::addExtraAttributes);
+    }
+
+    @Override
+    public boolean autoStart() {
+        return true;
+    }
+
+    private int addExtraAttributes(Symbol sym) {
+        System.out.println("Add attributes for " + sym);
+        Name testAttr = names.fromString("testAttr");
+        int alenIdx = classWriter.writeAttr(testAttr);
+        classWriter.databuf.appendChar(42);
+        classWriter.endAttr(alenIdx);
+        return 1;
+    }
+}
+