allocation-counter.c

#include <stdlib.h>
#include <roaring/memory.h>

#if defined(__APPLE__)
#include <malloc/malloc.h>
static size_t usable_size(void *block) { return malloc_size(block); }
#else
#include <malloc.h>
static size_t usable_size(void *block) { return malloc_usable_size(block); }
#endif

static size_t allocated;

static void *counting_malloc(size_t size) {
    void *block = malloc(size);
    if (block) allocated += usable_size(block);
    return block;
}

static void *counting_realloc(void *block, size_t size) {
    size_t before = block ? usable_size(block) : 0;
    void *moved = realloc(block, size);
    if (moved) {
        allocated -= before;
        allocated += usable_size(moved);
    }
    return moved;
}

static void *counting_calloc(size_t count, size_t size) {
    void *block = calloc(count, size);
    if (block) allocated += usable_size(block);
    return block;
}

static void counting_free(void *block) {
    if (block) allocated -= usable_size(block);
    free(block);
}

static void *counting_aligned_malloc(size_t alignment, size_t size) {
    void *block = NULL;
    if (posix_memalign(&block, alignment, size) != 0) return NULL;
    allocated += usable_size(block);
    return block;
}

void croaring_ffi_count_allocations(void) {
    roaring_memory_t hooks = {
        counting_malloc, counting_realloc, counting_calloc,
        counting_free, counting_aligned_malloc, counting_free,
    };
    roaring_init_memory_hook(hooks);
}

size_t croaring_ffi_allocated_bytes(void) { return allocated; }