--- /dev/null 2016-05-31 09:42:47.975716356 -0700 +++ new/src/jdk.vm.compiler/share/classes/org.graalvm.compiler.api.directives.test/src/org/graalvm/compiler/api/directives/test/TinyInstrumentor.java 2016-12-09 00:46:40.715882509 -0800 @@ -0,0 +1,242 @@ +/* + * Copyright (c) 2016, 2016, 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. + */ +package org.graalvm.compiler.api.directives.test; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; + +import org.graalvm.compiler.test.ExportingClassLoader; + +import jdk.internal.org.objectweb.asm.ClassReader; +import jdk.internal.org.objectweb.asm.ClassWriter; +import jdk.internal.org.objectweb.asm.Label; +import jdk.internal.org.objectweb.asm.MethodVisitor; +import jdk.internal.org.objectweb.asm.Opcodes; +import jdk.internal.org.objectweb.asm.tree.AbstractInsnNode; +import jdk.internal.org.objectweb.asm.tree.ClassNode; +import jdk.internal.org.objectweb.asm.tree.IincInsnNode; +import jdk.internal.org.objectweb.asm.tree.InsnList; +import jdk.internal.org.objectweb.asm.tree.JumpInsnNode; +import jdk.internal.org.objectweb.asm.tree.LabelNode; +import jdk.internal.org.objectweb.asm.tree.LineNumberNode; +import jdk.internal.org.objectweb.asm.tree.MethodNode; +import jdk.internal.org.objectweb.asm.tree.VarInsnNode; + +/** + * The {@code TinyInstrumentor} is a bytecode instrumentor using ASM bytecode manipulation + * framework. It injects given code snippet into a target method and creates a temporary class as + * the container. Because the target method is cloned into the temporary class, it is required that + * the target method is public static. Any referred method/field in the target method or the + * instrumentation snippet should be made public as well. + */ +public class TinyInstrumentor implements Opcodes { + + private InsnList instrumentationInstructions; + private int instrumentationMaxLocal; + + /** + * Create a instrumentor with a instrumentation snippet. The snippet is specified with the given + * class {@code instrumentationClass} and the given method name {@code methodName}. + */ + public TinyInstrumentor(Class instrumentationClass, String methodName) throws IOException { + MethodNode instrumentationMethod = getMethodNode(instrumentationClass, methodName); + assert instrumentationMethod != null; + assert (instrumentationMethod.access | ACC_STATIC) != 0; + assert "()V".equals(instrumentationMethod.desc); + instrumentationInstructions = cloneInstructions(instrumentationMethod.instructions); + instrumentationMaxLocal = instrumentationMethod.maxLocals; + // replace return instructions with a goto unless there is a single return at the end. In + // that case, simply remove the return. + List returnInstructions = new ArrayList<>(); + for (AbstractInsnNode instruction : selectAll(instrumentationInstructions)) { + if (instruction instanceof LineNumberNode) { + instrumentationInstructions.remove(instruction); + } else if (instruction.getOpcode() == RETURN) { + returnInstructions.add(instruction); + } + } + LabelNode exit = new LabelNode(); + if (returnInstructions.size() == 1) { + AbstractInsnNode returnInstruction = returnInstructions.get(0); + if (instrumentationInstructions.getLast() != returnInstruction) { + instrumentationInstructions.insertBefore(returnInstruction, new JumpInsnNode(GOTO, exit)); + } + instrumentationInstructions.remove(returnInstruction); + } else { + for (AbstractInsnNode returnInstruction : returnInstructions) { + instrumentationInstructions.insertBefore(returnInstruction, new JumpInsnNode(GOTO, exit)); + instrumentationInstructions.remove(returnInstruction); + } + } + instrumentationInstructions.add(exit); + } + + /** + * @return a {@link MethodNode} called {@code methodName} in the given class. + */ + private static MethodNode getMethodNode(Class clazz, String methodName) throws IOException { + ClassReader classReader = new ClassReader(clazz.getName()); + ClassNode classNode = new ClassNode(); + classReader.accept(classNode, ClassReader.SKIP_FRAMES); + + for (MethodNode methodNode : classNode.methods) { + if (methodNode.name.equals(methodName)) { + return methodNode; + } + } + return null; + } + + /** + * Create a {@link ClassNode} with empty constructor. + */ + private static ClassNode emptyClass(String name) { + ClassNode classNode = new ClassNode(); + classNode.visit(52, ACC_SUPER | ACC_PUBLIC, name.replace('.', '/'), null, "java/lang/Object", new String[]{}); + + MethodVisitor mv = classNode.visitMethod(ACC_PUBLIC, "", "()V", null, null); + mv.visitCode(); + Label l0 = new Label(); + mv.visitLabel(l0); + mv.visitVarInsn(ALOAD, 0); + mv.visitMethodInsn(INVOKESPECIAL, "java/lang/Object", "", "()V", false); + mv.visitInsn(RETURN); + Label l1 = new Label(); + mv.visitLabel(l1); + mv.visitMaxs(1, 1); + mv.visitEnd(); + + return classNode; + } + + /** + * Helper method for iterating the given {@link InsnList}. + */ + private static Iterable selectAll(InsnList instructions) { + return new Iterable() { + @Override + public Iterator iterator() { + return instructions.iterator(); + } + }; + } + + /** + * Make a clone of the given {@link InsnList}. + */ + private static InsnList cloneInstructions(InsnList instructions) { + Map labelMap = new HashMap<>(); + for (AbstractInsnNode instruction : selectAll(instructions)) { + if (instruction instanceof LabelNode) { + LabelNode clone = new LabelNode(new Label()); + LabelNode original = (LabelNode) instruction; + labelMap.put(original, clone); + } + } + InsnList clone = new InsnList(); + for (AbstractInsnNode insn : selectAll(instructions)) { + clone.add(insn.clone(labelMap)); + } + return clone; + } + + /** + * Shifts all local variable slot references by a specified amount. + */ + private static void shiftLocalSlots(InsnList instructions, int offset) { + for (AbstractInsnNode insn : selectAll(instructions)) { + if (insn instanceof VarInsnNode) { + VarInsnNode varInsn = (VarInsnNode) insn; + varInsn.var += offset; + + } else if (insn instanceof IincInsnNode) { + IincInsnNode iincInsn = (IincInsnNode) insn; + iincInsn.var += offset; + } + } + } + + /** + * Instrument the target method specified by the class {@code targetClass} and the method name + * {@code methodName}. For each occurrence of the {@code opcode}, the instrumentor injects a + * copy of the instrumentation snippet. + */ + public Class instrument(Class targetClass, String methodName, int opcode) throws IOException, ClassNotFoundException { + return instrument(targetClass, methodName, opcode, true); + } + + public Class instrument(Class targetClass, String methodName, int opcode, boolean insertAfter) throws IOException, ClassNotFoundException { + // create a container class + String className = targetClass.getName() + "$$" + methodName; + ClassNode classNode = emptyClass(className); + // duplicate the target method and add to the container class + MethodNode methodNode = getMethodNode(targetClass, methodName); + MethodNode newMethodNode = new MethodNode(methodNode.access, methodNode.name, methodNode.desc, methodNode.signature, methodNode.exceptions.toArray(new String[methodNode.exceptions.size()])); + methodNode.accept(newMethodNode); + classNode.methods.add(newMethodNode); + // perform bytecode instrumentation + for (AbstractInsnNode instruction : selectAll(newMethodNode.instructions)) { + if (instruction.getOpcode() == opcode) { + InsnList instrumentation = cloneInstructions(instrumentationInstructions); + shiftLocalSlots(instrumentation, newMethodNode.maxLocals); + newMethodNode.maxLocals += instrumentationMaxLocal; + if (insertAfter) { + newMethodNode.instructions.insert(instruction, instrumentation); + } else { + newMethodNode.instructions.insertBefore(instruction, instrumentation); + } + } + } + // dump a byte array and load the class with a dedicated loader to separate the namespace + ClassWriter classWriter = new ClassWriter(ClassWriter.COMPUTE_FRAMES); + classNode.accept(classWriter); + byte[] bytes = classWriter.toByteArray(); + return new Loader(className, bytes).findClass(className); + } + + private static class Loader extends ExportingClassLoader { + + private String className; + private byte[] bytes; + + Loader(String className, byte[] bytes) { + super(TinyInstrumentor.class.getClassLoader()); + this.className = className; + this.bytes = bytes; + } + + @Override + protected Class findClass(String name) throws ClassNotFoundException { + if (name.equals(className)) { + return defineClass(name, bytes, 0, bytes.length); + } else { + return super.findClass(name); + } + } + } + +}