/*
 * ProGuard -- shrinking, optimization, obfuscation, and preverification
 *             of Java bytecode.
 *
 * Copyright (c) 2002-2017 Eric Lafortune @ GuardSquare
 *
 * 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.preverify;

import proguard.classfile.*;
import proguard.classfile.attribute.*;
import proguard.classfile.attribute.preverification.*;
import proguard.classfile.attribute.visitor.AttributeVisitor;
import proguard.classfile.editor.*;
import proguard.classfile.instruction.InstructionConstants;
import proguard.classfile.util.SimplifiedVisitor;
import proguard.classfile.visitor.ClassPrinter;
import proguard.evaluation.*;
import proguard.evaluation.value.*;
import proguard.optimize.evaluation.*;

import java.util.*;

/**
 * This class can preverify methods in program class pools, according to a given
 * specification.
 *
 * @author Eric Lafortune
 */
public class CodePreverifier
extends      SimplifiedVisitor
implements   AttributeVisitor
{
    //*
    private static final boolean DEBUG = false;
    /*/
    private static       boolean DEBUG = true;
    //*/


    private final boolean microEdition;

    private final PartialEvaluator    partialEvaluator    = new PartialEvaluator();
    private final LivenessAnalyzer    livenessAnalyzer    = new LivenessAnalyzer(partialEvaluator);
    private final CodeAttributeEditor codeAttributeEditor = new CodeAttributeEditor();


    /**
     * Creates a new CodePreverifier.
     */
    public CodePreverifier(boolean microEdition)
    {
        this.microEdition = microEdition;
    }


    // Implementations for AttributeVisitor.

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


    public void visitCodeAttribute(Clazz clazz, Method method, CodeAttribute codeAttribute)
    {
        // TODO: Remove this when the preverifier has stabilized.
        // Catch any unexpected exceptions from the actual visiting method.
        try
        {
            // Process the code.
            visitCodeAttribute0(clazz, method, codeAttribute);
        }
        catch (RuntimeException ex)
        {
            System.err.println("Unexpected error while preverifying:");
            System.err.println("  Class       = ["+clazz.getName()+"]");
            System.err.println("  Method      = ["+method.getName(clazz)+method.getDescriptor(clazz)+"]");
            System.err.println("  Exception   = ["+ex.getClass().getName()+"] ("+ex.getMessage()+")");

            throw ex;
        }
    }


    public void visitCodeAttribute0(Clazz clazz, Method method, CodeAttribute codeAttribute)
    {
//        DEBUG =
//            clazz.getName().equals("abc/Def") &&
//            method.getName(clazz).equals("abc");

        ProgramClass  programClass  = (ProgramClass)clazz;
        ProgramMethod programMethod = (ProgramMethod)method;

        int codeLength = codeAttribute.u4codeLength;

        // Evaluate the method.
        //partialEvaluator.visitCodeAttribute(clazz, method, codeAttribute);
        livenessAnalyzer.visitCodeAttribute(clazz, method, codeAttribute);

        // We may have to remove unreachable code.
        codeAttributeEditor.reset(codeLength);

        // Collect the stack map frames.
        List stackMapFrameList = new ArrayList();

        for (int offset = 0; offset < codeLength; offset++)
        {
            // Only store frames at the beginning of code blocks.
            if (!partialEvaluator.isTraced(offset))
            {
                // Mark the unreachable instruction for deletion.
                codeAttributeEditor.deleteInstruction(offset);
            }
            else if (partialEvaluator.isBranchOrExceptionTarget(offset))
            {
                // Convert the variable values to types.
                VerificationType[] variableTypes =
                    correspondingVerificationTypes(programClass,
                                                   programMethod,
                                                   codeAttribute,
                                                   offset,
                                                   partialEvaluator.getVariablesBefore(offset));

                // Convert the stack values to types.
                VerificationType[] stackTypes =
                    correspondingVerificationTypes(programClass,
                                                   programMethod,
                                                   codeAttribute,
                                                   offset,
                                                   partialEvaluator.getStackBefore(offset));
                // Create and store a new frame.
                stackMapFrameList.add(new FullFrame(offset,
                                                    variableTypes,
                                                    stackTypes));
            }
        }

        // Compress the stack map frames if the target is not Java Micro Edition.
        if (!microEdition && !stackMapFrameList.isEmpty())
        {
            // Convert the initial variable values to types.
            VerificationType[] initialVariables =
                correspondingVerificationTypes(programClass,
                                               programMethod,
                                               codeAttribute,
                                               PartialEvaluator.AT_METHOD_ENTRY,
                                               partialEvaluator.getVariablesBefore(0));

            // Special case: the <init> method.
            if (method.getName(programClass).equals(ClassConstants.METHOD_NAME_INIT))
            {
                initialVariables[0] = VerificationTypeFactory.createUninitializedThisType();
            }

            compressStackMapFrames(initialVariables,
                                   stackMapFrameList);
        }

        // Get the proper name for the attribute to be added/replaced/deleted.
        String stackMapAttributeName = microEdition ?
             ClassConstants.ATTR_StackMap :
             ClassConstants.ATTR_StackMapTable;

        int frameCount = stackMapFrameList.size();

        if (DEBUG)
        {
            Attribute originalStackMapAttribute = codeAttribute.getAttribute(clazz,
                                                                             stackMapAttributeName);

            if (originalStackMapAttribute != null)
            {
                int originalFrameCount = microEdition ?
                    ((StackMapAttribute)originalStackMapAttribute).u2stackMapFramesCount :
                    ((StackMapTableAttribute)originalStackMapAttribute).u2stackMapFramesCount;

                StackMapFrame[] originalFrames = microEdition ?
                    ((StackMapAttribute)originalStackMapAttribute).stackMapFrames :
                    ((StackMapTableAttribute)originalStackMapAttribute).stackMapFrames;

                if (frameCount != originalFrameCount ||
                    !Arrays.equals(stackMapFrameList.toArray(), originalFrames))
                {
                    System.out.println("Original preverification ["+clazz.getName()+"]:");
                    new ClassPrinter().visitProgramMethod(programClass, programMethod);
                }
            }
            else if (frameCount != 0)
            {
                System.out.println("Original preverification empty ["+clazz.getName()+"."+method.getName(clazz)+"]");
            }
        }

        if (frameCount == 0)
        {
            // Remove any stack map (table) attribute from the code attribute.
            new AttributesEditor(programClass, programMethod, codeAttribute, true).deleteAttribute(stackMapAttributeName);
        }
        else
        {
            Attribute stackMapAttribute;

            // Create the appropriate attribute.
            if (microEdition)
            {
                // Copy the frames into an array.
                FullFrame[] stackMapFrames = new FullFrame[frameCount];
                stackMapFrameList.toArray(stackMapFrames);

                // Put the frames into a stack map attribute.
                stackMapAttribute = new StackMapAttribute(stackMapFrames);
            }
            else
            {
                // Copy the frames into an array.
                StackMapFrame[] stackMapFrames = new StackMapFrame[frameCount];
                stackMapFrameList.toArray(stackMapFrames);

                // Put the frames into a stack map table attribute.
                stackMapAttribute = new StackMapTableAttribute(stackMapFrames);
            }

            // Fill out the name of the stack map attribute.
            stackMapAttribute.u2attributeNameIndex =
                new ConstantPoolEditor(programClass).addUtf8Constant(stackMapAttributeName);

            // Add the new stack map (table) attribute to the code attribute.
            new AttributesEditor(programClass, programMethod, codeAttribute, true).addAttribute(stackMapAttribute);

            if (DEBUG)
            {
                System.out.println("Preverifier ["+programClass.getName()+"."+programMethod.getName(programClass)+"]:");
                stackMapAttribute.accept(programClass, programMethod, codeAttribute, new ClassPrinter());
            }
        }

        // Apply code modifications, deleting unreachable code.
        codeAttributeEditor.visitCodeAttribute(clazz, method, codeAttribute);
    }


    // Small utility methods.

    /**
     * Creates and returns the verification types corresponding to the given
     * variables. If necessary, class constants are added to the constant pool
     * of the given class.
     */
    private VerificationType[] correspondingVerificationTypes(ProgramClass    programClass,
                                                              ProgramMethod   programMethod,
                                                              CodeAttribute   codeAttribute,
                                                              int             offset,
                                                              TracedVariables variables)
    {
        int maximumVariablesSize = variables.size();
        int typeCount = 0;
        int typeIndex = 0;

        // Count the the number of verification types, ignoring any nulls at
        // the end.
        for (int index = 0; index < maximumVariablesSize; index++)
        {
            Value value = variables.getValue(index);

            typeIndex++;

            // Remember the maximum live type index.
            if (value != null &&
                (offset == PartialEvaluator.AT_METHOD_ENTRY ||
                 livenessAnalyzer.isAliveBefore(offset, index)))
            {
                typeCount = typeIndex;

                // Category 2 types that are alive are stored as single entries.
                if (value.isCategory2())
                {
                    index++;
                }
            }
        }

        // Create and fill out the verification types.
        VerificationType[] types = new VerificationType[typeCount];

        typeIndex = 0;

        // Note the slightly different terminating condition, because the
        // types may have been truncated.
        for (int index = 0; typeIndex < typeCount; index++)
        {
            Value value         = variables.getValue(index);
            Value producerValue = variables.getProducerValue(index);

            // Fill out the type.
            VerificationType type;

            if (value != null &&
                (offset == PartialEvaluator.AT_METHOD_ENTRY ||
                 livenessAnalyzer.isAliveBefore(offset, index)))
            {
                type = correspondingVerificationType(programClass,
                                                     programMethod,
                                                     codeAttribute,
                                                     offset,
                                                     index == 0,
                                                     value,
                                                     producerValue);

                // Category 2 types that are alive are stored as single entries.
                if (value.isCategory2())
                {
                    index++;
                }
            }
            else
            {
                type = VerificationTypeFactory.createTopType();
            }

            types[typeIndex++] = type;
        }

        return types;
    }


    /**
     * Creates and returns the verification types corresponding to the given
     * stack. If necessary, class constants are added to the constant pool
     * of the given class.
     */
    private VerificationType[] correspondingVerificationTypes(ProgramClass  programClass,
                                                              ProgramMethod programMethod,
                                                              CodeAttribute codeAttribute,
                                                              int           offset,
                                                              TracedStack   stack)
    {
        int maximumStackSize = stack.size();
        int typeCount = 0;

        // Count the the number of verification types.
        for (int index = 0; index < maximumStackSize; index++)
        {
            // We have to work down from the top of the stack.
            Value value = stack.getTop(index);

            typeCount++;

            // Category 2 types are stored as single entries.
            if (value.isCategory2())
            {
                index++;
            }
        }

        // Create and fill out the verification types.
        VerificationType[] types = new VerificationType[typeCount];

        int typeIndex = typeCount;

        for (int index = 0; index < maximumStackSize; index++)
        {
            // We have to work down from the top of the stack.
            Value value         = stack.getTop(index);
            Value producerValue = stack.getTopProducerValue(index);

            // Fill out the type.
            types[--typeIndex] =
                correspondingVerificationType(programClass,
                                              programMethod,
                                              codeAttribute,
                                              offset,
                                              false,
                                              value,
                                              producerValue);

            // Category 2 types are stored as single entries.
            if (value.isCategory2())
            {
                index++;
            }
        }

        return types;
    }


    /**
     * Creates and returns the verification type corresponding to the given
     * value. If necessary, a class constant is added to the constant pool of
     * the given class.
     */
    private VerificationType correspondingVerificationType(ProgramClass  programClass,
                                                           ProgramMethod programMethod,
                                                           CodeAttribute codeAttribute,
                                                           int           offset,
                                                           boolean       isVariable0,
                                                           Value         value,
                                                           Value         producerValue)
    {
        if (value == null)
        {
            return VerificationTypeFactory.createTopType();
        }

        int type = value.computationalType();

        switch (type)
        {
            case Value.TYPE_INSTRUCTION_OFFSET:
            case Value.TYPE_INTEGER:   return VerificationTypeFactory.createIntegerType();
            case Value.TYPE_LONG:      return VerificationTypeFactory.createLongType();
            case Value.TYPE_FLOAT:     return VerificationTypeFactory.createFloatType();
            case Value.TYPE_DOUBLE:    return VerificationTypeFactory.createDoubleType();
            case Value.TYPE_TOP:       return VerificationTypeFactory.createTopType();
            case Value.TYPE_REFERENCE:
                // Is it a Null type?
                ReferenceValue referenceValue = value.referenceValue();
                if (referenceValue.isNull() == Value.ALWAYS)
                {
                    return VerificationTypeFactory.createNullType();
                }

                // Does the reference type have a single producer?
                if (offset != PartialEvaluator.AT_METHOD_ENTRY)
                {
                    InstructionOffsetValue producers = producerValue.instructionOffsetValue();
                    if (producers.instructionOffsetCount() == 1)
                    {
                        int producerOffset = producers.instructionOffset(0);

                        // Follow any dup or swap instructions.
                        while (producerOffset != PartialEvaluator.AT_METHOD_ENTRY &&
                               isDupOrSwap(codeAttribute.code[producerOffset]))
                        {
                            producers      = partialEvaluator.getStackBefore(producerOffset).getTopProducerValue(0).instructionOffsetValue();
                            producerOffset = producers.minimumValue();
                        }

                        // Are we in an instance initialization method,
                        // before the super initialization, loading "this"?
                        if (partialEvaluator.isInitializer()                       &&
                            offset <= partialEvaluator.superInitializationOffset() &&
                            (isVariable0 ||
                             producerOffset > PartialEvaluator.AT_METHOD_ENTRY &&
                             codeAttribute.code[producerOffset] == InstructionConstants.OP_ALOAD_0))
                        {
                            // It's an UninitializedThis type.
                            return VerificationTypeFactory.createUninitializedThisType();
                        }

                        // Is the reference type newly created and still
                        // uninitialized?
                        if (producerOffset > PartialEvaluator.AT_METHOD_ENTRY &&
                            offset <= partialEvaluator.initializationOffset(producerOffset))
                        {
                            // It's an Uninitialized type.
                            return VerificationTypeFactory.createUninitializedType(producerOffset);
                        }
                    }
                }

                // It's an ordinary Object type.
                return VerificationTypeFactory.createObjectType(createClassConstant(programClass, referenceValue));
        }

        throw new IllegalArgumentException("Unknown computational type ["+type+"]");
    }


    /**
     * Finds or creates a class constant for the given reference value, and
     * returns its index in the constant pool.
     */
    private int createClassConstant(ProgramClass   programClass,
                                    ReferenceValue referenceValue)
    {
        return new ConstantPoolEditor(programClass).addClassConstant(referenceValue.getType(),
                                                                     referenceValue.getReferencedClass());
    }


    /**
     * Compresses the given list of full frames, for use in a stack map table.
     */
    private void compressStackMapFrames(VerificationType[] initialVariableTypes,
                                        List               stackMapFrameList)
    {
        int                previousVariablesCount = initialVariableTypes.length;
        VerificationType[] previousVariableTypes  = initialVariableTypes;

        int previousOffset = -1;

        for (int index = 0; index < stackMapFrameList.size(); index++)
        {
            FullFrame fullFrame = (FullFrame)stackMapFrameList.get(index);

            int                variablesCount = fullFrame.variablesCount;
            VerificationType[] variables      = fullFrame.variables;
            int                stackCount     = fullFrame.stackCount;
            VerificationType[] stack          = fullFrame.stack;

            // Start computing the compressed frame.
            // The default is the full frame.
            StackMapFrame compressedFrame = fullFrame;

            // Are all variables equal?
            if (variablesCount == previousVariablesCount &&
                equalVerificationTypes(variables, previousVariableTypes, variablesCount))
            {
                // Are the stacks equal?
                //if (stackCount == previousStackCount &&
                //    equalVerificationTypes(stack, previousStack, stackCount))
                //{
                //    // Remove the identical frame.
                //    stackMapFrameList.remove(index--);
                //
                //    // Move on to the next frame (at the same index).
                //    continue;
                //}
                // Is the new stack empty?
                //else
                if (stackCount == 0)
                {
                    compressedFrame = new SameZeroFrame();
                }
                // Does the new stack contain a single element?
                else if (stackCount == 1)
                {
                    compressedFrame = new SameOneFrame(stack[0]);
                }
            }
            // Is the stack empty?
            else if (stackCount == 0)
            {
                int additionalVariablesCount = variablesCount - previousVariablesCount;

                // Are the variables chopped?
                if (additionalVariablesCount < 0  &&
                    additionalVariablesCount > -4 &&
                    equalVerificationTypes(variables, previousVariableTypes, variablesCount))
                {
                    compressedFrame = new LessZeroFrame((byte)-additionalVariablesCount);
                }
                // Are the variables extended?
                else if (//previousVariablesCount   > 0 &&
                         additionalVariablesCount > 0 &&
                         additionalVariablesCount < 4 &&
                         equalVerificationTypes(variables, previousVariableTypes, previousVariablesCount))
                {
                    // Copy the additional variables into an array.
                    VerificationType[] additionalVariables = new VerificationType[additionalVariablesCount];
                    System.arraycopy(variables, variablesCount - additionalVariablesCount,
                                     additionalVariables, 0,
                                     additionalVariablesCount);

                    compressedFrame = new MoreZeroFrame(additionalVariables);
                }
            }

            // Compress the instruction offset.
            int offset = fullFrame.u2offsetDelta;
            compressedFrame.u2offsetDelta = offset - previousOffset - 1;
            previousOffset = offset;

            // Remember this frame.
            previousVariablesCount = fullFrame.variablesCount;
            previousVariableTypes     = fullFrame.variables;

            // Replace the full frame.
            stackMapFrameList.set(index, compressedFrame);
        }
    }


    /**
     * Returns whether the given arrays of verification types are equal, up to
     * the given length.
     */
    private boolean equalVerificationTypes(VerificationType[] types1,
                                           VerificationType[] types2,
                                           int                length)
    {
        if (length > 0 &&
            (types1.length < length ||
             types2.length < length))
        {
            return false;
        }

        for (int index = 0; index < length; index++)
        {
            if (!types1[index].equals(types2[index]))
            {
                return false;
            }
        }

        return true;
    }


    /**
     * Returns whether the given instruction opcode represents a dup or swap
     * instruction (dup, dup_x1, dup_x2, dup2, dup2_x1, dup2_x2, swap).
     */
    private boolean isDupOrSwap(int opcode)
    {
        return opcode >= InstructionConstants.OP_DUP &&
               opcode <= InstructionConstants.OP_SWAP;
    }
}
