/*
 * Copyright (c) 2014, 2019, 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.lir.stackslotalloc;

import static org.graalvm.compiler.debug.DebugContext.BASIC_LEVEL;
import static org.graalvm.compiler.lir.LIRValueUtil.asVirtualStackSlot;
import static org.graalvm.compiler.lir.LIRValueUtil.isVirtualStackSlot;
import static org.graalvm.compiler.lir.phases.LIRPhase.Options.LIROptimization;

import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Deque;
import java.util.EnumSet;
import java.util.PriorityQueue;
import java.util.function.Predicate;

import jdk.internal.vm.compiler.collections.EconomicSet;
import org.graalvm.compiler.core.common.LIRKind;
import org.graalvm.compiler.core.common.cfg.AbstractBlockBase;
import org.graalvm.compiler.debug.DebugCloseable;
import org.graalvm.compiler.debug.DebugContext;
import org.graalvm.compiler.debug.Indent;
import org.graalvm.compiler.debug.TimerKey;
import org.graalvm.compiler.lir.LIR;
import org.graalvm.compiler.lir.LIRInstruction;
import org.graalvm.compiler.lir.LIRInstruction.OperandFlag;
import org.graalvm.compiler.lir.LIRInstruction.OperandMode;
import org.graalvm.compiler.lir.ValueProcedure;
import org.graalvm.compiler.lir.VirtualStackSlot;
import org.graalvm.compiler.lir.framemap.FrameMap;
import org.graalvm.compiler.lir.framemap.FrameMapBuilderTool;
import org.graalvm.compiler.lir.framemap.SimpleVirtualStackSlot;
import org.graalvm.compiler.lir.framemap.VirtualStackSlotRange;
import org.graalvm.compiler.lir.gen.LIRGenerationResult;
import org.graalvm.compiler.lir.phases.AllocationPhase;
import org.graalvm.compiler.options.NestedBooleanOptionKey;
import org.graalvm.compiler.options.Option;
import org.graalvm.compiler.options.OptionType;

import jdk.vm.ci.code.CodeUtil;
import jdk.vm.ci.code.StackSlot;
import jdk.vm.ci.code.TargetDescription;
import jdk.vm.ci.meta.Value;
import jdk.vm.ci.meta.ValueKind;

/**
 * Linear Scan {@link StackSlotAllocatorUtil stack slot allocator}.
 * <p>
 * <b>Remark:</b> The analysis works under the assumption that a stack slot is no longer live after
 * its last usage. If an {@link LIRInstruction instruction} transfers the raw address of the stack
 * slot to another location, e.g. a registers, and this location is referenced later on, the
 * {@link org.graalvm.compiler.lir.LIRInstruction.Use usage} of the stack slot must be marked with
 * the {@link OperandFlag#UNINITIALIZED}. Otherwise the stack slot might be reused and its content
 * destroyed.
 */
public final class LSStackSlotAllocator extends AllocationPhase {

    public static class Options {
        // @formatter:off
        @Option(help = "Use linear scan stack slot allocation.", type = OptionType.Debug)
        public static final NestedBooleanOptionKey LIROptLSStackSlotAllocator = new NestedBooleanOptionKey(LIROptimization, true);
        // @formatter:on
    }

    private static final TimerKey MainTimer = DebugContext.timer("LSStackSlotAllocator");
    private static final TimerKey NumInstTimer = DebugContext.timer("LSStackSlotAllocator[NumberInstruction]");
    private static final TimerKey BuildIntervalsTimer = DebugContext.timer("LSStackSlotAllocator[BuildIntervals]");
    private static final TimerKey VerifyIntervalsTimer = DebugContext.timer("LSStackSlotAllocator[VerifyIntervals]");
    private static final TimerKey AllocateSlotsTimer = DebugContext.timer("LSStackSlotAllocator[AllocateSlots]");
    private static final TimerKey AssignSlotsTimer = DebugContext.timer("LSStackSlotAllocator[AssignSlots]");

    @Override
    protected void run(TargetDescription target, LIRGenerationResult lirGenRes, AllocationContext context) {
        allocateStackSlots((FrameMapBuilderTool) lirGenRes.getFrameMapBuilder(), lirGenRes);
        lirGenRes.buildFrameMap();
    }

    @SuppressWarnings("try")
    public static void allocateStackSlots(FrameMapBuilderTool builder, LIRGenerationResult res) {
        if (builder.getNumberOfStackSlots() > 0) {
            try (DebugCloseable t = MainTimer.start(res.getLIR().getDebug())) {
                new Allocator(res.getLIR(), builder).allocate();
            }
        }
    }

    private static final class Allocator {

        private final LIR lir;
        private final DebugContext debug;
        private final FrameMapBuilderTool frameMapBuilder;
        private final StackInterval[] stackSlotMap;
        private final PriorityQueue<StackInterval> unhandled;
        private final PriorityQueue<StackInterval> active;
        private final AbstractBlockBase<?>[] sortedBlocks;
        private final int maxOpId;

        @SuppressWarnings("try")
        private Allocator(LIR lir, FrameMapBuilderTool frameMapBuilder) {
            this.lir = lir;
            this.debug = lir.getDebug();
            this.frameMapBuilder = frameMapBuilder;
            this.stackSlotMap = new StackInterval[frameMapBuilder.getNumberOfStackSlots()];
            this.sortedBlocks = lir.getControlFlowGraph().getBlocks();

            // insert by from
            this.unhandled = new PriorityQueue<>((a, b) -> a.from() - b.from());
            // insert by to
            this.active = new PriorityQueue<>((a, b) -> a.to() - b.to());

            try (DebugCloseable t = NumInstTimer.start(debug)) {
                // step 1: number instructions
                this.maxOpId = numberInstructions(lir, sortedBlocks);
            }
        }

        @SuppressWarnings("try")
        private void allocate() {
            debug.dump(DebugContext.VERBOSE_LEVEL, lir, "After StackSlot numbering");

            boolean allocationFramesizeEnabled = StackSlotAllocatorUtil.allocatedFramesize.isEnabled(debug);
            long currentFrameSize = allocationFramesizeEnabled ? frameMapBuilder.getFrameMap().currentFrameSize() : 0;
            EconomicSet<LIRInstruction> usePos;
            // step 2: build intervals
            try (DebugContext.Scope s = debug.scope("StackSlotAllocationBuildIntervals"); Indent indent = debug.logAndIndent("BuildIntervals"); DebugCloseable t = BuildIntervalsTimer.start(debug)) {
                usePos = buildIntervals();
            }
            // step 3: verify intervals
            if (debug.areScopesEnabled()) {
                try (DebugCloseable t = VerifyIntervalsTimer.start(debug)) {
                    assert verifyIntervals();
                }
            }
            if (debug.isDumpEnabled(DebugContext.VERBOSE_LEVEL)) {
                dumpIntervals("Before stack slot allocation");
            }
            // step 4: allocate stack slots
            try (DebugCloseable t = AllocateSlotsTimer.start(debug)) {
                /*
                 * Allocate primitive spill slots before reference spill slots. This ensures a
                 * ReferenceMap will be as compact as possible and only exceed the encoding limit of
                 * a stack offset if there are really too many objects live on the stack at an
                 * instruction with a ReferenceMap (as opposed to the method simply having a very
                 * large frame).
                 */
                allocateStackSlots(IS_PRIMITIVE_INTERVAL);
                allocateStackSlots(IS_REFERENCE_INTERVAL);
            }
            if (debug.isDumpEnabled(DebugContext.VERBOSE_LEVEL)) {
                dumpIntervals("After stack slot allocation");
            }

            // step 5: assign stack slots
            try (DebugCloseable t = AssignSlotsTimer.start(debug)) {
                assignStackSlots(usePos);
            }
            if (allocationFramesizeEnabled) {
                StackSlotAllocatorUtil.allocatedFramesize.add(debug, frameMapBuilder.getFrameMap().currentFrameSize() - currentFrameSize);
            }
        }

        // ====================
        // step 1: number instructions
        // ====================

        /**
         * Numbers all instructions in all blocks.
         *
         * @return The id of the last operation.
         */
        private static int numberInstructions(LIR lir, AbstractBlockBase<?>[] sortedBlocks) {
            int opId = 0;
            int index = 0;
            for (AbstractBlockBase<?> block : sortedBlocks) {

                ArrayList<LIRInstruction> instructions = lir.getLIRforBlock(block);

                int numInst = instructions.size();
                for (int j = 0; j < numInst; j++) {
                    LIRInstruction op = instructions.get(j);
                    op.setId(opId);

                    index++;
                    opId += 2; // numbering of lirOps by two
                }
            }
            assert (index << 1) == opId : "must match: " + (index << 1);
            return opId - 2;
        }

        // ====================
        // step 2: build intervals
        // ====================

        private EconomicSet<LIRInstruction> buildIntervals() {
            return new FixPointIntervalBuilder(lir, stackSlotMap, maxOpId()).build();
        }

        // ====================
        // step 3: verify intervals
        // ====================

        private boolean verifyIntervals() {
            for (StackInterval interval : stackSlotMap) {
                if (interval != null) {
                    assert interval.verify(maxOpId());
                }
            }
            return true;
        }

        // ====================
        // step 4: allocate stack slots
        // ====================

        @SuppressWarnings("try")
        private void allocateStackSlots(Predicate<StackInterval> predicate) {
            for (StackInterval interval : stackSlotMap) {
                if (interval != null && (predicate == null || predicate.test(interval))) {
                    unhandled.add(interval);
                }
            }
            for (StackInterval current = activateNext(); current != null; current = activateNext()) {
                try (Indent indent = debug.logAndIndent("allocate %s", current)) {
                    allocateSlot(current);
                }
            }

            // Cannot re-use free slots between rounds of slot allocation
            freeSlots = null;
            active.clear();
        }

        private static final Predicate<StackInterval> IS_REFERENCE_INTERVAL = new Predicate<StackInterval>() {
            @Override
            public boolean test(StackInterval interval) {
                return !((LIRKind) interval.kind()).isValue();
            }
        };

        private static final Predicate<StackInterval> IS_PRIMITIVE_INTERVAL = new Predicate<StackInterval>() {
            @Override
            public boolean test(StackInterval interval) {
                return ((LIRKind) interval.kind()).isValue();
            }
        };

        private void allocateSlot(StackInterval current) {
            VirtualStackSlot virtualSlot = current.getOperand();
            final StackSlot location;
            if (virtualSlot instanceof VirtualStackSlotRange) {
                // No reuse of ranges (yet).
                VirtualStackSlotRange slotRange = (VirtualStackSlotRange) virtualSlot;
                location = frameMapBuilder.getFrameMap().allocateStackSlots(slotRange.getSlots());
                StackSlotAllocatorUtil.virtualFramesize.add(debug, frameMapBuilder.getFrameMap().spillSlotRangeSize(slotRange.getSlots()));
                StackSlotAllocatorUtil.allocatedSlots.increment(debug);
            } else {
                assert virtualSlot instanceof SimpleVirtualStackSlot : "Unexpected VirtualStackSlot type: " + virtualSlot;
                StackSlot slot = findFreeSlot((SimpleVirtualStackSlot) virtualSlot);
                if (slot != null) {
                    /*
                     * Free stack slot available. Note that we create a new one because the kind
                     * might not match.
                     */
                    location = StackSlot.get(current.kind(), slot.getRawOffset(), slot.getRawAddFrameSize());
                    StackSlotAllocatorUtil.reusedSlots.increment(debug);
                    debug.log(BASIC_LEVEL, "Reuse stack slot %s (reallocated from %s) for virtual stack slot %s", location, slot, virtualSlot);
                } else {
                    // Allocate new stack slot.
                    location = frameMapBuilder.getFrameMap().allocateSpillSlot(virtualSlot.getValueKind());
                    StackSlotAllocatorUtil.virtualFramesize.add(debug, frameMapBuilder.getFrameMap().spillSlotSize(virtualSlot.getValueKind()));
                    StackSlotAllocatorUtil.allocatedSlots.increment(debug);
                    debug.log(BASIC_LEVEL, "New stack slot %s for virtual stack slot %s", location, virtualSlot);
                }
            }
            debug.log("Allocate location %s for interval %s", location, current);
            current.setLocation(location);
        }

        /**
         * Map from log2 of {@link FrameMap#spillSlotSize(ValueKind) a spill slot size} to a list of
         * free stack slots.
         */
        private ArrayList<Deque<StackSlot>> freeSlots;

        /**
         * @return The list of free stack slots for {@code index} or {@code null} if there is none.
         */
        private Deque<StackSlot> getNullOrFreeSlots(int index) {
            if (freeSlots == null) {
                return null;
            }
            if (index < freeSlots.size()) {
                return freeSlots.get(index);
            }
            return null;
        }

        /**
         * @return the list of free stack slots for {@code index}. If there is none a list is
         *         created.
         */
        private Deque<StackSlot> getOrInitFreeSlots(int index) {
            Deque<StackSlot> freeList = null;
            if (freeSlots == null) {
                freeSlots = new ArrayList<>(6);
            } else if (index < freeSlots.size()) {
                freeList = freeSlots.get(index);
            }
            if (freeList == null) {
                int requiredSize = index + 1;
                for (int i = freeSlots.size(); i < requiredSize; i++) {
                    freeSlots.add(null);
                }
                freeList = new ArrayDeque<>();
                freeSlots.set(index, freeList);
            }
            return freeList;
        }

        /**
         * Gets a free stack slot for {@code slot} or {@code null} if there is none.
         */
        private StackSlot findFreeSlot(SimpleVirtualStackSlot slot) {
            assert slot != null;
            int size = log2SpillSlotSize(slot.getValueKind());
            Deque<StackSlot> freeList = getNullOrFreeSlots(size);
            if (freeList == null) {
                return null;
            }
            return freeList.pollLast();
        }

        /**
         * Adds a stack slot to the list of free slots.
         */
        private void freeSlot(StackSlot slot) {
            int size = log2SpillSlotSize(slot.getValueKind());
            getOrInitFreeSlots(size).addLast(slot);
        }

        private int log2SpillSlotSize(ValueKind<?> kind) {
            int size = frameMapBuilder.getFrameMap().spillSlotSize(kind);
            assert CodeUtil.isPowerOf2(size);
            return CodeUtil.log2(size);
        }

        /**
         * Gets the next unhandled interval and finishes handled intervals.
         */
        private StackInterval activateNext() {
            if (unhandled.isEmpty()) {
                return null;
            }
            StackInterval next = unhandled.poll();
            // finish handled intervals
            for (int id = next.from(); activePeekId() < id;) {
                finished(active.poll());
            }
            debug.log("active %s", next);
            active.add(next);
            return next;
        }

        /**
         * Gets the lowest {@link StackInterval#to() end position} of all active intervals. If there
         * is none {@link Integer#MAX_VALUE} is returned.
         */
        private int activePeekId() {
            StackInterval first = active.peek();
            if (first == null) {
                return Integer.MAX_VALUE;
            }
            return first.to();
        }

        /**
         * Finishes {@code interval} by adding its location to the list of free stack slots.
         */
        private void finished(StackInterval interval) {
            StackSlot location = interval.location();
            debug.log("finished %s (freeing %s)", interval, location);
            freeSlot(location);
        }

        // ====================
        // step 5: assign stack slots
        // ====================

        private void assignStackSlots(EconomicSet<LIRInstruction> usePos) {
            for (LIRInstruction op : usePos) {
                op.forEachInput(assignSlot);
                op.forEachAlive(assignSlot);
                op.forEachState(assignSlot);

                op.forEachTemp(assignSlot);
                op.forEachOutput(assignSlot);
            }
        }

        ValueProcedure assignSlot = new ValueProcedure() {
            @Override
            public Value doValue(Value value, OperandMode mode, EnumSet<OperandFlag> flags) {
                if (isVirtualStackSlot(value)) {
                    VirtualStackSlot slot = asVirtualStackSlot(value);
                    StackInterval interval = get(slot);
                    assert interval != null;
                    return interval.location();
                }
                return value;
            }
        };

        // ====================
        //
        // ====================

        /**
         * Gets the highest instruction id.
         */
        private int maxOpId() {
            return maxOpId;
        }

        private StackInterval get(VirtualStackSlot stackSlot) {
            return stackSlotMap[stackSlot.getId()];
        }

        private void dumpIntervals(String label) {
            debug.dump(DebugContext.VERBOSE_LEVEL, new StackIntervalDumper(Arrays.copyOf(stackSlotMap, stackSlotMap.length)), label);
        }

    }
}