/*
 * ProGuard -- shrinking, optimization, obfuscation, and preverification
 *             of Java bytecode.
 *
 * Copyright (c) 2002-2020 Guardsquare NV
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation; either version 2 of the License, or (at your option)
 * any later version.
 *
 * This program 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 for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
 */

package proguard.optimize.peephole;

import proguard.classfile.*;
import proguard.classfile.attribute.Attribute;
import proguard.classfile.attribute.CodeAttribute;
import proguard.classfile.attribute.visitor.AttributeVisitor;
import proguard.classfile.constant.ClassConstant;
import proguard.classfile.constant.Constant;
import proguard.classfile.constant.FieldrefConstant;
import proguard.classfile.constant.MethodrefConstant;
import proguard.classfile.constant.visitor.ConstantVisitor;
import proguard.classfile.editor.CodeAttributeEditor;
import proguard.classfile.instruction.*;
import proguard.classfile.instruction.visitor.InstructionVisitor;
import proguard.classfile.visitor.ClassVisitor;
import proguard.optimize.info.WrapperClassMarker;

/**
 * This AttributeVisitor simplifies the use of retargeted wrapper classes in
 * the code attributes that it visits. More specifically, it replaces
 *     "new Wrapper(targetClass)" by "targetClass", and
 *     "wrapper.field" by "wrapper".
 * You still need to retarget all class references, replacing references to
 * the wrapper class by references to the target class.
 *
 * @see WrapperClassMarker
 * @see ClassMerger
 * @see RetargetedClassFilter
 * @see TargetClassChanger
 * @author Eric Lafortune
 */
public class WrapperClassUseSimplifier
implements   AttributeVisitor,
             InstructionVisitor,
             ConstantVisitor,
             ClassVisitor
{
    //*
    private static final boolean DEBUG = false;
    /*/
    private static       boolean DEBUG = System.getProperty("wc") != null;
    //*/


    private final InstructionVisitor extraInstructionVisitor;

    private final CodeAttributeEditor codeAttributeEditor = new CodeAttributeEditor(true, true);

    // Fields acting as parameters and return values for the visitor methods.
    private Clazz wrappedClass;
    private Instruction popInstruction;


    /**
     * Creates a new WrapperClassUseSimplifier.
     */
    public WrapperClassUseSimplifier()
    {
        this(null);
    }


    /**
     * Creates a new WrapperClassUseSimplifier.
     * @param extraInstructionVisitor an optional extra visitor for all
     *                                simplified instructions.
     */
    public WrapperClassUseSimplifier(InstructionVisitor extraInstructionVisitor)
    {
        this.extraInstructionVisitor = extraInstructionVisitor;
    }


    // Implementations for AttributeVisitor.

    public void visitAnyAttribute(Clazz clazz, Attribute attribute) {}


    public void visitCodeAttribute(Clazz clazz, Method method, CodeAttribute codeAttribute)
    {
        if (DEBUG)
        {
            System.out.println("WrapperClassUseSimplifier: "+clazz.getName()+"."+method.getName(clazz)+method.getDescriptor(clazz));
        }

        int codeLength = codeAttribute.u4codeLength;

        // Reset the code changes.
        codeAttributeEditor.reset(codeLength);

        // Edit the instructions.
        codeAttribute.instructionsAccept(clazz, method, this);

        // Apply all accumulated changes to the code.
        codeAttributeEditor.visitCodeAttribute(clazz, method, codeAttribute);
    }


    // Implementations for InstructionVisitor.

    public void visitAnyInstruction(Clazz clazz, Method method, CodeAttribute codeAttribute, int offset, Instruction instruction) {}


    public void visitConstantInstruction(Clazz clazz, Method method, CodeAttribute codeAttribute, int offset, ConstantInstruction constantInstruction)
    {
        switch (constantInstruction.opcode)
        {
            case Instruction.OP_NEW:
            {
                // Is it instantiating a wrapper class?
                if (isReferencingWrapperClass(clazz, constantInstruction.constantIndex))
                {
                    // Is the next instruction a typical dup instruction?
                    int nextOffset = offset + constantInstruction.length(offset);
                    popInstruction = InstructionFactory.create(codeAttribute.code, nextOffset);
                    switch (popInstruction.canonicalOpcode())
                    {
                        case Instruction.OP_DUP:
                        {
                            // Delete the new and dup instructions:
                            //     new Wrapper
                            //     dup
                            codeAttributeEditor.deleteInstruction(offset);
                            codeAttributeEditor.deleteInstruction(nextOffset);
                            popInstruction = null;
                            break;
                        }
                        case Instruction.OP_ASTORE:
                        {
                            // Replace the new instance by a dummy null
                            // and remember to store the target instance:
                            //     new Wrapper -> aconst_null
                            //     astore x    -> remember
                            //     aload x
                            codeAttributeEditor.replaceInstruction(offset, new SimpleInstruction(Instruction.OP_ACONST_NULL));
                            break;
                        }
                        default:
                        {
                            // Replace the new instance by a dummy null
                            // and remember to pop the target instance:
                            //     new Wrapper -> aconst_null
                            codeAttributeEditor.replaceInstruction(offset, new SimpleInstruction(Instruction.OP_ACONST_NULL));
                            popInstruction = new SimpleInstruction(Instruction.OP_POP);
                        }
                    }

                    if (extraInstructionVisitor != null)
                    {
                        extraInstructionVisitor.visitConstantInstruction(clazz, method, codeAttribute, offset, constantInstruction);
                    }
                }
                break;
            }
            case Instruction.OP_INVOKESPECIAL:
            {
                // Is it initializing a wrapper class?
                if (isReferencingWrapperClass(clazz, constantInstruction.constantIndex))
                {
                    // Do we have a special pop instruction?
                    // TODO: May still fail with nested initializers.
                    if (popInstruction == null)
                    {
                        // Delete the initializer invocation (with the
                        // wrapper instance no longer on the stack):
                        //     Wrapper.<init>(target) -> target
                        codeAttributeEditor.deleteInstruction(offset);
                    }
                    else
                    {
                        // Delete the initializer invocation, and store
                        // the target instance again:
                        //     invokespecial Wrapper.<init>(target) -> astore x / pop
                        codeAttributeEditor.replaceInstruction(offset, new Instruction[]
                        {
                            popInstruction,
                            new SimpleInstruction(Instruction.OP_POP),
                        });
                    }
                }
                break;
            }
            case Instruction.OP_GETFIELD:
            {
                // Is it retrieving the field of the wrapper class?
                if (isReferencingWrapperClass(clazz, constantInstruction.constantIndex))
                {
                    // Delete the retrieval:
                    //     wrapper.field -> wrapper.
                    codeAttributeEditor.deleteInstruction(offset);
                }
                break;
            }
        }
    }


    // Implementations for ConstantVisitor.

    public void visitAnyConstant(Clazz clazz, Constant constant) {}


    public void visitFieldrefConstant(Clazz clazz, FieldrefConstant fieldrefConstant)
    {
        // Is the constant retrieving from a wrapper class?
        fieldrefConstant.referencedClassAccept(this);
    }


    public void visitMethodrefConstant(Clazz clazz, MethodrefConstant methodrefConstant)
    {
        if (methodrefConstant.getName(clazz).equals(ClassConstants.METHOD_NAME_INIT))
        {
            // Is the constant referring to a wrapper class?
            methodrefConstant.referencedClassAccept(this);
        }
    }


    public void visitClassConstant(Clazz clazz, ClassConstant classConstant)
    {
        // Is the constant referring to a wrapper class?
        classConstant.referencedClassAccept(this);
    }


    // Implementations for ClassVisitor.

    @Override
    public void visitAnyClass(Clazz clazz) { }


    @Override
    public void visitProgramClass(ProgramClass programClass)
    {
        wrappedClass = ClassMerger.getTargetClass(programClass);
    }


    // Small utility methods.

    /**
     * Returns whether the constant at the given offset is referencing a
     * wrapper class (different from the given class) that is being retargeted.
     */
    private boolean isReferencingWrapperClass(Clazz clazz, int constantIndex)
    {
        wrappedClass = null;

        clazz.constantPoolEntryAccept(constantIndex, this);

        return wrappedClass != null;
    }
}