package jdk.test;

import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Group;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Warmup;
import org.openjdk.jmh.infra.Blackhole;

import java.lang.invoke.MethodHandles;
import java.lang.invoke.VarHandle;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.concurrent.TimeUnit;
import java.util.function.IntUnaryOperator;

@BenchmarkMode(Mode.AverageTime)
@Fork(value = 1, warmups = 0)
@Warmup(iterations = 10)
@Measurement(iterations = 20)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
public class GetAndUpdateBench2 {

    // allocate a direct buffer that starts at the start of cache line
    // so that INT_VIEW's indices 0 and 4 address elements in the same cache-line
    static final ByteBuffer BB = ByteBuffer.allocateDirect(128)
                                           .alignedSlice(64)
                                           .order(ByteOrder.nativeOrder());

    static final VarHandle INT_VIEW =
        MethodHandles.byteBufferViewVarHandle(int[].class, ByteOrder.nativeOrder());

    static final IntUnaryOperator UPDATEFN = i -> {
        Blackhole.consumeCPU(20L);
        return i + 1;
    };

    @Benchmark
    @Group("dflt")
    public int getAndUpdate1_dflt() {
        return getAndUpdate_dflt(0);
    }

    @Benchmark
    @Group("dflt")
    public int getAndUpdate2_dflt() {
        return getAndUpdate_dflt(Integer.BYTES);
    }

    private int getAndUpdate_dflt(int i) {
        int prev = (int) INT_VIEW.get(BB, i), next = 0;
        for (boolean haveNext = false; ; ) {
            if (!haveNext)
                next = UPDATEFN.applyAsInt(prev);
            if (INT_VIEW.weakCompareAndSet(BB, i, prev, next))
                return prev;
            haveNext = (prev == (prev = (int) INT_VIEW.get(BB, i)));
        }
    }

    @Benchmark
    @Group("martin")
    public int getAndUpdate1_martin() {
        return getAndUpdate_martin(0);
    }

    @Benchmark
    @Group("martin")
    public int getAndUpdate2_martin() {
        return getAndUpdate_martin(Integer.BYTES);
    }

    private int getAndUpdate_martin(int i) {
        for (int prev = (int) INT_VIEW.get(BB, i); ; ) {
            int next = UPDATEFN.applyAsInt(prev);
            do {
                if (INT_VIEW.weakCompareAndSet(BB, i, prev, next))
                    return prev;
            } while (prev == (prev = (int) INT_VIEW.get(BB, i)));
        }
    }

    @Benchmark
    @Group("shade")
    public int getAndUpdate1_shade() {
        return getAndUpdate_shade(0);
    }

    @Benchmark
    @Group("shade")
    public int getAndUpdate2_shade() {
        return getAndUpdate_shade(Integer.BYTES);
    }

    private int getAndUpdate_shade(int i) {
        int prev, next;
        do {
            prev = (int) INT_VIEW.get(BB, i);
            next = UPDATEFN.applyAsInt(prev);
        } while (!INT_VIEW.weakCompareAndSet(BB, i, prev, next));
        return prev;
    }

    @Benchmark
    @Group("strong")
    public int getAndUpdate1_strong() {
        return getAndUpdate_strong(0);
    }

    @Benchmark
    @Group("strong")
    public int getAndUpdate2_strong() {
        return getAndUpdate_strong(Integer.BYTES);
    }

    private int getAndUpdate_strong(int i) {
        int prev, next;
        do {
            prev = (int) INT_VIEW.get(BB, i);
            next = UPDATEFN.applyAsInt(prev);
        } while (!INT_VIEW.compareAndSet(BB, i, prev, next));
        return prev;
    }
}