// Copyright 2016 The Fuchsia Authors
// Copyright (c) 2008-2012 Travis Geiselbrecht
//
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT

#include "tests.h"

#include <arch/ops.h>
#include <err.h>
#include <inttypes.h>
#include <kernel/brwlock.h>
#include <kernel/mp.h>
#include <kernel/mutex.h>
#include <kernel/spinlock.h>
#include <kernel/thread.h>
#include <ktl/type_traits.h>
#include <platform.h>
#include <rand.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <trace.h>

const size_t BUFSIZE = (3 * 1024 * 1024); // must be smaller than max allowed heap allocation
const size_t ITER = (1UL * 1024 * 1024 * 1024 / BUFSIZE); // enough iterations to have to copy/set 1GB of memory

__NO_INLINE static void bench_set_overhead() {
    uint32_t* buf = (uint32_t*)malloc(BUFSIZE);
    if (buf == nullptr) {
        TRACEF("error: malloc failed\n");
        return;
    }

    spin_lock_saved_state_t state;
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);
    uint64_t count = arch_cycle_count();
    for (size_t i = 0; i < ITER; i++) {
        __asm__ volatile("");
    }
    count = arch_cycle_count() - count;
    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    printf("took %" PRIu64 " cycles overhead to loop %zu times\n",
           count, ITER);

    free(buf);
}

__NO_INLINE static void bench_memset() {
    uint8_t* buf = (uint8_t*)memalign(PAGE_SIZE, BUFSIZE);
    if (buf == nullptr) {
        TRACEF("error: memalign failed\n");
        return;
    }

    spin_lock_saved_state_t state;
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);
    uint64_t count = arch_cycle_count();
    for (size_t i = 0; i < ITER; i++) {
        memset(buf, 0, BUFSIZE);
    }
    count = arch_cycle_count() - count;
    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    uint64_t bytes_cycle = (BUFSIZE * ITER * 1000ULL) / count;
    printf("took %" PRIu64 " cycles to memset a buffer of size %zu %zu times "
           "(%" PRIu64 " bytes), %" PRIu64 ".%03" PRIu64 " bytes/cycle\n",
           count, BUFSIZE, ITER, BUFSIZE * ITER, bytes_cycle / 1000, bytes_cycle % 1000);

    free(buf);
}

__NO_INLINE static void bench_memset_per_page() {
    uint8_t* buf = (uint8_t*)memalign(PAGE_SIZE, BUFSIZE);
        if (buf == nullptr) {
        TRACEF("error: memalign failed\n");
        return;
    }

    spin_lock_saved_state_t state;
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);
    uint64_t count = arch_cycle_count();
    for (size_t i = 0; i < ITER; i++) {
        for (size_t j = 0; j < BUFSIZE; j += PAGE_SIZE) {
            memset(buf + j, 0, PAGE_SIZE);
        }
    }
    count = arch_cycle_count() - count;
    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    uint64_t bytes_cycle = (BUFSIZE * ITER * 1000ULL) / count;
    printf("took %" PRIu64 " cycles to per-page memset a buffer of size %zu %zu times "
           "(%" PRIu64 " bytes), %" PRIu64 ".%03" PRIu64 " bytes/cycle\n",
           count, BUFSIZE, ITER, BUFSIZE * ITER, bytes_cycle / 1000, bytes_cycle % 1000);

    free(buf);
}

__NO_INLINE static void bench_zero_page() {
    uint8_t* buf = (uint8_t*)memalign(PAGE_SIZE, BUFSIZE);
        if (buf == nullptr) {
        TRACEF("error: memalign failed\n");
        return;
    }

    spin_lock_saved_state_t state;
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);
    uint64_t count = arch_cycle_count();
    for (size_t i = 0; i < ITER; i++) {
        for (size_t j = 0; j < BUFSIZE; j += PAGE_SIZE) {
            arch_zero_page(buf + j);
        }
    }
    count = arch_cycle_count() - count;
    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    uint64_t bytes_cycle = (BUFSIZE * ITER * 1000ULL) / count;
    printf("took %" PRIu64 " cycles to arch_zero_page a buffer of size %zu %zu times "
           "(%" PRIu64 " bytes), %" PRIu64 ".%03" PRIu64 " bytes/cycle\n",
           count, BUFSIZE, ITER, BUFSIZE * ITER, bytes_cycle / 1000, bytes_cycle % 1000);

    free(buf);
}

template <typename T>
__NO_INLINE static void bench_cset() {
    T* buf = (T*)malloc(BUFSIZE);
    if (buf == nullptr) {
        TRACEF("error: malloc failed\n");
        return;
    }

    spin_lock_saved_state_t state;
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);
    uint64_t count = arch_cycle_count();
    for (size_t i = 0; i < ITER; i++) {
        for (size_t j = 0; j < BUFSIZE / sizeof(T); j++) {
            buf[j] = 0;
        }
    }
    count = arch_cycle_count() - count;
    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    uint64_t bytes_cycle = (BUFSIZE * ITER * 1000ULL) / count;
    printf("took %" PRIu64 " cycles to clear a buffer using wordsize %zu of size %zu %zu times "
           "(%" PRIu64 " bytes), %" PRIu64 ".%03" PRIu64 " bytes/cycle\n",
           count, sizeof(*buf), BUFSIZE, ITER, BUFSIZE * ITER, bytes_cycle / 1000, bytes_cycle % 1000);

    free(buf);
}

__NO_INLINE static void bench_cset_wide() {
    uint32_t* buf = (uint32_t*)malloc(BUFSIZE);
    if (buf == nullptr) {
        TRACEF("error: malloc failed\n");
        return;
    }

    spin_lock_saved_state_t state;
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);
    uint64_t count = arch_cycle_count();
    for (size_t i = 0; i < ITER; i++) {
        for (size_t j = 0; j < BUFSIZE / sizeof(*buf) / 8; j++) {
            buf[j * 8] = 0;
            buf[j * 8 + 1] = 0;
            buf[j * 8 + 2] = 0;
            buf[j * 8 + 3] = 0;
            buf[j * 8 + 4] = 0;
            buf[j * 8 + 5] = 0;
            buf[j * 8 + 6] = 0;
            buf[j * 8 + 7] = 0;
        }
    }
    count = arch_cycle_count() - count;
    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    uint64_t bytes_cycle = (BUFSIZE * ITER * 1000ULL) / count;
    printf("took %" PRIu64 " cycles to clear a buffer of size %zu %zu times 8 words at a time "
           "(%" PRIu64 " bytes), %" PRIu64 ".%03" PRIu64 " bytes/cycle\n",
           count, BUFSIZE, ITER, BUFSIZE * ITER, bytes_cycle / 1000, bytes_cycle % 1000);

    free(buf);
}

__NO_INLINE static void bench_memcpy() {
    uint8_t* buf = (uint8_t*)calloc(1, BUFSIZE);
    if (buf == nullptr) {
        TRACEF("error: calloc failed\n");
        return;
    }

    spin_lock_saved_state_t state;
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);
    uint64_t count = arch_cycle_count();
    for (size_t i = 0; i < ITER; i++) {
        memcpy(buf, buf + BUFSIZE / 2, BUFSIZE / 2);
    }
    count = arch_cycle_count() - count;
    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    uint64_t bytes_cycle = (BUFSIZE / 2 * ITER * 1000ULL) / count;
    printf("took %" PRIu64 " cycles to memcpy a buffer of size %zu %zu times "
           "(%zu source bytes), %" PRIu64 ".%03" PRIu64 " source bytes/cycle\n",
           count, BUFSIZE / 2, ITER, BUFSIZE / 2 * ITER, bytes_cycle / 1000, bytes_cycle % 1000);

    free(buf);
}

__NO_INLINE static void bench_spinlock() {
    spin_lock_saved_state_t state;
    spin_lock_saved_state_t state2;
    spin_lock_t lock;

    spin_lock_init(&lock);

#define COUNT (128 * 1024 * 1024)
    // test 1: acquire/release a spinlock with interrupts already disabled
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    uint64_t c = arch_cycle_count();
    for (size_t i = 0; i < COUNT; i++) {
        spin_lock(&lock);
        spin_unlock(&lock);
    }
    c = arch_cycle_count() - c;

    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    printf("%" PRIu64 " cycles to acquire/release spinlock %d times (%" PRIu64 " cycles per)\n", c, COUNT, c / COUNT);

    // test 2: acquire/release a spinlock with irq save and irqs already disabled
    arch_interrupt_save(&state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    c = arch_cycle_count();
    for (size_t i = 0; i < COUNT; i++) {
        spin_lock_irqsave(&lock, state2);
        spin_unlock_irqrestore(&lock, state2);
    }
    c = arch_cycle_count() - c;

    arch_interrupt_restore(state, ARCH_DEFAULT_SPIN_LOCK_FLAG_INTERRUPTS);

    printf("%" PRIu64 " cycles to acquire/release spinlock w/irqsave (already disabled) %d times (%" PRIu64 " cycles per)\n", c, COUNT, c / COUNT);

    // test 2: acquire/release a spinlock with irq save and irqs enabled
    c = arch_cycle_count();
    for (size_t i = 0; i < COUNT; i++) {
        spin_lock_irqsave(&lock, state2);
        spin_unlock_irqrestore(&lock, state2);
    }
    c = arch_cycle_count() - c;

    printf("%" PRIu64 " cycles to acquire/release spinlock w/irqsave %d times (%" PRIu64 " cycles per)\n", c, COUNT, c / COUNT);
#undef COUNT
}

__NO_INLINE static void bench_mutex() {
    Mutex m;

    static const uint count = 128 * 1024 * 1024;
    uint64_t c = arch_cycle_count();
    for (size_t i = 0; i < count; i++) {
        m.Acquire();
        m.Release();
    }
    c = arch_cycle_count() - c;

    printf("%" PRIu64 " cycles to acquire/release uncontended mutex %u times (%" PRIu64 " cycles per)\n", c, count, c / count);
}

template<typename LockType>
__NO_INLINE static void bench_rwlock() {
    LockType rw;
    static const uint count = 128 * 1024 * 1024;
    uint64_t c = arch_cycle_count();
    for (size_t i = 0; i < count; i++) {
        rw.ReadAcquire();
        rw.ReadRelease();
    }
    c = arch_cycle_count() - c;

    printf("%" PRIu64
           " cycles to acquire/release uncontended brwlock(PI: %d) for read %u times (%" PRIu64
           " cycles per)\n",
           c, ktl::is_same_v<LockType, BrwLockPi>, count, c / count);

    c = arch_cycle_count();
    for (size_t i = 0; i < count; i++) {
        rw.WriteAcquire();
        rw.WriteRelease();
    }
    c = arch_cycle_count() - c;

    printf("%" PRIu64
           " cycles to acquire/release uncontended brwlock(PI: %d) for write %u times (%" PRIu64
           " cycles per)\n",
           c, ktl::is_same_v<LockType, BrwLockPi>, count, c / count);
}

int benchmarks(int, const cmd_args*, uint32_t) {
    bench_set_overhead();
    bench_memcpy();
    bench_memset();

    bench_memset_per_page();
    bench_zero_page();

    bench_cset<uint8_t>();
    bench_cset<uint16_t>();
    bench_cset<uint32_t>();
    bench_cset<uint64_t>();
    bench_cset_wide();

    bench_spinlock();
    bench_mutex();
    bench_rwlock<BrwLockPi>();
    bench_rwlock<BrwLockNoPi>();

    return 0;
}