/*
 * Copyright (c) 2019, 2024, 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.
 *
 */

/*
 * @test
 * @run testng TestSegmentCopy
 */

import java.lang.foreign.Arena;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;
import java.lang.invoke.VarHandle;
import java.nio.ByteOrder;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.IntFunction;

import org.testng.SkipException;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;

import static java.lang.foreign.ValueLayout.JAVA_BYTE;
import static org.testng.Assert.*;

public class TestSegmentCopy {

    static final int TEST_BYTE_SIZE = 16;

    @Test(dataProvider = "segmentKinds")
    public void testByteCopy(SegmentKind kind1, SegmentKind kind2) {
        MemorySegment s1 = kind1.makeSegment(TEST_BYTE_SIZE);
        MemorySegment s2 = kind2.makeSegment(TEST_BYTE_SIZE);

        // for all offsets
        for (int s1Offset = 0; s1Offset < s1.byteSize(); s1Offset++) {
            for (int s2Offset = 0; s2Offset < s2.byteSize(); s2Offset++) {
                long slice1ByteSize = s1.byteSize() - s1Offset;
                long slice2ByteSize = s2.byteSize() - s2Offset;

                long copySize = Math.min(slice1ByteSize, slice2ByteSize);

                //prepare source slice
                for (int i = 0 ; i < copySize; i++) {
                    Type.BYTE.set(s1, s1Offset, i, i);
                }
                //perform copy
                MemorySegment.copy(s1, Type.BYTE.layout, s1Offset, s2, Type.BYTE.layout, s2Offset, copySize);
                //check that copy actually worked
                for (int i = 0; i < copySize; i++) {
                    Type.BYTE.check(s2, s2Offset, i, i);
                }
            }
        }
    }

    @Test(dataProvider = "conjunctSegments")
    public void testCopy5ArgInvariants(MemorySegment src, MemorySegment dst) {
        assertThrows(IndexOutOfBoundsException.class, () -> MemorySegment.copy(src, 0, dst, 0, -1));
        assertThrows(IndexOutOfBoundsException.class, () -> MemorySegment.copy(src, -1, dst, 0, src.byteSize()));
        assertThrows(IndexOutOfBoundsException.class, () -> MemorySegment.copy(src, 0, dst, -1, src.byteSize()));
        assertThrows(IndexOutOfBoundsException.class, () -> MemorySegment.copy(src, 1, dst, 0, src.byteSize()));
        assertThrows(IndexOutOfBoundsException.class, () -> MemorySegment.copy(src, 0, dst, 1, src.byteSize()));
    }

    @Test(dataProvider = "conjunctSegments")
    public void testConjunctCopy7ArgRight(MemorySegment src, MemorySegment dst) {
        testConjunctCopy(src, 0, dst, 1, CopyOp.of7Arg());
    }

    @Test(dataProvider = "conjunctSegments")
    public void testConjunctCopy5ArgRight(MemorySegment src, MemorySegment dst) {
        testConjunctCopy(src, 0, dst, 1, CopyOp.of5Arg());
    }

    @Test(dataProvider = "conjunctSegments")
    public void testConjunctCopy7ArgLeft(MemorySegment src, MemorySegment dst) {
        testConjunctCopy(src, 1, dst, 0, CopyOp.of7Arg());
    }

    @Test(dataProvider = "conjunctSegments")
    public void testConjunctCopy5ArgLeft(MemorySegment src, MemorySegment dst) {
        testConjunctCopy(src, 1, dst, 0, CopyOp.of5Arg());
    }

    void testConjunctCopy(MemorySegment src, long srcOffset, MemorySegment dst, long dstOffset, CopyOp op) {
        if (src.byteSize() < 4 || src.address() != dst.address()) {
            // Only test larger segments where the skew is zero
            return;
        }

        try (var arena = Arena.ofConfined()) {
            // Create a disjoint segment for expected behavior
            MemorySegment disjoint = arena.allocate(dst.byteSize());
            disjoint.copyFrom(src);
            op.copy(src, srcOffset, disjoint, dstOffset, 3);
            byte[] expected = disjoint.toArray(JAVA_BYTE);

            // Do a conjoint copy
            op.copy(src, srcOffset, dst, dstOffset, 3);
            byte[] actual = dst.toArray(JAVA_BYTE);

            assertEquals(actual, expected);
        }
    }

    @FunctionalInterface
    interface CopyOp {
        void copy(MemorySegment src, long srcOffset, MemorySegment dst, long dstOffset, long bytes);

        static CopyOp of5Arg() {
            return MemorySegment::copy;
        }

        static CopyOp of7Arg() {
            return (MemorySegment src, long srcOffset, MemorySegment dst, long dstOffset, long bytes) ->
                    MemorySegment.copy(src, JAVA_BYTE, srcOffset, dst, JAVA_BYTE, dstOffset, bytes);
        }

    }

    @Test(dataProvider = "segmentKinds")
    public void testByteCopySizes(SegmentKind kind1, SegmentKind kind2) {

        record Offsets(int src, int dst){}

        for (Offsets offsets : List.of(new Offsets(3, 7), new Offsets(7, 3))) {
            for (int size = 0; size < 513; size++) {
                MemorySegment src = kind1.makeSegment(size + offsets.src());
                MemorySegment dst = kind2.makeSegment(size + offsets.dst());
                //prepare source slice
                for (int i = 0; i < size; i++) {
                    src.set(JAVA_BYTE, i + offsets.src(), (byte) i);
                }
                //perform copy
                MemorySegment.copy(src, offsets.src(), dst, offsets.dst(), size);
                //check that copy actually worked
                for (int i = 0; i < size; i++) {
                    assertEquals(dst.get(JAVA_BYTE, i + offsets.dst()), (byte) i);
                }
            }
        }
    }

    @Test(expectedExceptions = IllegalArgumentException.class, dataProvider = "segmentKinds")
    public void testReadOnlyCopy(SegmentKind kind1, SegmentKind kind2) {
        MemorySegment s1 = kind1.makeSegment(TEST_BYTE_SIZE);
        MemorySegment s2 = kind2.makeSegment(TEST_BYTE_SIZE);
        // check failure with read-only dest
        MemorySegment.copy(s1, Type.BYTE.layout, 0, s2.asReadOnly(), Type.BYTE.layout, 0, 0);
    }

    @Test(expectedExceptions = IllegalArgumentException.class,
            expectedExceptionsMessageRegExp = ".*Attempt to write a read-only segment.*")
    public void badCopy6Arg() {
        try (Arena scope = Arena.ofConfined()) {
            MemorySegment dest = scope.allocate(ValueLayout.JAVA_INT).asReadOnly();
            MemorySegment.copy(new int[1],0, dest, ValueLayout.JAVA_INT, 0 ,1); // should throw
        }
    }

    @Test(expectedExceptions = IndexOutOfBoundsException.class, dataProvider = "types")
    public void testBadOverflow(Type type) {
        if (type.layout.byteSize() > 1) {
            MemorySegment segment = MemorySegment.ofArray(new byte[100]);
            MemorySegment.copy(segment, type.layout, 0, segment, type.layout, 0, Long.MAX_VALUE);
        } else {
            throw new SkipException("Byte layouts do not overflow");
        }
    }

    @Test(dataProvider = "segmentKindsAndTypes")
    public void testElementCopy(SegmentKind kind1, SegmentKind kind2, Type type1, Type type2) {
        MemorySegment s1 = kind1.makeSegment(TEST_BYTE_SIZE);
        MemorySegment s2 = kind2.makeSegment(TEST_BYTE_SIZE);

        // for all offsets
        for (int s1Offset = 0; s1Offset < s1.byteSize(); s1Offset++) {
            for (int s2Offset = 0; s2Offset < s2.byteSize(); s2Offset++) {
                long slice1ByteSize = s1.byteSize() - s1Offset;
                long slice2ByteSize = s2.byteSize() - s2Offset;

                long slice1ElementSize = slice1ByteSize / type1.size();
                long slice2ElementSize = slice2ByteSize / type2.size();

                long copySize = Math.min(slice1ElementSize, slice2ElementSize);

                //prepare source slice
                for (int i = 0 ; i < copySize; i++) {
                    type1.set(s1, s1Offset, i, i);
                }
                //perform copy
                MemorySegment.copy(s1, type1.layout, s1Offset, s2, type2.layout, s2Offset, copySize);
                //check that copy actually worked
                for (int i = 0; i < copySize; i++) {
                    type2.check(s2, s2Offset, i, i);
                }
            }
        }
    }

    @Test(expectedExceptions = IllegalArgumentException.class)
    public void testHyperAlignedSrc() {
        MemorySegment segment = MemorySegment.ofArray(new byte[] {1, 2, 3, 4});
        MemorySegment.copy(segment, 0, segment, JAVA_BYTE.withByteAlignment(2), 0, 4);
    }

    @Test(expectedExceptions = IllegalArgumentException.class)
    public void testHyperAlignedDst() {
        MemorySegment segment = MemorySegment.ofArray(new byte[] {1, 2, 3, 4});
        MemorySegment.copy(segment, JAVA_BYTE.withByteAlignment(2), 0, segment, 0, 4);
    }

    @Test
    public void testCopy5ArgWithNegativeValues() {
        MemorySegment src = MemorySegment.ofArray(new byte[] {1, 2, 3, 4});
        MemorySegment dst = MemorySegment.ofArray(new byte[] {1, 2, 3, 4});
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, -1, dst, 0, 4)
        );
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, 0, dst, -1, 4)
        );
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, 0, dst, 0, -1)
        );
    }

    @Test
    public void testCopy7ArgWithNegativeValues() {
        MemorySegment src = MemorySegment.ofArray(new byte[] {1, 2, 3, 4});
        MemorySegment dst = MemorySegment.ofArray(new byte[] {1, 2, 3, 4});
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, JAVA_BYTE, -1, dst, JAVA_BYTE, 0, 4)
        );
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, JAVA_BYTE, 0, dst, JAVA_BYTE, -1, 4)
        );
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, JAVA_BYTE, 0, dst, JAVA_BYTE, 0, -1)
        );
    }

    @Test
    public void testCopyFromArrayWithNegativeValues() {
        MemorySegment src = MemorySegment.ofArray(new byte[] {1, 2, 3, 4});
        byte[] dst = new byte[] {1, 2, 3, 4};
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, JAVA_BYTE, -1, dst, 0, 4)
        );
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, JAVA_BYTE, 0, dst, -1, 4)
        );
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, JAVA_BYTE, 0, dst, 0, -1)
        );
    }

    @Test
    public void testCopyToArrayWithNegativeValues() {
        byte[] src = new byte[] {1, 2, 3, 4};
        MemorySegment dst = MemorySegment.ofArray(new byte[] {1, 2, 3, 4});
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, -1, dst, JAVA_BYTE, 0, 4)
        );
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, 0, dst, JAVA_BYTE, -1, 4)
        );
        assertThrows(IndexOutOfBoundsException.class, () ->
                MemorySegment.copy(src, 0, dst, JAVA_BYTE, 0, -1)
        );
    }

    enum Type {
        // Byte
        BYTE(byte.class, JAVA_BYTE, i -> (byte)i),
        //LE
        SHORT_LE(short.class, ValueLayout.JAVA_SHORT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), i -> (short)i),
        CHAR_LE(char.class, ValueLayout.JAVA_CHAR_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), i -> (char)i),
        INT_LE(int.class, ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), i -> i),
        FLOAT_LE(float.class, ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), i -> (float)i),
        LONG_LE(long.class, ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), i -> (long)i),
        DOUBLE_LE(double.class, ValueLayout.JAVA_DOUBLE_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN), i -> (double)i),
        //BE
        SHORT_BE(short.class, ValueLayout.JAVA_SHORT_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN), i -> (short)i),
        CHAR_BE(char.class, ValueLayout.JAVA_CHAR_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN), i -> (char)i),
        INT_BE(int.class, ValueLayout.JAVA_INT_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN), i -> i),
        FLOAT_BE(float.class, ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN), i -> (float)i),
        LONG_BE(long.class, ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN), i -> (long)i),
        DOUBLE_BE(double.class, ValueLayout.JAVA_DOUBLE_UNALIGNED.withOrder(ByteOrder.BIG_ENDIAN), i -> (double)i);

        final ValueLayout layout;
        final IntFunction<Object> valueConverter;
        final Class<?> carrier;

        @SuppressWarnings("unchecked")
        <Z> Type(Class<Z> carrier, ValueLayout layout, IntFunction<Z> valueConverter) {
            this.carrier = carrier;
            this.layout = layout;
            this.valueConverter = (IntFunction<Object>)valueConverter;
        }

        long size() {
            return layout.byteSize();
        }

        VarHandle handle() {
            return layout.varHandle();
        }

        void set(MemorySegment segment, long offset, int index, int val) {
            handle().set(segment, offset + (index * size()), valueConverter.apply(val));
        }

        void check(MemorySegment segment, long offset, int index, int val) {
            assertEquals(handle().get(segment, offset + (index * size())), valueConverter.apply(val));
        }
    }

    enum SegmentKind {
        NATIVE(i -> Arena.ofAuto().allocate(i, 1)),
        ARRAY(i -> MemorySegment.ofArray(new byte[i]));

        final IntFunction<MemorySegment> segmentFactory;

        SegmentKind(IntFunction<MemorySegment> segmentFactory) {
            this.segmentFactory = segmentFactory;
        }

        MemorySegment makeSegment(int size) {
            return segmentFactory.apply(size);
        }
    }

    @DataProvider
    static Object[][] segmentKinds() {
        List<Object[]> cases = new ArrayList<>();
        for (SegmentKind kind1 : SegmentKind.values()) {
            for (SegmentKind kind2 : SegmentKind.values()) {
                cases.add(new Object[] {kind1, kind2});
            }
        }
        return cases.toArray(Object[][]::new);
    }

    @DataProvider
    static Object[][] conjunctSegments() {
        List<Object[]> cases = new ArrayList<>();
        for (SegmentKind kind : SegmentKind.values()) {
            // Different paths might be taken in the implementation depending on the
            // size, type, and address of the underlying segments.
            for (int len : new int[]{0, 1, 7, 512}) {
                for (int offset : new int[]{-1, 0, 1}) {
                    MemorySegment segment = kind.makeSegment(len + 2);
                    MemorySegment src = segment.asSlice(1 + offset, len);
                    MemorySegment dst = segment.asSlice(1, len);
                    for (int i = 0; i < len; i++) {
                        src.set(JAVA_BYTE, i, (byte) i);
                    }
                    // src = 0, 1, ... , len-1
                    cases.add(new Object[]{src, dst});
                }
            }
        }
        return cases.toArray(Object[][]::new);
    }

    @DataProvider
    static Object[][] types() {
        return Arrays.stream(Type.values())
                .map(t -> new Object[] { t })
                .toArray(Object[][]::new);
    }

    @DataProvider
    static Object[][] segmentKindsAndTypes() {
        List<Object[]> cases = new ArrayList<>();
        for (Object[] segmentKinds : segmentKinds()) {
            for (Type type1 : Type.values()) {
                for (Type type2 : Type.values()) {
                    if (type1.layout.carrier() == type2.layout.carrier()) {
                        cases.add(new Object[]{segmentKinds[0], segmentKinds[1], type1, type2});
                    }
                }
            }
        }
        return cases.toArray(Object[][]::new);
    }
}