package com.arm.benchmarks.vectorapi;

import java.util.Arrays;
import java.util.concurrent.TimeUnit;

import jdk.incubator.vector.IntVector;
import jdk.incubator.vector.Vector;
import jdk.incubator.vector.VectorMask;
import jdk.incubator.vector.VectorSpecies;

import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;

@BenchmarkMode(Mode.Throughput)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@State(Scope.Benchmark)
@Warmup(iterations = 3, time = 10)
@Measurement(iterations = 5, time = 10)
@Fork(value = 1, jvmArgsPrepend = {"--add-modules=jdk.incubator.vector"})
public class IntVectorJmhTest {

    @Param({"1000", "10000", "100000"})
    public int fastPathCount;

    @Param({"false", "true"})
    public boolean hasSlowPath;

    private static final VectorSpecies<Integer> SPECIES =
        IntVector.SPECIES_PREFERRED;

    private int slowPathCount;
    private int arrayLength;
    private int[] array;
    private int[] result;
    private VectorMask<Integer> mask;

    /**
     * Setup before benchmark.
     */
    @Setup
    public void init() {
        slowPathCount = hasSlowPath ? 1 : 0;
        arrayLength = SPECIES.length() * fastPathCount + slowPathCount;

        boolean[] m = new boolean[SPECIES.length()];
        Arrays.fill(m, 0, 1, true);
        mask = VectorMask.fromValues(SPECIES, m);

        array = new int[arrayLength];
        for (int i = 0; i < arrayLength; i++) {
            array[i] = i;
        }

        result = new int[arrayLength];
    }

    /**
     * Test against fromArray/intoArray API.
     */
    @Benchmark
    public void fromIntoArray() {
        for (int i = 0; i < array.length; i += SPECIES.length()) {
            IntVector av = IntVector.fromArray(SPECIES, array, i, mask);
            av.intoArray(result, i, mask);
        }
    }
}