/*
 * ProGuard -- shrinking, optimization, obfuscation, and preverification
 *             of Java bytecode.
 *
 * Copyright (c) 2002-2018 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.classfile.util;

import proguard.classfile.*;
import proguard.classfile.attribute.CodeAttribute;
import proguard.classfile.constant.Constant;
import proguard.classfile.instruction.*;
import proguard.evaluation.TracedStack;
import proguard.evaluation.value.*;
import proguard.optimize.evaluation.PartialEvaluator;
import proguard.optimize.peephole.InstructionSequenceReplacer;

/**
 * This class finds sequences of instructions that correspond to primitive
 * array initializations. Such initializations may be represented more
 * efficiently in other bytecode languages.
 *
 * @author Eric Lafortune
 */
public class ArrayInitializationMatcher
{
    private static final int X = InstructionSequenceReplacer.X;

    private final PartialEvaluator partialEvaluator;

    private int    arrayInitializationStart;
    private int    arrayInitializationEnd;
    private Object array;

    private final Constant[] CONSTANTS = new Constant[0];

    // Conversion with dex2jar might result in arrays
    // being pre-stored to a variable before initialization:
    //
    //   newarray
    //   astore X
    //   aload  X
    //   initialization start
    //
    private final Instruction[] ARRAY_PRESTORE_INSTRUCTIONS = new Instruction[]
        {
            new VariableInstruction(InstructionConstants.OP_ASTORE, X),
            new VariableInstruction(InstructionConstants.OP_ALOAD, X)
        };

    private final InstructionSequenceMatcher arrayPreStoreMatcher =
        new InstructionSequenceMatcher(CONSTANTS, ARRAY_PRESTORE_INSTRUCTIONS);


    /**
     * Creates a new ArrayInitializationMatcher.
     */
    public ArrayInitializationMatcher()
    {
        this(new PartialEvaluator());
    }


    /**
     * Creates a new ArrayInitializationMatcher that will use the given partial
     * evaluator. The evaluator must have been initialized before trying to
     * match any array initializations.
     * @param partialEvaluator the evaluator to be used for the analysis.
     */
    public ArrayInitializationMatcher(PartialEvaluator partialEvaluator)
    {
        this.partialEvaluator = partialEvaluator;
    }


    /**
     * Returns whether the code fragment starting at the specified newarray
     * instruction is followed by a static array initialization.
     * @param clazz               the class.
     * @param method              the method.
     * @param codeAttribute       the code attribute.
     * @param newArrayOffset      the offset of the newarray instruction.
     * @param newArrayInstruction the newarray instruction.
     * @return whether there is a static array initialization.
     */
    public boolean matchesArrayInitialization(Clazz             clazz,
                                              Method            method,
                                              CodeAttribute     codeAttribute,
                                              int               newArrayOffset,
                                              SimpleInstruction newArrayInstruction)
    {
        array = null;

        TracedStack stackBefore =
            partialEvaluator.getStackBefore(newArrayOffset);

        IntegerValue integerValue =
            stackBefore.getTop(0).integerValue();

        if (!integerValue.isParticular())
        {
            return false;
        }


        int arrayLength = integerValue.value();

        int newArrayType     = newArrayInstruction.constant;
        int arrayStoreOpcode = arrayStoreOpcode(newArrayType);

        byte[] code = codeAttribute.code;

        int         offset      = newArrayOffset;
        Instruction instruction = newArrayInstruction;

        int skipOffset = skipPreStoreInstructions(clazz, method, codeAttribute, offset + instruction.length(offset));
        if (skipOffset > 0)
        {
            offset         = skipOffset;
            newArrayOffset = offset;
            instruction    = InstructionFactory.create(code, offset);
        }

        // Remember the potential initialization start.
        int tmpInitializationStart = offset + instruction.length(offset);

        // Check if all the elements in the array are initialized.
        for (int index = 0; index < arrayLength; index++)
        {
            // Check if the array reference is pushed.
            instruction = InstructionFactory.create(code, offset += instruction.length(offset));

            if (instruction.stackPushCount(clazz) < 1 ||
                !partialEvaluator.getStackAfter(offset).getTopActualProducerValue(0).instructionOffsetValue().contains(newArrayOffset))
            {
                return false;
            }

            // Check that the array index is pushed.
            instruction = InstructionFactory.create(code, offset += instruction.length(offset));
            if (instruction.stackPushCount(clazz) != 1)
            {
                return false;
            }

            Value indexValue = partialEvaluator.getStackAfter(offset).getTop(0);
            if (indexValue.computationalType() != Value.TYPE_INTEGER ||
                !indexValue.integerValue().isParticular()            ||
                indexValue.integerValue().value() != index)
            {
                return false;
            }

            // Check if a particular value is pushed.
            instruction = InstructionFactory.create(code, offset += instruction.length(offset));
            if (instruction.stackPushCount(clazz) < 1 ||
                !partialEvaluator.getStackAfter(offset).getTop(0).isParticular())
            {
                return false;
            }

            Value elementValue = partialEvaluator.getStackAfter(offset).getTop(0);

            // Check if the value is stored in the array.
            instruction = InstructionFactory.create(code, offset += instruction.length(offset));
            if (instruction.opcode != arrayStoreOpcode)
            {
                return false;
            }

            if (index == 0)
            {
                array = newArray(newArrayType, arrayLength);
            }

            arrayStore(newArrayType, array, index, elementValue);
        }

        arrayInitializationStart = tmpInitializationStart;
        arrayInitializationEnd   = offset;

        return offset > newArrayOffset;
    }


    /**
     * Returns the offset to skip to after a new-array instruction.
     * <p>
     * This is a work-around for code converted by dex2jar. In some
     * cases, after an array has been created, it is immediately
     * stored into a variable and loaded again:
     * <pre>
     *   newarray
     *   astore X
     *   aload  X
     *   initialization
     * </pre>
     *
     * @param clazz            the class.
     * @param method           the method.
     * @param codeAttribute    the code attribute.
     * @param startOffset      the start offset.
     * @return the offset to skip to
     */
    private int skipPreStoreInstructions(Clazz         clazz,
                                         Method        method,
                                         CodeAttribute codeAttribute,
                                         int           startOffset)
    {
        final int instructionCount = arrayPreStoreMatcher.instructionCount();

        arrayPreStoreMatcher.reset();
        for (int count = 0, offset = startOffset;
             count < instructionCount && offset < codeAttribute.u4codeLength;
             count++)
        {
            Instruction instruction =
                InstructionFactory.create(codeAttribute.code, offset);

            instruction.accept(clazz, method, codeAttribute, offset,
                               arrayPreStoreMatcher);

            offset += instruction.length(offset);
        }

        if (arrayPreStoreMatcher.isMatching())
        {
            return arrayPreStoreMatcher.matchedInstructionOffset(instructionCount - 1);
        }

        return -1;
    }

    /**
     * Returns the first offset of the recent static array initialization, i.e. the first
     * initialization instruction after the newarray.
     * @see #matchesArrayInitialization
     */
    public int arrayInitializationStart()
    {
        return arrayInitializationStart;
    }

    /**
     * Returns the last offset of the recent static array initialization.
     * @see #matchesArrayInitialization
     */
    public int arrayInitializationEnd()
    {
        return arrayInitializationEnd;
    }

    /**
     * Returns the recent static array initialization.
     * @see #matchesArrayInitialization
     */
    public Object array()
    {
        return array;
    }


    private byte internalType(int newArrayType)
    {
        switch (newArrayType)
        {
            case InstructionConstants.ARRAY_T_BOOLEAN:
            case InstructionConstants.ARRAY_T_BYTE:
            case InstructionConstants.ARRAY_T_CHAR:
            case InstructionConstants.ARRAY_T_SHORT:
            case InstructionConstants.ARRAY_T_INT:     return Value.TYPE_INTEGER;
            case InstructionConstants.ARRAY_T_LONG:    return Value.TYPE_LONG;
            case InstructionConstants.ARRAY_T_FLOAT:   return Value.TYPE_FLOAT;
            case InstructionConstants.ARRAY_T_DOUBLE:  return Value.TYPE_DOUBLE;
            default:
                throw new IllegalArgumentException("Unexpected new array type ["+newArrayType+"]");
        }
    }


    private byte arrayStoreOpcode(int newArrayType)
    {
        switch (newArrayType)
        {
            case InstructionConstants.ARRAY_T_BOOLEAN:
            case InstructionConstants.ARRAY_T_BYTE:    return InstructionConstants.OP_BASTORE;
            case InstructionConstants.ARRAY_T_CHAR:    return InstructionConstants.OP_CASTORE;
            case InstructionConstants.ARRAY_T_SHORT:   return InstructionConstants.OP_SASTORE;
            case InstructionConstants.ARRAY_T_INT:     return InstructionConstants.OP_IASTORE;
            case InstructionConstants.ARRAY_T_LONG:    return InstructionConstants.OP_LASTORE;
            case InstructionConstants.ARRAY_T_FLOAT:   return InstructionConstants.OP_FASTORE;
            case InstructionConstants.ARRAY_T_DOUBLE:  return InstructionConstants.OP_DASTORE;
            default:
                throw new IllegalArgumentException("Unexpected new array type ["+newArrayType+"]");
        }
    }


    private Object newArray(int newArrayType, int arrayLength)
    {
        switch (newArrayType)
        {
            case InstructionConstants.ARRAY_T_BOOLEAN: return new boolean[arrayLength];
            case InstructionConstants.ARRAY_T_BYTE:    return new byte[arrayLength];
            case InstructionConstants.ARRAY_T_CHAR:    return new char[arrayLength];
            case InstructionConstants.ARRAY_T_SHORT:   return new short[arrayLength];
            case InstructionConstants.ARRAY_T_INT:     return new int[arrayLength];
            case InstructionConstants.ARRAY_T_LONG:    return new long[arrayLength];
            case InstructionConstants.ARRAY_T_FLOAT:   return new float[arrayLength];
            case InstructionConstants.ARRAY_T_DOUBLE:  return new double[arrayLength];
            default:
                throw new IllegalArgumentException("Unexpected new array type ["+newArrayType+"]");
        }
    }


    private void arrayStore(int newArrayType, Object array, int index, Value value)
    {
        switch (newArrayType)
        {
            case InstructionConstants.ARRAY_T_BOOLEAN: ((boolean[])array)[index] = 0 !=   value.integerValue().value(); break;
            case InstructionConstants.ARRAY_T_BYTE:    ((byte   [])array)[index] = (byte) value.integerValue().value(); break;
            case InstructionConstants.ARRAY_T_CHAR:    ((char   [])array)[index] = (char) value.integerValue().value(); break;
            case InstructionConstants.ARRAY_T_SHORT:   ((short  [])array)[index] = (short)value.integerValue().value(); break;
            case InstructionConstants.ARRAY_T_INT:     ((int    [])array)[index] =        value.integerValue().value(); break;
            case InstructionConstants.ARRAY_T_LONG:    ((long   [])array)[index] =        value.longValue().value();    break;
            case InstructionConstants.ARRAY_T_FLOAT:   ((float  [])array)[index] =        value.floatValue().value();   break;
            case InstructionConstants.ARRAY_T_DOUBLE:  ((double [])array)[index] =        value.doubleValue().value();  break;
            default:
                throw new IllegalArgumentException("Unexpected new array type ["+newArrayType+"]");
        }
    }
}