#include <cheriintrin.h>
#include <stdint.h>
#include <stddef.h>

#define HEAP_SIZE 4096

static char heap[HEAP_SIZE];
static size_t bump = 0;

// Add delta value for TLB translation
static inline void * __capability add_delta(void * __capability cap, int offset) {
    void * __capability result;

    asm volatile (
        "candperm %0, %1, %2"
        : "=C" (result)  // Output: %0 (result)
        : "C" (cap),     // Input:  %1 (original cap)
          "r" (offset)   // Input:  %2 (offset register)
        :                // No clobbered registers
    );

    return result;
}


void * __capability malloc(size_t size) {

    // Align to 8 bytes (important for capability safety)
    size = (size + 7) & ~7;

    if (bump + size > HEAP_SIZE)
        return NULL;

    void * __capability base = cheri_ddc_get();

    uintptr_t addr = (uintptr_t)(heap + bump);

    // Create capability to this region
    void * __capability cap = cheri_address_set(base, addr);

    // Enforce bounds (this is the key CHERI feature)
    cap = cheri_bounds_set(cap, size);

    // Hard-coded delta value 
    cap = add_delta(cap, 4);

    bump += size;

    return cap;
}

void bump_reset(void) {
    bump = 0;
}

int main(void) {
    char * __capability a = malloc(30);
    char * __capability b = malloc(16);

    if (!a || !b) return -1;

    a[0] = 'A';
    a[1] = 'C';
    b[0] = 'B';

    // This will fault (out-of-bounds)
    // a[20] = 'X';

    if (a[1] != 'C') {
        while (1);
    }

    return 0;
}