/*

 * 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.virtual.phases.ea;

import java.util.ArrayList;

import java.util.BitSet;

import java.util.Iterator;

import java.util.List;

import java.util.function.IntUnaryOperator;

import jdk.internal.vm.compiler.collections.EconomicMap;

import jdk.internal.vm.compiler.collections.EconomicSet;

import jdk.internal.vm.compiler.collections.Equivalence;

import org.graalvm.compiler.core.common.GraalOptions;

import org.graalvm.compiler.core.common.cfg.Loop;

import org.graalvm.compiler.core.common.spi.ConstantFieldProvider;

import org.graalvm.compiler.core.common.type.Stamp;

import org.graalvm.compiler.core.common.type.StampFactory;

import org.graalvm.compiler.debug.CounterKey;

import org.graalvm.compiler.debug.DebugContext;

import org.graalvm.compiler.graph.Node;

import org.graalvm.compiler.graph.NodeBitMap;

import org.graalvm.compiler.graph.Position;

import org.graalvm.compiler.graph.spi.Canonicalizable;

import org.graalvm.compiler.nodes.AbstractEndNode;

import org.graalvm.compiler.nodes.CallTargetNode;

import org.graalvm.compiler.nodes.ConstantNode;

import org.graalvm.compiler.nodes.ControlSinkNode;

import org.graalvm.compiler.nodes.FixedNode;

import org.graalvm.compiler.nodes.FixedWithNextNode;

import org.graalvm.compiler.nodes.FrameState;

import org.graalvm.compiler.nodes.Invoke;

import org.graalvm.compiler.nodes.LoopBeginNode;

import org.graalvm.compiler.nodes.LoopExitNode;

import org.graalvm.compiler.nodes.NodeView;

import org.graalvm.compiler.nodes.PhiNode;

import org.graalvm.compiler.nodes.ProxyNode;

import org.graalvm.compiler.nodes.StructuredGraph;

import org.graalvm.compiler.nodes.StructuredGraph.ScheduleResult;

import org.graalvm.compiler.nodes.UnwindNode;

import org.graalvm.compiler.nodes.ValueNode;

import org.graalvm.compiler.nodes.ValuePhiNode;

import org.graalvm.compiler.nodes.ValueProxyNode;

import org.graalvm.compiler.nodes.VirtualState;

import org.graalvm.compiler.nodes.VirtualState.NodeClosure;

import org.graalvm.compiler.nodes.cfg.Block;

import org.graalvm.compiler.nodes.spi.LoweringProvider;

import org.graalvm.compiler.nodes.spi.NodeWithState;

import org.graalvm.compiler.nodes.spi.Virtualizable;

import org.graalvm.compiler.nodes.spi.VirtualizableAllocation;

import org.graalvm.compiler.nodes.spi.VirtualizerTool;

import org.graalvm.compiler.nodes.virtual.AllocatedObjectNode;

import org.graalvm.compiler.nodes.virtual.VirtualObjectNode;

import org.graalvm.compiler.virtual.nodes.VirtualObjectState;

import jdk.vm.ci.meta.ConstantReflectionProvider;

import jdk.vm.ci.meta.JavaConstant;

import jdk.vm.ci.meta.JavaKind;

import jdk.vm.ci.meta.MetaAccessProvider;

public abstract class PartialEscapeClosure<BlockT extends PartialEscapeBlockState<BlockT>> extends EffectsClosure<BlockT> {

    public static final CounterKey COUNTER_MATERIALIZATIONS = DebugContext.counter("Materializations");

    public static final CounterKey COUNTER_MATERIALIZATIONS_PHI = DebugContext.counter("MaterializationsPhi");

    public static final CounterKey COUNTER_MATERIALIZATIONS_MERGE = DebugContext.counter("MaterializationsMerge");

    public static final CounterKey COUNTER_MATERIALIZATIONS_UNHANDLED = DebugContext.counter("MaterializationsUnhandled");

    public static final CounterKey COUNTER_MATERIALIZATIONS_LOOP_REITERATION = DebugContext.counter("MaterializationsLoopReiteration");

    public static final CounterKey COUNTER_MATERIALIZATIONS_LOOP_END = DebugContext.counter("MaterializationsLoopEnd");

    public static final CounterKey COUNTER_ALLOCATION_REMOVED = DebugContext.counter("AllocationsRemoved");

    public static final CounterKey COUNTER_MEMORYCHECKPOINT = DebugContext.counter("MemoryCheckpoint");

    /**

     * Nodes with inputs that were modified during analysis are marked in this bitset - this way

     * nodes that are not influenced at all by analysis can be rejected quickly.

     */

    private final NodeBitMap hasVirtualInputs;

    /**

     * This is handed out to implementers of {@link Virtualizable}.

     */

    protected final VirtualizerToolImpl tool;

    /**

     * The indexes into this array correspond to {@link VirtualObjectNode#getObjectId()}.

     */

    public final ArrayList<VirtualObjectNode> virtualObjects = new ArrayList<>();

    @Override

    public boolean needsApplyEffects() {

        if (hasChanged()) {

            return true;

        }

        /*

         * If there is a mismatch between the number of materializations and the number of

         * virtualizations, we need to apply effects, even if there were no other significant

         * changes to the graph. This applies to each block, since moving from one block to the

         * other can also be important (if the probabilities of the block differ).

         */

        for (Block block : cfg.getBlocks()) {

            GraphEffectList effects = blockEffects.get(block);

            if (effects != null) {

                if (effects.getVirtualizationDelta() != 0) {

                    return true;

                }

            }

        }

        return false;

    }

    private final class CollectVirtualObjectsClosure extends NodeClosure<ValueNode> {

        private final EconomicSet<VirtualObjectNode> virtual;

        private final GraphEffectList effects;

        private final BlockT state;

        private CollectVirtualObjectsClosure(EconomicSet<VirtualObjectNode> virtual, GraphEffectList effects, BlockT state) {

            this.virtual = virtual;

            this.effects = effects;

            this.state = state;

        }

        @Override

        public void apply(Node usage, ValueNode value) {

            if (value instanceof VirtualObjectNode) {

                VirtualObjectNode object = (VirtualObjectNode) value;

                if (object.getObjectId() != -1 && state.getObjectStateOptional(object) != null) {

                    virtual.add(object);

                }

            } else {

                ValueNode alias = getAlias(value);

                if (alias instanceof VirtualObjectNode) {

                    VirtualObjectNode object = (VirtualObjectNode) alias;

                    virtual.add(object);

                    effects.replaceFirstInput(usage, value, object);

                }

            }

        }

    }

    /**

     * Final subclass of PartialEscapeClosure, for performance and to make everything behave nicely

     * with generics.

     */

    public static final class Final extends PartialEscapeClosure<PartialEscapeBlockState.Final> {

        public Final(ScheduleResult schedule, MetaAccessProvider metaAccess, ConstantReflectionProvider constantReflection, ConstantFieldProvider constantFieldProvider,

                        LoweringProvider loweringProvider) {

            super(schedule, metaAccess, constantReflection, constantFieldProvider, loweringProvider);

        }

        @Override

        protected PartialEscapeBlockState.Final getInitialState() {

            return new PartialEscapeBlockState.Final(tool.getOptions(), tool.getDebug());

        }

        @Override

        protected PartialEscapeBlockState.Final cloneState(PartialEscapeBlockState.Final oldState) {

            return new PartialEscapeBlockState.Final(oldState);

        }

    }

    public PartialEscapeClosure(ScheduleResult schedule, MetaAccessProvider metaAccess, ConstantReflectionProvider constantReflection, ConstantFieldProvider constantFieldProvider) {

        this(schedule, metaAccess, constantReflection, constantFieldProvider, null);

    }

    public PartialEscapeClosure(ScheduleResult schedule, MetaAccessProvider metaAccess, ConstantReflectionProvider constantReflection, ConstantFieldProvider constantFieldProvider,

                    LoweringProvider loweringProvider) {

        super(schedule, schedule.getCFG());

        StructuredGraph graph = schedule.getCFG().graph;

        this.hasVirtualInputs = graph.createNodeBitMap();

        this.tool = new VirtualizerToolImpl(metaAccess, constantReflection, constantFieldProvider, this, graph.getAssumptions(), graph.getOptions(), debug, loweringProvider);

    }

    /**

     * @return true if the node was deleted, false otherwise

     */

    @Override

    protected boolean processNode(Node node, BlockT state, GraphEffectList effects, FixedWithNextNode lastFixedNode) {

        /*

         * These checks make up for the fact that an earliest schedule moves CallTargetNodes upwards

         * and thus materializes virtual objects needlessly. Also, FrameStates and ConstantNodes are

         * scheduled, but can safely be ignored.

         */

        if (node instanceof CallTargetNode || node instanceof FrameState || node instanceof ConstantNode) {

            return false;

        } else if (node instanceof Invoke) {

            processNodeInternal(((Invoke) node).callTarget(), state, effects, lastFixedNode);

        }

        return processNodeInternal(node, state, effects, lastFixedNode);

    }

    private boolean processNodeInternal(Node node, BlockT state, GraphEffectList effects, FixedWithNextNode lastFixedNode) {

        FixedNode nextFixedNode = lastFixedNode == null ? null : lastFixedNode.next();

        VirtualUtil.trace(node.getOptions(), debug, "%s", node);

        if (requiresProcessing(node)) {

            if (processVirtualizable((ValueNode) node, nextFixedNode, state, effects) == false) {

                return false;

            }

            if (tool.isDeleted()) {

                VirtualUtil.trace(node.getOptions(), debug, "deleted virtualizable allocation %s", node);

                return true;

            }

        }

        if (hasVirtualInputs.isMarked(node) && node instanceof ValueNode) {

            if (node instanceof Virtualizable) {

                if (processVirtualizable((ValueNode) node, nextFixedNode, state, effects) == false) {

                    return false;

                }

                if (tool.isDeleted()) {

                    VirtualUtil.trace(node.getOptions(), debug, "deleted virtualizable node %s", node);

                    return true;

                }

            }

            processNodeInputs((ValueNode) node, nextFixedNode, state, effects);

        }

        if (hasScalarReplacedInputs(node) && node instanceof ValueNode) {

            if (processNodeWithScalarReplacedInputs((ValueNode) node, nextFixedNode, state, effects)) {

                return true;

            }

        }

        return false;

    }

    protected boolean requiresProcessing(Node node) {

        return node instanceof VirtualizableAllocation;

    }

    private boolean processVirtualizable(ValueNode node, FixedNode insertBefore, BlockT state, GraphEffectList effects) {

        tool.reset(state, node, insertBefore, effects);

        return virtualize(node, tool);

    }

    protected boolean virtualize(ValueNode node, VirtualizerTool vt) {

        ((Virtualizable) node).virtualize(vt);

        return true; // request further processing

    }

    /**

     * This tries to canonicalize the node based on improved (replaced) inputs.

     */

    @SuppressWarnings("unchecked")

    private boolean processNodeWithScalarReplacedInputs(ValueNode node, FixedNode insertBefore, BlockT state, GraphEffectList effects) {

        ValueNode canonicalizedValue = node;

        if (node instanceof Canonicalizable.Unary<?>) {

            Canonicalizable.Unary<ValueNode> canonicalizable = (Canonicalizable.Unary<ValueNode>) node;

            ObjectState valueObj = getObjectState(state, canonicalizable.getValue());

            ValueNode valueAlias = valueObj != null ? valueObj.getMaterializedValue() : getScalarAlias(canonicalizable.getValue());

            if (valueAlias != canonicalizable.getValue()) {

                canonicalizedValue = (ValueNode) canonicalizable.canonical(tool, valueAlias);

            }

        } else if (node instanceof Canonicalizable.Binary<?>) {

            Canonicalizable.Binary<ValueNode> canonicalizable = (Canonicalizable.Binary<ValueNode>) node;

            ObjectState xObj = getObjectState(state, canonicalizable.getX());

            ValueNode xAlias = xObj != null ? xObj.getMaterializedValue() : getScalarAlias(canonicalizable.getX());

            ObjectState yObj = getObjectState(state, canonicalizable.getY());

            ValueNode yAlias = yObj != null ? yObj.getMaterializedValue() : getScalarAlias(canonicalizable.getY());

            if (xAlias != canonicalizable.getX() || yAlias != canonicalizable.getY()) {

                canonicalizedValue = (ValueNode) canonicalizable.canonical(tool, xAlias, yAlias);

            }

        } else {

            return false;

        }

        if (canonicalizedValue != node && canonicalizedValue != null) {

            if (canonicalizedValue.isAlive()) {

                ValueNode alias = getAliasAndResolve(state, canonicalizedValue);

                if (alias instanceof VirtualObjectNode) {

                    addVirtualAlias((VirtualObjectNode) alias, node);

                    effects.deleteNode(node);

                } else {

                    effects.replaceAtUsages(node, alias, insertBefore);

                    addScalarAlias(node, alias);

                }

            } else {

                if (!prepareCanonicalNode(canonicalizedValue, state, effects)) {

                    VirtualUtil.trace(node.getOptions(), debug, "replacement via canonicalization too complex: %s -> %s", node, canonicalizedValue);

                    return false;

                }

                if (canonicalizedValue instanceof ControlSinkNode) {

                    effects.replaceWithSink((FixedWithNextNode) node, (ControlSinkNode) canonicalizedValue);

                    state.markAsDead();

                } else {

                    effects.replaceAtUsages(node, canonicalizedValue, insertBefore);

                    addScalarAlias(node, canonicalizedValue);

                }

            }

            VirtualUtil.trace(node.getOptions(), debug, "replaced via canonicalization: %s -> %s", node, canonicalizedValue);

            return true;

        }

        return false;

    }

    /**

     * Nodes created during canonicalizations need to be scanned for values that were replaced.

     */

    private boolean prepareCanonicalNode(ValueNode node, BlockT state, GraphEffectList effects) {

        assert !node.isAlive();

        for (Position pos : node.inputPositions()) {

            Node input = pos.get(node);

            if (input instanceof ValueNode) {

                if (input.isAlive()) {

                    if (!(input instanceof VirtualObjectNode)) {

                        ObjectState obj = getObjectState(state, (ValueNode) input);

                        if (obj != null) {

                            if (obj.isVirtual()) {

                                return false;

                            } else {

                                pos.initialize(node, obj.getMaterializedValue());

                            }

                        } else {

                            pos.initialize(node, getScalarAlias((ValueNode) input));

                        }

                    }

                } else {

                    if (!prepareCanonicalNode((ValueNode) input, state, effects)) {

                        return false;

                    }

                }

            }

        }

        return true;

    }

    /**

     * This replaces all inputs that point to virtual or materialized values with the actual value,

     * materializing if necessary. Also takes care of frame states, adding the necessary

     * {@link VirtualObjectState}.

     */

    private void processNodeInputs(ValueNode node, FixedNode insertBefore, BlockT state, GraphEffectList effects) {

        VirtualUtil.trace(node.getOptions(), debug, "processing nodewithstate: %s", node);

        for (Node input : node.inputs()) {

            if (input instanceof ValueNode) {

                ValueNode alias = getAlias((ValueNode) input);

                if (alias instanceof VirtualObjectNode) {

                    int id = ((VirtualObjectNode) alias).getObjectId();

                    ensureMaterialized(state, id, insertBefore, effects, COUNTER_MATERIALIZATIONS_UNHANDLED);

                    effects.replaceFirstInput(node, input, state.getObjectState(id).getMaterializedValue());

                    VirtualUtil.trace(node.getOptions(), debug, "replacing input %s at %s", input, node);

                }

            }

        }

        if (node instanceof NodeWithState) {

            processNodeWithState((NodeWithState) node, state, effects);

        }

    }

    private void processNodeWithState(NodeWithState nodeWithState, BlockT state, GraphEffectList effects) {

        for (FrameState fs : nodeWithState.states()) {

            FrameState frameState = getUniqueFramestate(nodeWithState, fs);

            EconomicSet<VirtualObjectNode> virtual = EconomicSet.create(Equivalence.IDENTITY_WITH_SYSTEM_HASHCODE);

            frameState.applyToNonVirtual(new CollectVirtualObjectsClosure(virtual, effects, state));

            collectLockedVirtualObjects(state, virtual);

            collectReferencedVirtualObjects(state, virtual);

            addVirtualMappings(frameState, virtual, state, effects);

        }

    }

    private static FrameState getUniqueFramestate(NodeWithState nodeWithState, FrameState frameState) {

        if (frameState.hasMoreThanOneUsage()) {

            // Can happen for example from inlined snippets with multiple state split nodes.

            FrameState copy = (FrameState) frameState.copyWithInputs();

            nodeWithState.asNode().replaceFirstInput(frameState, copy);

            return copy;

        }

        return frameState;

    }

    private void addVirtualMappings(FrameState frameState, EconomicSet<VirtualObjectNode> virtual, BlockT state, GraphEffectList effects) {

        for (VirtualObjectNode obj : virtual) {

            effects.addVirtualMapping(frameState, state.getObjectState(obj).createEscapeObjectState(debug, obj));

        }

    }

    private void collectReferencedVirtualObjects(BlockT state, EconomicSet<VirtualObjectNode> virtual) {

        Iterator<VirtualObjectNode> iterator = virtual.iterator();

        while (iterator.hasNext()) {

            VirtualObjectNode object = iterator.next();

            int id = object.getObjectId();

            if (id != -1) {

                ObjectState objState = state.getObjectStateOptional(id);

                if (objState != null && objState.isVirtual()) {

                    for (ValueNode entry : objState.getEntries()) {

                        if (entry instanceof VirtualObjectNode) {

                            VirtualObjectNode entryVirtual = (VirtualObjectNode) entry;

                            if (!virtual.contains(entryVirtual)) {

                                virtual.add(entryVirtual);

                            }

                        }

                    }

                }

            }

        }

    }

    private void collectLockedVirtualObjects(BlockT state, EconomicSet<VirtualObjectNode> virtual) {

        for (int i = 0; i < state.getStateCount(); i++) {

            ObjectState objState = state.getObjectStateOptional(i);

            if (objState != null && objState.isVirtual() && objState.hasLocks()) {

                virtual.add(virtualObjects.get(i));

            }

        }

    }

    /**

     * @return true if materialization happened, false if not.

     */

    protected boolean ensureMaterialized(PartialEscapeBlockState<?> state, int object, FixedNode materializeBefore, GraphEffectList effects, CounterKey counter) {

        if (state.getObjectState(object).isVirtual()) {

            counter.increment(debug);

            VirtualObjectNode virtual = virtualObjects.get(object);

            state.materializeBefore(materializeBefore, virtual, effects);

            assert !updateStatesForMaterialized(state, virtual, state.getObjectState(object).getMaterializedValue()) : "method must already have been called before";

            return true;

        } else {

            return false;

        }

    }

    public static boolean updateStatesForMaterialized(PartialEscapeBlockState<?> state, VirtualObjectNode virtual, ValueNode materializedValue) {

        // update all existing states with the newly materialized object

        boolean change = false;

        for (int i = 0; i < state.getStateCount(); i++) {

            ObjectState objState = state.getObjectStateOptional(i);

            if (objState != null && objState.isVirtual()) {

                ValueNode[] entries = objState.getEntries();

                for (int i2 = 0; i2 < entries.length; i2++) {

                    if (entries[i2] == virtual) {

                        state.setEntry(i, i2, materializedValue);

                        change = true;

                    }

                }

            }

        }

        return change;

    }

    @Override

    protected BlockT stripKilledLoopLocations(Loop<Block> loop, BlockT originalInitialState) {

        BlockT initialState = super.stripKilledLoopLocations(loop, originalInitialState);

        if (loop.getDepth() > GraalOptions.EscapeAnalysisLoopCutoff.getValue(cfg.graph.getOptions())) {

            /*

             * After we've reached the maximum loop nesting, we'll simply materialize everything we

             * can to make sure that the loops only need to be iterated one time. Care is taken here

             * to not materialize virtual objects that have the "ensureVirtualized" flag set.

             */

            LoopBeginNode loopBegin = (LoopBeginNode) loop.getHeader().getBeginNode();

            AbstractEndNode end = loopBegin.forwardEnd();

            Block loopPredecessor = loop.getHeader().getFirstPredecessor();

            assert loopPredecessor.getEndNode() == end;

            int length = initialState.getStateCount();

            boolean change;

            BitSet ensureVirtualized = new BitSet(length);

            for (int i = 0; i < length; i++) {

                ObjectState state = initialState.getObjectStateOptional(i);

                if (state != null && state.isVirtual() && state.getEnsureVirtualized()) {

                    ensureVirtualized.set(i);

                }

            }

            do {

                // propagate "ensureVirtualized" flag

                change = false;

                for (int i = 0; i < length; i++) {

                    if (!ensureVirtualized.get(i)) {

                        ObjectState state = initialState.getObjectStateOptional(i);

                        if (state != null && state.isVirtual()) {

                            for (ValueNode entry : state.getEntries()) {

                                if (entry instanceof VirtualObjectNode) {

                                    if (ensureVirtualized.get(((VirtualObjectNode) entry).getObjectId())) {

                                        change = true;

                                        ensureVirtualized.set(i);

                                        break;

                                    }

                                }

                            }

                        }

                    }

                }

            } while (change);

            for (int i = 0; i < length; i++) {

                ObjectState state = initialState.getObjectStateOptional(i);

                if (state != null && state.isVirtual() && !ensureVirtualized.get(i)) {

                    initialState.materializeBefore(end, virtualObjects.get(i), blockEffects.get(loopPredecessor));

                }

            }

        }

        return initialState;

    }

    @Override