#include <cassert>
#include <jni.h>
#include <string>

#include "PoolAllocator.h"

void testPoolAllocator();

extern "C" JNIEXPORT jstring JNICALL
Java_com_example_poolallocator_MainActivity_testPoolAllocator(
	JNIEnv* env,
	jobject thiz) {
    testPoolAllocator();
    return env->NewStringUTF("Pool Allocator test finished");
}

struct Particle {
    int x;
    int y;
};


// Allocate and free in the same order.
void testFreeInOrder() {
    PoolAllocator<Particle> allocator(100);
    allocator.Init();
   
    Particle* ptr1 = allocator.Allocate();
    ptr1->x = ptr1->y = 1;
    Particle* ptr2 = allocator.Allocate();
    ptr2->x = ptr2->y = 2;
    allocator.Free(ptr1);
    allocator.Free(ptr2);
}

// Free allocated objects in reverse order.
void testFreeReverseOrder() {
    PoolAllocator<Particle> allocator(100);
    allocator.Init();

    // Allocate and free in reverse order.
    Particle* ptr1 = allocator.Allocate();
    Particle* ptr2 = allocator.Allocate();
    allocator.Free(ptr2);
    allocator.Free(ptr1);

    // Verify that the next two allocs still work.
    Particle* ptr3 = allocator.Allocate();
    Particle* ptr4 = allocator.Allocate();
    allocator.Free(ptr3);
    allocator.Free(ptr4);
}

// Allocate and free calls are interleaved.
void testFreeInterleaved() {
    PoolAllocator<Particle> allocator(100);
    allocator.Init();

    Particle* ptr1 = allocator.Allocate();

    Particle* ptr2 = allocator.Allocate();
    allocator.Free(ptr1);
    Particle* ptr3 = allocator.Allocate();
    allocator.Free(ptr2);
    Particle* ptr4 = allocator.Allocate();
    allocator.Free(ptr3);

    allocator.Free(ptr4);
}

void testAllocateUntilFull() {
    PoolAllocator<Particle> allocator(100);
    allocator.Init();

    for (int i = 0; i < allocator.GetCapacity(); i++) {
        Particle* ptr1 = allocator.Allocate();
        // do not free ptr1 here.
        ptr1->x = 1;
        ptr1->y = 2;
    }

    // The next allocation should be nullptr.
    Particle* ptr2 = allocator.Allocate();
    assert(ptr2 == nullptr);
}

void testInitNotCalled() {
    PoolAllocator<Particle> allocator(100);
    // Nothing else. Just the destructor.
}

// Modifying an object that was already freed.
void testUseAfterFree() {
    PoolAllocator<Particle> allocator(100);
    allocator.Init();

    Particle* ptr1 = allocator.Allocate();
    allocator.Free(ptr1);

    // Illegal access.
    ptr1->x = 1;
    ptr1->y = 2;
}

// Calling Free on an already freed object.
void testDoubleFree() {
    PoolAllocator<Particle> allocator(100);
    allocator.Init();

    Particle* ptr1 = allocator.Allocate();
    ptr1->x = 1;
    ptr1->y = 2;
    allocator.Free(ptr1);

    allocator.Free(ptr1);
}


void testPoolAllocator() {
    // These are valid operations. HWAsan should not crash here.
    testFreeInOrder();
    testFreeReverseOrder();
    testFreeInterleaved();
    testAllocateUntilFull();
    testInitNotCalled();

    // These methods trigger invalid memory accesses that HWAsan catches.
    // Without HWAsan, they can be difficult to debug.
    //testUseAfterFree();
    //testDoubleFree();
}