/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.bcel.verifier.structurals;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import org.apache.bcel.generic.ASTORE;
import org.apache.bcel.generic.ATHROW;
import org.apache.bcel.generic.BranchInstruction;
import org.apache.bcel.generic.CodeExceptionGen;
import org.apache.bcel.generic.GotoInstruction;
import org.apache.bcel.generic.IndexedInstruction;
import org.apache.bcel.generic.Instruction;
import org.apache.bcel.generic.InstructionHandle;
import org.apache.bcel.generic.JsrInstruction;
import org.apache.bcel.generic.LocalVariableInstruction;
import org.apache.bcel.generic.MethodGen;
import org.apache.bcel.generic.RET;
import org.apache.bcel.generic.ReturnInstruction;
import org.apache.bcel.generic.Select;
import org.apache.bcel.verifier.exc.AssertionViolatedException;
import org.apache.bcel.verifier.exc.StructuralCodeConstraintException;

/**
 * Instances of this class contain information about the subroutines found in a code array of a method. This
 * implementation considers the top-level (the instructions reachable without a JSR or JSR_W starting off from the first
 * instruction in a code array of a method) being a special subroutine; see getTopLevel() for that. Please note that the
 * definition of subroutines in the Java Virtual Machine Specification, Second Edition is somewhat incomplete.
 * Therefore, JustIce uses an own, more rigid notion. Basically, a subroutine is a piece of code that starts at the
 * target of a JSR of JSR_W instruction and ends at a corresponding RET instruction. Note also that the control flow of
 * a subroutine may be complex and non-linear; and that subroutines may be nested. JustIce also mandates subroutines not
 * to be protected by exception handling code (for the sake of control flow predictability). To understand JustIce's
 * notion of subroutines, please read
 *
 * TODO: refer to the paper.
 *
 * @see #getTopLevel()
 */
public class Subroutines {
    // Node coloring constants
    private enum ColourConstants {
        WHITE, GRAY, BLACK
    }

    /**
     * This inner class implements the Subroutine interface.
     */
    private final class SubroutineImpl implements Subroutine {

        /**
         * UNSET, a symbol for an uninitialized localVariable field. This is used for the "top-level" Subroutine; for example no
         * subroutine.
         */
        private static final int UNSET = -1;

        private final SubroutineImpl[] EMPTY_ARRAY = {};

        /**
         * The Local Variable slot where the first instruction of this subroutine (an ASTORE) stores the JsrInstruction's
         * ReturnAddress in and the RET of this subroutine operates on.
         */
        private int localVariable = UNSET;

        /** The instructions that belong to this subroutine. */
        private final Set<InstructionHandle> instructions = new HashSet<>(); // Elements: InstructionHandle

        /**
         * The JSR or JSR_W instructions that define this subroutine by targeting it.
         */
        private final Set<InstructionHandle> theJSRs = new HashSet<>();

        /**
         * The RET instruction that leaves this subroutine.
         */
        private InstructionHandle theRET;

        /**
         * Constructs a new instance.
         */
        SubroutineImpl() {
            // empty
        }

        /**
         * Adds a new JSR or JSR_W that has this subroutine as its target.
         */
        public void addEnteringJsrInstruction(final InstructionHandle jsrInst) {
            if (jsrInst == null || !(jsrInst.getInstruction() instanceof JsrInstruction)) {
                throw new AssertionViolatedException("Expecting JsrInstruction InstructionHandle.");
            }
            if (localVariable == UNSET) {
                throw new AssertionViolatedException("Set the localVariable first!");
            }
            // Something is wrong when an ASTORE is targeted that does not operate on the same local variable than the rest of the
            // JsrInstruction-targets and the RET.
            // (We don't know out leader here so we cannot check if we're really targeted!)
            if (localVariable != ((ASTORE) ((JsrInstruction) jsrInst.getInstruction()).getTarget().getInstruction()).getIndex()) {
                throw new AssertionViolatedException("Setting a wrong JsrInstruction.");
            }
            theJSRs.add(jsrInst);
        }

        /*
         * Adds an instruction to this subroutine. All instructions must have been added before invoking setLeavingRET().
         *
         * @see #setLeavingRET
         */
        void addInstruction(final InstructionHandle ih) {
            if (theRET != null) {
                throw new AssertionViolatedException("All instructions must have been added before invoking setLeavingRET().");
            }
            instructions.add(ih);
        }

        /*
         * Refer to the Subroutine interface for documentation.
         */
        @Override
        public boolean contains(final InstructionHandle inst) {
            return instructions.contains(inst);
        }

        /*
         * Satisfies Subroutine.getAccessedLocalIndices().
         */
        @Override
        public int[] getAccessedLocalsIndices() {
            // TODO: Implement caching.
            final Set<Integer> acc = new HashSet<>();
            if (theRET == null && this != getTopLevel()) {
                throw new AssertionViolatedException("This subroutine object must be built up completely before calculating accessed locals.");
            }
            {
                for (final InstructionHandle ih : instructions) {
                    // RET is not a LocalVariableInstruction in the current version of BCEL.
                    if (ih.getInstruction() instanceof LocalVariableInstruction || ih.getInstruction() instanceof RET) {
                        final int idx = ((IndexedInstruction) ih.getInstruction()).getIndex();
                        acc.add(Integer.valueOf(idx));
                        // LONG? DOUBLE?.
                        try {
                            // LocalVariableInstruction instances are typed without the need to look into
                            // the constant pool.
                            if (ih.getInstruction() instanceof LocalVariableInstruction) {
                                final int s = ((LocalVariableInstruction) ih.getInstruction()).getType(null).getSize();
                                if (s == 2) {
                                    acc.add(Integer.valueOf(idx + 1));
                                }
                            }
                        } catch (final RuntimeException re) {
                            throw new AssertionViolatedException("BCEL did not like NULL as a ConstantPoolGen object.", re);
                        }
                    }
                }
            }

            {
                final int[] ret = new int[acc.size()];
                int j = -1;
                for (final Integer accessedLocal : acc) {
                    j++;
                    ret[j] = accessedLocal.intValue();
                }
                return ret;
            }
        }

        /*
         * Refer to the Subroutine interface for documentation.
         */
        @Override
        public InstructionHandle[] getEnteringJsrInstructions() {
            if (this == getTopLevel()) {
                throw new AssertionViolatedException("getLeavingRET() called on top level pseudo-subroutine.");
            }
            return theJSRs.toArray(InstructionHandle.EMPTY_ARRAY);
        }

        /*
         * Refer to the Subroutine interface for documentation.
         */
        @Override
        public InstructionHandle[] getInstructions() {
            return instructions.toArray(InstructionHandle.EMPTY_ARRAY);
        }

        /*
         * Refer to the Subroutine interface for documentation.
         */
        @Override
        public InstructionHandle getLeavingRET() {
            if (this == getTopLevel()) {
                throw new AssertionViolatedException("getLeavingRET() called on top level pseudo-subroutine.");
            }
            return theRET;
        }

        /* Satisfies Subroutine.getRecursivelyAccessedLocalsIndices(). */
        @Override
        public int[] getRecursivelyAccessedLocalsIndices() {
            final Set<Integer> s = new HashSet<>();
            final int[] lvs = getAccessedLocalsIndices();
            for (final int lv : lvs) {
                s.add(Integer.valueOf(lv));
            }
            getRecursivelyAccessedLocalsIndicesHelper(s, subSubs());
            final int[] ret = new int[s.size()];
            int j = -1;
            for (final Integer index : s) {
                j++;
                ret[j] = index.intValue();
            }
            return ret;
        }

        /**
         * A recursive helper method for getRecursivelyAccessedLocalsIndices().
         *
         * @see #getRecursivelyAccessedLocalsIndices()
         */
        private void getRecursivelyAccessedLocalsIndicesHelper(final Set<Integer> set, final Subroutine[] subs) {
            for (final Subroutine sub : subs) {
                final int[] lvs = sub.getAccessedLocalsIndices();
                for (final int lv : lvs) {
                    set.add(Integer.valueOf(lv));
                }
                if (sub.subSubs().length != 0) {
                    getRecursivelyAccessedLocalsIndicesHelper(set, sub.subSubs());
                }
            }
        }

        /**
         * Sets the leaving RET instruction. Must be invoked after all instructions are added. Must not be invoked for top-level
         * 'subroutine'.
         */
        void setLeavingRET() {
            if (localVariable == UNSET) {
                throw new AssertionViolatedException("setLeavingRET() called for top-level 'subroutine' or forgot to set local variable first.");
            }
            InstructionHandle ret = null;
            for (final InstructionHandle actual : instructions) {
                if (actual.getInstruction() instanceof RET) {
                    if (ret != null) {
                        throw new StructuralCodeConstraintException("Subroutine with more then one RET detected: '" + ret + "' and '" + actual + "'.");
                    }
                    ret = actual;
                }
            }
            if (ret == null) {
                throw new StructuralCodeConstraintException("Subroutine without a RET detected.");
            }
            if (((RET) ret.getInstruction()).getIndex() != localVariable) {
                throw new StructuralCodeConstraintException(
                    "Subroutine uses '" + ret + "' which does not match the correct local variable '" + localVariable + "'.");
            }
            theRET = ret;
        }

        /*
         * Sets the local variable slot the ASTORE that is targeted by the JsrInstructions of this subroutine operates on. This
         * subroutine's RET operates on that same local variable slot, of course.
         */
        void setLocalVariable(final int i) {
            if (localVariable != UNSET) {
                throw new AssertionViolatedException("localVariable set twice.");
            }
            localVariable = i;
        }

        /*
         * Satisfies Subroutine.subSubs().
         */
        @Override
        public Subroutine[] subSubs() {
            final Set<Subroutine> h = new HashSet<>();

            for (final InstructionHandle ih : instructions) {
                final Instruction inst = ih.getInstruction();
                if (inst instanceof JsrInstruction) {
                    final InstructionHandle targ = ((JsrInstruction) inst).getTarget();
                    h.add(getSubroutine(targ));
                }
            }
            return h.toArray(EMPTY_ARRAY);
        }

        /**
         * Returns a String representation of this object, merely for debugging purposes. (Internal) Warning: Verbosity on a
         * problematic subroutine may cause stack overflow errors due to recursive subSubs() calls. Don't use this, then.
         */
        @Override
        public String toString() {
            final StringBuilder ret = new StringBuilder();
            ret.append("Subroutine: Local variable is '").append(localVariable);
            ret.append("', JSRs are '").append(theJSRs);
            ret.append("', RET is '").append(theRET);
            ret.append("', Instructions: '").append(instructions).append("'.");

            ret.append(" Accessed local variable slots: '");
            int[] alv = getAccessedLocalsIndices();
            for (final int element : alv) {
                ret.append(element);
                ret.append(" ");
            }
            ret.append("'.");

            ret.append(" Recursively (via subsub...routines) accessed local variable slots: '");
            alv = getRecursivelyAccessedLocalsIndices();
            for (final int element : alv) {
                ret.append(element);
                ret.append(" ");
            }
            ret.append("'.");

            return ret.toString();
        }

    } // end Inner Class SubrouteImpl

    /**
     * A utility method that calculates the successors of a given InstructionHandle <strong>in the same subroutine</strong>. That
     * means, a RET does not have any successors as defined here. A JsrInstruction has its physical successor as its
     * successor (opposed to its target) as defined here.
     */
    private static InstructionHandle[] getSuccessors(final InstructionHandle instruction) {
        final InstructionHandle[] single = new InstructionHandle[1];

        final Instruction inst = instruction.getInstruction();

        // Terminates method normally.
        // Terminates method abnormally, because JustIce mandates
        // subroutines not to be protected by exception handlers.
        if (inst instanceof RET || inst instanceof ReturnInstruction || inst instanceof ATHROW) {
            return InstructionHandle.EMPTY_ARRAY;
        }

        // See method comment.
        if (inst instanceof JsrInstruction) {
            single[0] = instruction.getNext();
            return single;
        }

        if (inst instanceof GotoInstruction) {
            single[0] = ((GotoInstruction) inst).getTarget();
            return single;
        }

        if (inst instanceof BranchInstruction) {
            if (inst instanceof Select) {
                // BCEL's getTargets() returns only the non-default targets,
                // thanks to Eli Tilevich for reporting.
                final InstructionHandle[] matchTargets = ((Select) inst).getTargets();
                final InstructionHandle[] ret = new InstructionHandle[matchTargets.length + 1];
                ret[0] = ((Select) inst).getTarget();
                System.arraycopy(matchTargets, 0, ret, 1, matchTargets.length);
                return ret;
            }
            final InstructionHandle[] pair = new InstructionHandle[2];
            pair[0] = instruction.getNext();
            pair[1] = ((BranchInstruction) inst).getTarget();
            return pair;
        }

        // default case: Fall through.
        single[0] = instruction.getNext();
        return single;
    }

    /**
     * The map containing the subroutines found. Key: InstructionHandle of the leader of the subroutine. Elements:
     * SubroutineImpl objects.
     */
    private final Map<InstructionHandle, Subroutine> subroutines = new HashMap<>();

    /**
     * This is referring to a special subroutine, namely the top level. This is not really a subroutine but we use it to
     * distinguish between top level instructions and unreachable instructions.
     */
    // CHECKSTYLE:OFF
    public final Subroutine TOPLEVEL; // TODO can this be made private?
    // CHECKSTYLE:ON

    /**
     * Constructs a new instance.
     *
     * @param mg A MethodGen object representing method to create the Subroutine objects of. Assumes that JustIce strict
     *        checks are needed.
     */
    public Subroutines(final MethodGen mg) {
        this(mg, true);
    }

    /**
     * Constructs a new instance.
     *
     * @param mg A MethodGen object representing method to create the Subroutine objects of.
     * @param enableJustIceCheck whether to enable additional JustIce checks.
     * @since 6.0
     */
    public Subroutines(final MethodGen mg, final boolean enableJustIceCheck) {
        final InstructionHandle[] all = mg.getInstructionList().getInstructionHandles();
        final CodeExceptionGen[] handlers = mg.getExceptionHandlers();

        // Define our "Toplevel" fake subroutine.
        TOPLEVEL = new SubroutineImpl();

        // Calculate "real" subroutines.
        final Set<InstructionHandle> subLeaders = new HashSet<>(); // Elements: InstructionHandle
        for (final InstructionHandle element : all) {
            final Instruction inst = element.getInstruction();
            if (inst instanceof JsrInstruction) {
                subLeaders.add(((JsrInstruction) inst).getTarget());
            }
        }

        // Build up the database.
        for (final InstructionHandle astore : subLeaders) {
            final SubroutineImpl sr = new SubroutineImpl();
            sr.setLocalVariable(((ASTORE) astore.getInstruction()).getIndex());
            subroutines.put(astore, sr);
        }

        // Fake it a bit. We want a virtual "TopLevel" subroutine.
        subroutines.put(all[0], TOPLEVEL);
        subLeaders.add(all[0]);

        // Tell the subroutines about their JsrInstructions.
        // Note that there cannot be a JSR targeting the top-level
        // since "Jsr 0" is disallowed in Pass 3a.
        // Instructions shared by a subroutine and the toplevel are
        // disallowed and checked below, after the BFS.
        for (final InstructionHandle element : all) {
            final Instruction inst = element.getInstruction();
            if (inst instanceof JsrInstruction) {
                final InstructionHandle leader = ((JsrInstruction) inst).getTarget();
                ((SubroutineImpl) getSubroutine(leader)).addEnteringJsrInstruction(element);
            }
        }

        // Now do a BFS from every subroutine leader to find all the
        // instructions that belong to a subroutine.
        // we don't want to assign an instruction to two or more Subroutine objects.
        final Set<InstructionHandle> instructionsAssigned = new HashSet<>();

        // Graph coloring. Key: InstructionHandle, Value: ColourConstants enum.
        final Map<InstructionHandle, ColourConstants> colors = new HashMap<>();

        final List<InstructionHandle> qList = new ArrayList<>();
        for (final InstructionHandle actual : subLeaders) {
            // Do some BFS with "actual" as the root of the graph.
            // Init colors
            for (final InstructionHandle element : all) {
                colors.put(element, ColourConstants.WHITE);
            }
            colors.put(actual, ColourConstants.GRAY);
            // Init Queue

            qList.clear();
            qList.add(actual); // add(Obj) adds to the end, remove(0) removes from the start.

            /*
             * BFS ALGORITHM MODIFICATION: Start out with multiple "root" nodes, as exception handlers are starting points of
             * top-level code, too. [why top-level? TODO: Refer to the special JustIce notion of subroutines.]
             */
            if (actual == all[0]) {
                for (final CodeExceptionGen handler : handlers) {
                    colors.put(handler.getHandlerPC(), ColourConstants.GRAY);
                    qList.add(handler.getHandlerPC());
                }
            }
            /* CONTINUE NORMAL BFS ALGORITHM */

            // Loop until Queue is empty
            while (!qList.isEmpty()) {
                final InstructionHandle u = qList.remove(0);
                final InstructionHandle[] successors = getSuccessors(u);
                for (final InstructionHandle successor : successors) {
                    if (colors.get(successor) == ColourConstants.WHITE) {
                        colors.put(successor, ColourConstants.GRAY);
                        qList.add(successor);
                    }
                }
                colors.put(u, ColourConstants.BLACK);
            }
            // BFS ended above.
            for (final InstructionHandle element : all) {
                if (colors.get(element) == ColourConstants.BLACK) {
                    ((SubroutineImpl) (actual == all[0] ? getTopLevel() : getSubroutine(actual))).addInstruction(element);
                    if (instructionsAssigned.contains(element)) {
                        throw new StructuralCodeConstraintException(
                            "Instruction '" + element + "' is part of more than one subroutine (or of the top level and a subroutine).");
                    }
                    instructionsAssigned.add(element);
                }
            }
            if (actual != all[0]) { // If we don't deal with the top-level 'subroutine'
                ((SubroutineImpl) getSubroutine(actual)).setLeavingRET();
            }
        }

        if (enableJustIceCheck) {
            // Now make sure no instruction of a Subroutine is protected by exception handling code
            // as is mandated by JustIces notion of subroutines.
            for (final CodeExceptionGen handler : handlers) {
                InstructionHandle protectedIh = handler.getStartPC();
                while (protectedIh != handler.getEndPC().getNext()) {
                    // Note the inclusive/inclusive notation of "generic API" exception handlers!
                    for (final Subroutine sub : subroutines.values()) {
                        if (sub != subroutines.get(all[0]) && sub.contains(protectedIh)) {
                            throw new StructuralCodeConstraintException("Subroutine instruction '" + protectedIh + "' is protected by an exception handler, '"
                                + handler + "'. This is forbidden by the JustIce verifier due to its clear definition of subroutines.");
                        }
                    }
                    protectedIh = protectedIh.getNext();
                }
            }
        }

        // Now make sure no subroutine is calling a subroutine
        // that uses the same local variable for the RET as themselves
        // (recursively).
        // This includes that subroutines may not call themselves
        // recursively, even not through intermediate calls to other
        // subroutines.
        noRecursiveCalls(getTopLevel(), new HashSet<>());

    }

    /**
     * Returns the Subroutine object associated with the given leader (that is, the first instruction of the subroutine).
     * You must not use this to get the top-level instructions modeled as a Subroutine object.
     *
     * @see #getTopLevel()
     */
    public Subroutine getSubroutine(final InstructionHandle leader) {
        final Subroutine ret = subroutines.get(leader);

        if (ret == null) {
            throw new AssertionViolatedException("Subroutine requested for an InstructionHandle that is not a leader of a subroutine.");
        }

        if (ret == TOPLEVEL) {
            throw new AssertionViolatedException("TOPLEVEL special subroutine requested; use getTopLevel().");
        }

        return ret;
    }

    /**
     * For easy handling, the piece of code that is <strong>not</strong> a subroutine, the top-level, is also modeled as a Subroutine
     * object. It is a special Subroutine object where <B>you must not invoke getEnteringJsrInstructions() or
     * getLeavingRET()</B>.
     *
     * @see Subroutine#getEnteringJsrInstructions()
     * @see Subroutine#getLeavingRET()
     */
    public Subroutine getTopLevel() {
        return TOPLEVEL;
    }

    /**
     * This (recursive) utility method makes sure that no subroutine is calling a subroutine that uses the same local
     * variable for the RET as themselves (recursively). This includes that subroutines may not call themselves recursively,
     * even not through intermediate calls to other subroutines.
     *
     * @throws StructuralCodeConstraintException if the above constraint is not satisfied.
     */
    private void noRecursiveCalls(final Subroutine sub, final Set<Integer> set) {
        final Subroutine[] subs = sub.subSubs();

        for (final Subroutine sub2 : subs) {
            final int index = ((RET) sub2.getLeavingRET().getInstruction()).getIndex();

            if (!set.add(Integer.valueOf(index))) {
                // Don't use toString() here because of possibly infinite recursive subSubs() calls then.
                final SubroutineImpl si = (SubroutineImpl) sub2;
                throw new StructuralCodeConstraintException("Subroutine with local variable '" + si.localVariable + "', JSRs '" + si.theJSRs + "', RET '"
                    + si.theRET + "' is called by a subroutine which uses the same local variable index as itself; maybe even a recursive call?"
                    + " JustIce's clean definition of a subroutine forbids both.");
            }

            noRecursiveCalls(sub2, set);

            set.remove(Integer.valueOf(index));
        }
    }

    /**
     * Returns the subroutine object associated with the given instruction. This is a costly operation, you should consider
     * using getSubroutine(InstructionHandle). Returns 'null' if the given InstructionHandle lies in so-called 'dead code',
     * for example code that can never be executed.
     *
     * @see #getSubroutine(InstructionHandle)
     * @see #getTopLevel()
     */
    public Subroutine subroutineOf(final InstructionHandle any) {
        for (final Subroutine s : subroutines.values()) {
            if (s.contains(any)) {
                return s;
            }
        }
        System.err.println("DEBUG: Please verify '" + any.toString(true) + "' lies in dead code.");
        return null;
        // throw new AssertionViolatedException("No subroutine for InstructionHandle found (DEAD CODE?).");
    }

    /**
     * Returns a String representation of this object; merely for debugging puposes.
     */
    @Override
    public String toString() {
        return "---\n" + subroutines + "\n---\n";
    }
}