import java.io.IOException;
import java.util.Iterator;
 
import org.apache.bcel.Repository;
import org.apache.bcel.classfile.ClassFormatException;
import org.apache.bcel.classfile.Method;
import org.apache.bcel.generic.*;
 
/**
 * Instruments byte-code by adding: <code>
 * System.out.println("About to call: METHOD_NAME METHOD_SIGNATURE");
 * </code>
 * before method invocation.
 *
 * @author Radoslaw Urbas
 * @since Oct 27, 2007
 *
 */
public class Transform
{
 
    private ClassGen classGen = null;
 
    /**
     *
     * @param className name of class to instrument
     * @throws ClassNotFoundException
     * @throws ClassFormatException
     * @throws IOException
     */
    public Transform(String className) throws ClassNotFoundException,
            ClassFormatException, IOException
    {
        this.classGen = new ClassGen(Repository.lookupClass(className));
    }
 
    /**
     * Instruments the class by displaying method information before invocation.
     *
     * @return instrumented ClassGen object
     */
    private ClassGen instrument()
    {
        ConstantPoolGen pgen = classGen.getConstantPool();
        String cname = classGen.getClassName();
        for (Method method : classGen.getMethods())
        {
            MethodGen methgen = instrumentMethod(pgen, cname, method);
            classGen.replaceMethod(method, methgen.getMethod());
        }
        return classGen;
    }
 
    /**
     * Instruments method by displaying information before method invocation.
     *
     * @param constantPoolGen constant pool for class containing method to
     *            instrument
     * @param className name of the class containing method to instrument
     * @param method method to instrument
     * @return instrumented MethodGen object
     */
    private MethodGen instrumentMethod(ConstantPoolGen constantPoolGen,
            String className, Method method)
    {
        InstructionFactory instructionFactory = new InstructionFactory(classGen);
        MethodGen methgen = new MethodGen(method, className, constantPoolGen);
        InstructionList originalInstructionList = methgen.getInstructionList();
        Iterator instructionIterator = originalInstructionList.iterator();
        while (instructionIterator.hasNext())
        {
            InstructionHandle ih = (InstructionHandle) instructionIterator
                    .next();
            if (ih.getInstruction() instanceof InvokeInstruction)
            {
                originalInstructionList.insert(ih, instructionFactory
                        .createPrintln(getMessage((InvokeInstruction) ih
                                .getInstruction(), constantPoolGen)));
            }
        }
        methgen.setMaxStack();
        return methgen;
    }
 
    /**
     *
     * @param invokeInstruction object representing invoke instruction
     * @param constantPoolGen
     * @return message based on method name and signature
     */
    private String getMessage(InvokeInstruction invokeInstruction,
            ConstantPoolGen constantPoolGen)
    {
        String text = "About to call: "
                + invokeInstruction.getMethodName(constantPoolGen)
                + invokeInstruction.getSignature(constantPoolGen);
        return text;
    }
 
    /**
     *
     * @param args requires one parameter: class name. Class should be available
     *            in directory ./classes
     * @throws ClassNotFoundException
     * @throws IOException
     */
    public static void main(String[] args) throws ClassNotFoundException,
            IOException
    {
        String className = args[0];
        Transform transform = new Transform(className);
        ClassGen classGen = transform.instrument();
        classGen.getJavaClass().dump(className + ".class");
    }
}