/*

 * Copyright (c) 2013, 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.asm.aarch64;

import static jdk.vm.ci.aarch64.AArch64.CPU;

import static jdk.vm.ci.aarch64.AArch64.rscratch1;

import static jdk.vm.ci.aarch64.AArch64.rscratch2;

import static jdk.vm.ci.aarch64.AArch64.sp;

import static jdk.vm.ci.aarch64.AArch64.zr;

import static org.graalvm.compiler.asm.aarch64.AArch64Address.AddressingMode.BASE_REGISTER_ONLY;

import static org.graalvm.compiler.asm.aarch64.AArch64Address.AddressingMode.EXTENDED_REGISTER_OFFSET;

import static org.graalvm.compiler.asm.aarch64.AArch64Address.AddressingMode.IMMEDIATE_SCALED;

import static org.graalvm.compiler.asm.aarch64.AArch64Address.AddressingMode.IMMEDIATE_UNSCALED;

import static org.graalvm.compiler.asm.aarch64.AArch64Address.AddressingMode.REGISTER_OFFSET;

import static org.graalvm.compiler.asm.aarch64.AArch64MacroAssembler.AddressGenerationPlan.WorkPlan.ADD_TO_BASE;

import static org.graalvm.compiler.asm.aarch64.AArch64MacroAssembler.AddressGenerationPlan.WorkPlan.ADD_TO_INDEX;

import static org.graalvm.compiler.asm.aarch64.AArch64MacroAssembler.AddressGenerationPlan.WorkPlan.NO_WORK;

import org.graalvm.compiler.asm.BranchTargetOutOfBoundsException;

import org.graalvm.compiler.asm.Label;

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

import org.graalvm.compiler.debug.GraalError;

import jdk.vm.ci.aarch64.AArch64;

import jdk.vm.ci.code.Register;

import jdk.vm.ci.code.TargetDescription;

public class AArch64MacroAssembler extends AArch64Assembler {

    private final ScratchRegister[] scratchRegister = new ScratchRegister[]{new ScratchRegister(rscratch1), new ScratchRegister(rscratch2)};

    // Points to the next free scratch register

    private int nextFreeScratchRegister = 0;

    public AArch64MacroAssembler(TargetDescription target) {

        super(target);

    }

    public class ScratchRegister implements AutoCloseable {

        private final Register register;

        public ScratchRegister(Register register) {

            this.register = register;

        }

        public Register getRegister() {

            return register;

        }

        @Override

        public void close() {

            assert nextFreeScratchRegister > 0 : "Close called too often";

            nextFreeScratchRegister--;

        }

    }

    public ScratchRegister getScratchRegister() {

        return scratchRegister[nextFreeScratchRegister++];

    }

    /**

     * Specifies what actions have to be taken to turn an arbitrary address of the form

     * {@code base + displacement [+ index [<< scale]]} into a valid AArch64Address.

     */

    public static class AddressGenerationPlan {

        public final WorkPlan workPlan;

        public final AArch64Address.AddressingMode addressingMode;

        public final boolean needsScratch;

        public enum WorkPlan {

            /**

             * Can be used as-is without extra work.

             */

            NO_WORK,

            /**

             * Add scaled displacement to index register.

             */

            ADD_TO_INDEX,

            /**

             * Add unscaled displacement to base register.

             */

            ADD_TO_BASE,

        }

        /**

         * @param workPlan Work necessary to generate a valid address.

         * @param addressingMode Addressing mode of generated address.

         * @param needsScratch True if generating address needs a scatch register, false otherwise.

         */

        public AddressGenerationPlan(WorkPlan workPlan, AArch64Address.AddressingMode addressingMode, boolean needsScratch) {

            this.workPlan = workPlan;

            this.addressingMode = addressingMode;

            this.needsScratch = needsScratch;

        }

    }

    /**

     * Generates an addressplan for an address of the form

     * {@code base + displacement [+ index [<< log2(transferSize)]]} with the index register and

     * scaling being optional.

     *

     * @param displacement an arbitrary displacement.

     * @param hasIndexRegister true if the address uses an index register, false otherwise. non null

     * @param transferSize the memory transfer size in bytes. The log2 of this specifies how much

     *            the index register is scaled. If 0 no scaling is assumed. Can be 0, 1, 2, 4 or 8.

     * @return AddressGenerationPlan that specifies the actions necessary to generate a valid

     *         AArch64Address for the given parameters.

     */

    public static AddressGenerationPlan generateAddressPlan(long displacement, boolean hasIndexRegister, int transferSize) {

        assert transferSize == 0 || transferSize == 1 || transferSize == 2 || transferSize == 4 || transferSize == 8;

        boolean indexScaled = transferSize != 0;

        int log2Scale = NumUtil.log2Ceil(transferSize);

        long scaledDisplacement = displacement >> log2Scale;

        boolean displacementScalable = indexScaled && (displacement & (transferSize - 1)) == 0;

        if (displacement == 0) {

            // register offset without any work beforehand.

            return new AddressGenerationPlan(NO_WORK, REGISTER_OFFSET, false);

        } else {

            if (hasIndexRegister) {

                if (displacementScalable) {

                    boolean needsScratch = !isArithmeticImmediate(scaledDisplacement);

                    return new AddressGenerationPlan(ADD_TO_INDEX, REGISTER_OFFSET, needsScratch);

                } else {

                    boolean needsScratch = !isArithmeticImmediate(displacement);

                    return new AddressGenerationPlan(ADD_TO_BASE, REGISTER_OFFSET, needsScratch);

                }

            } else {

                if (displacementScalable && NumUtil.isUnsignedNbit(12, scaledDisplacement)) {

                    return new AddressGenerationPlan(NO_WORK, IMMEDIATE_SCALED, false);

                } else if (NumUtil.isSignedNbit(9, displacement)) {

                    return new AddressGenerationPlan(NO_WORK, IMMEDIATE_UNSCALED, false);

                } else {

                    boolean needsScratch = !isArithmeticImmediate(displacement);

                    return new AddressGenerationPlan(ADD_TO_BASE, REGISTER_OFFSET, needsScratch);

                }

            }

        }

    }

    /**

     * Returns an AArch64Address pointing to

     * {@code base + displacement + index << log2(transferSize)}.

     *

     * @param base general purpose register. May not be null or the zero register.

     * @param displacement arbitrary displacement added to base.

     * @param index general purpose register. May not be null or the stack pointer.

     * @param signExtendIndex if true consider index register a word register that should be

     *            sign-extended before being added.

     * @param transferSize the memory transfer size in bytes. The log2 of this specifies how much

     *            the index register is scaled. If 0 no scaling is assumed. Can be 0, 1, 2, 4 or 8.

     * @param additionalReg additional register used either as a scratch register or as part of the

     *            final address, depending on whether allowOverwrite is true or not. May not be null

     *            or stackpointer.

     * @param allowOverwrite if true allows to change value of base or index register to generate

     *            address.

     * @return AArch64Address pointing to memory at

     *         {@code base + displacement + index << log2(transferSize)}.

     */

    public AArch64Address makeAddress(Register base, long displacement, Register index, boolean signExtendIndex, int transferSize, Register additionalReg, boolean allowOverwrite) {

        AddressGenerationPlan plan = generateAddressPlan(displacement, !index.equals(zr), transferSize);

        assert allowOverwrite || !zr.equals(additionalReg) || plan.workPlan == NO_WORK;

        assert !plan.needsScratch || !zr.equals(additionalReg);

        int log2Scale = NumUtil.log2Ceil(transferSize);

        long scaledDisplacement = displacement >> log2Scale;

        Register newIndex = index;

        Register newBase = base;

        int immediate;

        switch (plan.workPlan) {

            case NO_WORK:

                if (plan.addressingMode == IMMEDIATE_SCALED) {

                    immediate = (int) scaledDisplacement;

                } else {

                    immediate = (int) displacement;

                }

                break;

            case ADD_TO_INDEX:

                newIndex = allowOverwrite ? index : additionalReg;

                assert !newIndex.equals(sp) && !newIndex.equals(zr);

                if (plan.needsScratch) {

                    mov(additionalReg, scaledDisplacement);

                    add(signExtendIndex ? 32 : 64, newIndex, index, additionalReg);

                } else {

                    add(signExtendIndex ? 32 : 64, newIndex, index, (int) scaledDisplacement);

                }

                immediate = 0;

                break;

            case ADD_TO_BASE:

                newBase = allowOverwrite ? base : additionalReg;

                assert !newBase.equals(sp) && !newBase.equals(zr);

                if (plan.needsScratch) {

                    mov(additionalReg, displacement);

                    add(64, newBase, base, additionalReg);

                } else {

                    add(64, newBase, base, (int) displacement);

                }

                immediate = 0;

                break;

            default:

                throw GraalError.shouldNotReachHere();

        }

        AArch64Address.AddressingMode addressingMode = plan.addressingMode;

        ExtendType extendType = null;

        if (addressingMode == REGISTER_OFFSET) {

            if (newIndex.equals(zr)) {

                addressingMode = BASE_REGISTER_ONLY;

            } else if (signExtendIndex) {

                addressingMode = EXTENDED_REGISTER_OFFSET;

                extendType = ExtendType.SXTW;

            }

        }

        return AArch64Address.createAddress(addressingMode, newBase, newIndex, immediate, transferSize != 0, extendType);

    }

    /**

     * Returns an AArch64Address pointing to {@code base + displacement}. Specifies the memory

     * transfer size to allow some optimizations when building the address.

     *

     * @param base general purpose register. May not be null or the zero register.

     * @param displacement arbitrary displacement added to base.

     * @param transferSize the memory transfer size in bytes.

     * @param additionalReg additional register used either as a scratch register or as part of the

     *            final address, depending on whether allowOverwrite is true or not. May not be

     *            null, zero register or stackpointer.

     * @param allowOverwrite if true allows to change value of base or index register to generate

     *            address.

     * @return AArch64Address pointing to memory at {@code base + displacement}.

     */

    public AArch64Address makeAddress(Register base, long displacement, Register additionalReg, int transferSize, boolean allowOverwrite) {

        assert additionalReg.getRegisterCategory().equals(CPU);

        return makeAddress(base, displacement, zr, /* sign-extend */false, transferSize, additionalReg, allowOverwrite);

    }

    /**

     * Returns an AArch64Address pointing to {@code base + displacement}. Fails if address cannot be

     * represented without overwriting base register or using a scratch register.

     *

     * @param base general purpose register. May not be null or the zero register.

     * @param displacement arbitrary displacement added to base.

     * @param transferSize the memory transfer size in bytes. The log2 of this specifies how much

     *            the index register is scaled. If 0 no scaling is assumed. Can be 0, 1, 2, 4 or 8.

     * @return AArch64Address pointing to memory at {@code base + displacement}.

     */

    public AArch64Address makeAddress(Register base, long displacement, int transferSize) {

        return makeAddress(base, displacement, zr, /* signExtend */false, //

                        transferSize, zr, /* allowOverwrite */false);

    }

    /**

     * Loads memory address into register.

     *

     * @param dst general purpose register. May not be null, zero-register or stackpointer.

     * @param address address whose value is loaded into dst. May not be null,

     *            {@link org.graalvm.compiler.asm.aarch64.AArch64Address.AddressingMode#IMMEDIATE_POST_INDEXED

     *            POST_INDEXED} or

     *            {@link org.graalvm.compiler.asm.aarch64.AArch64Address.AddressingMode#IMMEDIATE_PRE_INDEXED

     *            IMMEDIATE_PRE_INDEXED}

     * @param transferSize the memory transfer size in bytes. The log2 of this specifies how much

     *            the index register is scaled. Can be 1, 2, 4 or 8.

     */

    public void loadAddress(Register dst, AArch64Address address, int transferSize) {

        assert transferSize == 1 || transferSize == 2 || transferSize == 4 || transferSize == 8;

        assert dst.getRegisterCategory().equals(CPU);

        int shiftAmt = NumUtil.log2Ceil(transferSize);

        switch (address.getAddressingMode()) {

            case IMMEDIATE_SCALED:

                int scaledImmediate = address.getImmediateRaw() << shiftAmt;

                int lowerBits = scaledImmediate & NumUtil.getNbitNumberInt(12);

                int higherBits = scaledImmediate & ~NumUtil.getNbitNumberInt(12);

                boolean firstAdd = true;

                if (lowerBits != 0) {

                    add(64, dst, address.getBase(), lowerBits);

                    firstAdd = false;

                }

                if (higherBits != 0) {

                    Register src = firstAdd ? address.getBase() : dst;

                    add(64, dst, src, higherBits);

                }

                break;

            case IMMEDIATE_UNSCALED:

                int immediate = address.getImmediateRaw();

                add(64, dst, address.getBase(), immediate);

                break;

            case REGISTER_OFFSET:

                add(64, dst, address.getBase(), address.getOffset(), ShiftType.LSL, address.isScaled() ? shiftAmt : 0);

                break;

            case EXTENDED_REGISTER_OFFSET:

                add(64, dst, address.getBase(), address.getOffset(), address.getExtendType(), address.isScaled() ? shiftAmt : 0);

                break;

            case PC_LITERAL: {

                addressOf(dst);

                break;

            }

            case BASE_REGISTER_ONLY:

                movx(dst, address.getBase());

                break;

            default:

                throw GraalError.shouldNotReachHere();

        }

    }

    public void movx(Register dst, Register src) {

        mov(64, dst, src);

    }

    public void mov(int size, Register dst, Register src) {

        if (dst.equals(sp) || src.equals(sp)) {

            add(size, dst, src, 0);

        } else {

            or(size, dst, zr, src);

        }

    }

    /**

     * Generates a 64-bit immediate move code sequence.

     *

     * @param dst general purpose register. May not be null, stackpointer or zero-register.

     * @param imm the value to move into the register

     * @param annotateImm Flag denoting if annotation should be added.

     */

    private void mov64(Register dst, long imm, boolean annotateImm) {

        // We have to move all non zero parts of the immediate in 16-bit chunks

        int numMovs = 0;

        int pos = position();

        boolean firstMove = true;

        for (int offset = 0; offset < 64; offset += 16) {

            int chunk = (int) (imm >> offset) & NumUtil.getNbitNumberInt(16);

            if (chunk == 0) {

                continue;

            }

            if (firstMove) {

                movz(64, dst, chunk, offset);

                firstMove = false;

            } else {

                movk(64, dst, chunk, offset);

            }

            ++numMovs;

        }

        assert !firstMove;

        if (annotateImm) {

            annotateImmediateMovSequence(pos, numMovs);

        }

    }

    /**

     * Loads immediate into register.

     *

     * @param dst general purpose register. May not be null, zero-register or stackpointer.

     * @param imm immediate loaded into register.

     */

    public void mov(Register dst, long imm) {

        mov(dst, imm, false);

    }

    /**

     * Loads immediate into register.

     *

     * @param dst general purpose register. May not be null, zero-register or stackpointer.

     * @param imm immediate loaded into register.

     * @param annotateImm Flag to signal of the immediate value should be annotated.

     */

    public void mov(Register dst, long imm, boolean annotateImm) {

        assert dst.getRegisterCategory().equals(CPU);

        if (imm == 0L) {

            movx(dst, zr);

        } else if (LogicalImmediateTable.isRepresentable(true, imm) != LogicalImmediateTable.Representable.NO) {

            or(64, dst, zr, imm);

        } else if (imm >> 32 == -1L && (int) imm < 0 && LogicalImmediateTable.isRepresentable((int) imm) != LogicalImmediateTable.Representable.NO) {

            // If the higher 32-bit are 1s and the sign bit of the lower 32-bits is set *and* we can

            // represent the lower 32 bits as a logical immediate we can create the lower 32-bit and

            // then sign extend

            // them. This allows us to cover immediates like ~1L with 2 instructions.

            mov(dst, (int) imm);

            sxt(64, 32, dst, dst);

        } else {

            mov64(dst, imm, annotateImm);

        }

    }

    /**

     * Loads immediate into register.

     *

     * @param dst general purpose register. May not be null, zero-register or stackpointer.

     * @param imm immediate loaded into register.

     */

    public void mov(Register dst, int imm) {

        mov(dst, imm & 0xFFFF_FFFFL);

    }

    /**

     * Generates a 48-bit immediate move code sequence. The immediate may later be updated by

     * HotSpot.

     *

     * In AArch64 mode the virtual address space is 48-bits in size, so we only need three

     * instructions to create a patchable instruction sequence that can reach anywhere.

     *

     * @param dst general purpose register. May not be null, stackpointer or zero-register.

     * @param imm

     */

    public void movNativeAddress(Register dst, long imm) {

        movNativeAddress(dst, imm, false);

    }

    /**

     * Generates a 48-bit immediate move code sequence. The immediate may later be updated by

     * HotSpot.

     *

     * In AArch64 mode the virtual address space is 48-bits in size, so we only need three

     * instructions to create a patchable instruction sequence that can reach anywhere.

     *

     * @param dst general purpose register. May not be null, stackpointer or zero-register.

     * @param imm The immediate address

     * @param annotateImm Flag to signal of the immediate value should be annotated.

     */

    public void movNativeAddress(Register dst, long imm, boolean annotateImm) {

        assert (imm & 0xFFFF_0000_0000_0000L) == 0;

        // We have to move all non zero parts of the immediate in 16-bit chunks

        boolean firstMove = true;

        int pos = position();

        for (int offset = 0; offset < 48; offset += 16) {

            int chunk = (int) (imm >> offset) & NumUtil.getNbitNumberInt(16);

            if (firstMove) {

                movz(64, dst, chunk, offset);

                firstMove = false;

            } else {

                movk(64, dst, chunk, offset);

            }

        }

        if (annotateImm) {

            annotateImmediateMovSequence(pos, 3);

        }

        assert !firstMove;

    }

    /**

     * Generates a 32-bit immediate move code sequence. The immediate may later be updated by

     * HotSpot.

     *

     * @param dst general purpose register. May not be null, stackpointer or zero-register.

     * @param imm

     */

    public void movNarrowAddress(Register dst, long imm) {

        assert (imm & 0xFFFF_FFFF_0000_0000L) == 0;

        movz(64, dst, (int) (imm >>> 16), 16);

        movk(64, dst, (int) (imm & 0xffff), 0);

    }

    /**

     * @return Number of instructions necessary to load immediate into register.

     */

    public static int nrInstructionsToMoveImmediate(long imm) {

        if (imm == 0L || LogicalImmediateTable.isRepresentable(true, imm) != LogicalImmediateTable.Representable.NO) {

            return 1;

        }

        if (imm >> 32 == -1L && (int) imm < 0 && LogicalImmediateTable.isRepresentable((int) imm) != LogicalImmediateTable.Representable.NO) {

            // If the higher 32-bit are 1s and the sign bit of the lower 32-bits is set *and* we can

            // represent the lower 32 bits as a logical immediate we can create the lower 32-bit and

            // then sign extend

            // them. This allows us to cover immediates like ~1L with 2 instructions.

            return 2;

        }

        int nrInstructions = 0;

        for (int offset = 0; offset < 64; offset += 16) {

            int part = (int) (imm >> offset) & NumUtil.getNbitNumberInt(16);

            if (part != 0) {

                nrInstructions++;

            }

        }

        return nrInstructions;

    }

    /**

     * Loads a srcSize value from address into rt sign-extending it if necessary.

     *

     * @param targetSize size of target register in bits. Must be 32 or 64.

     * @param srcSize size of memory read in bits. Must be 8, 16 or 32 and smaller or equal to

     *            targetSize.

     * @param rt general purpose register. May not be null or stackpointer.

     * @param address all addressing modes allowed. May not be null.

     */

    @Override

    public void ldrs(int targetSize, int srcSize, Register rt, AArch64Address address) {

        assert targetSize == 32 || targetSize == 64;

        assert srcSize <= targetSize;

        if (targetSize == srcSize) {

            super.ldr(srcSize, rt, address);

        } else {

            super.ldrs(targetSize, srcSize, rt, address);

        }

    }

    /**

     * Loads a srcSize value from address into rt zero-extending it if necessary.

     *

     * @param srcSize size of memory read in bits. Must be 8, 16 or 32 and smaller or equal to

     *            targetSize.

     * @param rt general purpose register. May not be null or stackpointer.

     * @param address all addressing modes allowed. May not be null.