/*
 * Copyright (c) 2011, 2018, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code 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
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */


package org.graalvm.compiler.phases.graph;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.Iterator;
import java.util.List;

import jdk.internal.vm.compiler.collections.EconomicMap;
import jdk.internal.vm.compiler.collections.Equivalence;
import jdk.internal.vm.compiler.collections.MapCursor;
import org.graalvm.compiler.graph.Node;
import org.graalvm.compiler.nodes.AbstractBeginNode;
import org.graalvm.compiler.nodes.AbstractEndNode;
import org.graalvm.compiler.nodes.AbstractMergeNode;
import org.graalvm.compiler.nodes.EndNode;
import org.graalvm.compiler.nodes.FixedNode;
import org.graalvm.compiler.nodes.FixedWithNextNode;
import org.graalvm.compiler.nodes.LoopBeginNode;
import org.graalvm.compiler.nodes.LoopEndNode;
import org.graalvm.compiler.nodes.LoopExitNode;

public final class ReentrantNodeIterator {

    public static class LoopInfo<StateT> {

        public final EconomicMap<LoopEndNode, StateT> endStates;
        public final EconomicMap<LoopExitNode, StateT> exitStates;

        public LoopInfo(int endCount, int exitCount) {
            endStates = EconomicMap.create(Equivalence.IDENTITY, endCount);
            exitStates = EconomicMap.create(Equivalence.IDENTITY, exitCount);
        }
    }

    public abstract static class NodeIteratorClosure<StateT> {

        protected abstract StateT processNode(FixedNode node, StateT currentState);

        protected abstract StateT merge(AbstractMergeNode merge, List<StateT> states);

        protected abstract StateT afterSplit(AbstractBeginNode node, StateT oldState);

        protected abstract EconomicMap<LoopExitNode, StateT> processLoop(LoopBeginNode loop, StateT initialState);

        
Determine whether iteration should continue in the current state.
Params:
currentState – /**
         * Determine whether iteration should continue in the current state.
         *
         * @param currentState
         */
        protected boolean continueIteration(StateT currentState) {
            return true;
        }
    }

    private ReentrantNodeIterator() {
        // no instances allowed
    }

    public static <StateT> LoopInfo<StateT> processLoop(NodeIteratorClosure<StateT> closure, LoopBeginNode loop, StateT initialState) {
        EconomicMap<FixedNode, StateT> blockEndStates = apply(closure, loop, initialState, loop);

        LoopInfo<StateT> info = new LoopInfo<>(loop.loopEnds().count(), loop.loopExits().count());
        for (LoopEndNode end : loop.loopEnds()) {
            if (blockEndStates.containsKey(end)) {
                info.endStates.put(end, blockEndStates.get(end));
            }
        }
        for (LoopExitNode exit : loop.loopExits()) {
            if (blockEndStates.containsKey(exit)) {
                info.exitStates.put(exit, blockEndStates.get(exit));
            }
        }
        return info;
    }

    public static <StateT> void apply(NodeIteratorClosure<StateT> closure, FixedNode start, StateT initialState) {
        apply(closure, start, initialState, null);
    }

    private static <StateT> EconomicMap<FixedNode, StateT> apply(NodeIteratorClosure<StateT> closure, FixedNode start, StateT initialState, LoopBeginNode boundary) {
        assert start != null;
        Deque<AbstractBeginNode> nodeQueue = new ArrayDeque<>();
        EconomicMap<FixedNode, StateT> blockEndStates = EconomicMap.create(Equivalence.IDENTITY);

        StateT state = initialState;
        FixedNode current = start;
        do {
            while (current instanceof FixedWithNextNode) {
                if (boundary != null && current instanceof LoopExitNode && ((LoopExitNode) current).loopBegin() == boundary) {
                    blockEndStates.put(current, state);
                    current = null;
                } else {
                    FixedNode next = ((FixedWithNextNode) current).next();
                    state = closure.processNode(current, state);
                    current = closure.continueIteration(state) ? next : null;
                }
            }

            if (current != null) {
                state = closure.processNode(current, state);

                if (closure.continueIteration(state)) {
                    Iterator<Node> successors = current.successors().iterator();
                    if (!successors.hasNext()) {
                        if (current instanceof LoopEndNode) {
                            blockEndStates.put(current, state);
                        } else if (current instanceof EndNode) {
                            // add the end node and see if the merge is ready for processing
                            AbstractMergeNode merge = ((EndNode) current).merge();
                            if (merge instanceof LoopBeginNode) {
                                EconomicMap<LoopExitNode, StateT> loopExitState = closure.processLoop((LoopBeginNode) merge, state);
                                MapCursor<LoopExitNode, StateT> entry = loopExitState.getEntries();
                                while (entry.advance()) {
                                    blockEndStates.put(entry.getKey(), entry.getValue());
                                    nodeQueue.add(entry.getKey());
                                }
                            } else {
                                boolean endsVisited = true;
                                for (AbstractEndNode forwardEnd : merge.forwardEnds()) {
                                    if (forwardEnd != current && !blockEndStates.containsKey(forwardEnd)) {
                                        endsVisited = false;
                                        break;
                                    }
                                }
                                if (endsVisited) {
                                    ArrayList<StateT> states = new ArrayList<>(merge.forwardEndCount());
                                    for (int i = 0; i < merge.forwardEndCount(); i++) {
                                        AbstractEndNode forwardEnd = merge.forwardEndAt(i);
                                        assert forwardEnd == current || blockEndStates.containsKey(forwardEnd);
                                        StateT other = forwardEnd == current ? state : blockEndStates.removeKey(forwardEnd);
                                        states.add(other);
                                    }
                                    state = closure.merge(merge, states);
                                    current = closure.continueIteration(state) ? merge : null;
                                    continue;
                                } else {
                                    assert !blockEndStates.containsKey(current);
                                    blockEndStates.put(current, state);
                                }
                            }
                        }
                    } else {
                        FixedNode firstSuccessor = (FixedNode) successors.next();
                        if (!successors.hasNext()) {
                            current = firstSuccessor;
                            continue;
                        } else {
                            do {
                                AbstractBeginNode successor = (AbstractBeginNode) successors.next();
                                StateT successorState = closure.afterSplit(successor, state);
                                if (closure.continueIteration(successorState)) {
                                    blockEndStates.put(successor, successorState);
                                    nodeQueue.add(successor);
                                }
                            } while (successors.hasNext());

                            state = closure.afterSplit((AbstractBeginNode) firstSuccessor, state);
                            current = closure.continueIteration(state) ? firstSuccessor : null;
                            continue;
                        }
                    }
                }
            }

            // get next queued block
            if (nodeQueue.isEmpty()) {
                return blockEndStates;
            } else {
                current = nodeQueue.removeFirst();
                assert blockEndStates.containsKey(current);
                state = blockEndStates.removeKey(current);
                assert !(current instanceof AbstractMergeNode) && current instanceof AbstractBeginNode;
            }
        } while (true);
    }
}