Bump allocator with tracked relocations in C

relocation.c

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>

// for type names, assert and alignment helpers
// https://github.com/frogtoss/ftg_toolbox_public/blob/main/ftg_core.h
#define FTG_IMPLEMENT_CORE
#include "../grain/ftg_toolbox/ftg_core.h"


// foo and bar are types that will be serialized with pointers in them
struct bar_s;
typedef struct foo_s {
    int dummy;
    struct bar_s *bar;
} foo_t;

typedef struct bar_s {
    struct foo_s *foo;
    struct bar_s *bar;
    
    int dummy;
} bar_t;

#define NULL_PADDING_BYTES 8

#define BUMP_SZ FTG_MB(4)

#define WRITE_FILENAME "out.bin"

#define MAX_RELOCATIONS ((usize)1024)
u32 g_relocation_table[MAX_RELOCATIONS];
u64 g_num_relocations = 0;


// bump is like arena, but contiguous, not paged
typedef struct {
    u8 *start;
    u8 *ptr;
    u8 *end;
} bump_t;


void *bump_alloc(bump_t *b, size_t bytes) {
    if ((usize)(b->end - b->ptr) < bytes) {    
        FTG_ASSERT_ALWAYS(!"bump oom");
    }

    // 8-byte align all bump allocations
    bytes = FTG_ALIGN_UP(bytes, 8);

    void *ptr = b->ptr;
    b->ptr += bytes;
    
    return ptr;
}

bump_t
bump_alloc_init(usize bytes) {
    bump_t b;

    // fixme: malloc does not guarantee alignment
    b.start = b.ptr = malloc(bytes);
    b.end = b.start + bytes;

    // intent is serialization, so enforce zeroing
    // of the block including padding bytes.
    //
    // this calls explicit_bzero where possible to
    // avoid optimization culling it.
    ftg_bzero(b.start, bytes);

    // handle NULL edge case: the first 8 bytes
    // are already allocated, so a relocation
    // offset of 0 is not a contradiction.
    bump_alloc(&b, NULL_PADDING_BYTES);

    return b;
}

u32
ptr_to_relocation_offset(void *base, void *ptr) {
    if (!ptr)
        return 0;
    
    // basic validity test
    FTG_ASSERT(base);
    
    // bump allocator ptrs must be 8-byte aligned
    FTG_ASSERT(FTG_IS_PTR_ALIGNED(ptr, 8));
    FTG_ASSERT(FTG_IS_PTR_ALIGNED(base, 8));

    // signed difference to catch overflow
    ptrdiff_t d = (u8*)ptr - (u8*)base;
    FTG_ASSERT(d >= 0);

    // ensure the shifted offset fits in 32GB of memory
    FTG_ASSERT(((usize)d >> 3) <= UINT32_MAX);    

    // redundant with alignment check, but placed
    // here to underscore a point:
    // 
    // first three bits of 8-byte aligned pointer is 0
    FTG_ASSERT((d & 7) == 0);

    // offsets can support up to 32GB of memory
    u32 offset = (u32)((usize)d >> 3);
    
    return offset;
}

void *
relocation_offset_to_ptr(void *base, u32 offset) {
    if (offset == 0)
        return NULL;
    
    // basic validity test
    FTG_ASSERT(base);

    uptr p = (uptr)base + ((uptr)offset << 3);
    FTG_ASSERT(p);

    return (void*)p;
}


void
bump_write_to_disk(const bump_t *bump) {

    // note: a real world serializer would take into account
    // differences that happen when moving a file from one machine to
    // another; supply a magic number, relocation table offset size
    // difference. But this is a demo.
    FILE *fh = fopen(WRITE_FILENAME, "wb");
    FTG_ASSERT_ALWAYS(fh);

    // relocation table contains offsets into the memory where
    // pointers are stored.  Convert the absolute pointers in memory
    // into offsets.
    for (u64 i = 0; i < g_num_relocations; i++) {
        u32 field_offset = g_relocation_table[i];
        void **field_ptr = (void **)(bump->start + ((uptr)field_offset << 3));

        void *abs_ptr = *field_ptr;
        u32 rel_offset = ptr_to_relocation_offset(bump->start, abs_ptr);

        // patch the pointer into a relative offset
        *field_ptr = (void*)((uptr)rel_offset);
    }
    

    u64 heap_bytes = bump->ptr - bump->start;
    usize records = 0;
    records += fwrite(&g_num_relocations, sizeof(u64), 1, fh);
    records += fwrite(&heap_bytes, sizeof(u64), 1, fh);
    records += fwrite(&g_relocation_table, sizeof(u32), g_num_relocations, fh);
    records += fwrite(bump->start, 1, heap_bytes, fh);
    FTG_ASSERT_ALWAYS(records == heap_bytes + 2 + g_num_relocations);

    fclose(fh);
    
    // This overwrite all the pointers in bump.  A production system could
    // restore them in place.
}

bump_t bump_init_and_read_from_disk(void) {
    FILE *fh = fopen(WRITE_FILENAME, "rb");
    FTG_ASSERT_ALWAYS(fh);

    usize records = 0;
    u64 heap_bytes;
    records += fread(&g_num_relocations, sizeof(u64), 1, fh);
    records += fread(&heap_bytes, sizeof(u64), 1, fh);    
    records += fread(&g_relocation_table, sizeof(u32), g_num_relocations, fh);

    // demo: trust that the heap bytes in the file is the correct
    // amount, and won't overflow or overallocate
    bump_t bump = bump_alloc_init(heap_bytes);

    records += fread(bump.start, 1, heap_bytes, fh);
    FTG_ASSERT_ALWAYS(records == heap_bytes + 2 + g_num_relocations);    

    for (u64 i = 0; i < g_num_relocations; i++) {
        u32 field_offset = g_relocation_table[i];
        void **field_ptr = (void **)(bump.start + ((uptr)field_offset << 3));

        uptr rel_ptr = (uptr)*field_ptr;
        *field_ptr = relocation_offset_to_ptr(bump.start, (u32)rel_ptr);
    }

    fclose(fh);

    return bump;
}

void *bump_get_first_alloc(const bump_t *bump) {
    return bump->start + NULL_PADDING_BYTES;
}

//
// type 'ctors' that track addresses of pointers in the struct
foo_t *
new_foo(bump_t *bump) {
    foo_t *foo = bump_alloc(bump, sizeof(foo_t));

    /* &foo->bar */
    g_relocation_table[g_num_relocations++] =
        ptr_to_relocation_offset(bump->start, (u8*)foo + offsetof(foo_t, bar));

    return foo;
}

bar_t *
new_bar(bump_t *bump) {
    bar_t *bar = bump_alloc(bump, sizeof(bar_t));


    /* &bar->foo */
    g_relocation_table[g_num_relocations++] =
        ptr_to_relocation_offset(bump->start, (u8 *)bar + offsetof(bar_t, foo));

    /* &bar->bar */
    g_relocation_table[g_num_relocations++] =
        ptr_to_relocation_offset(bump->start, (u8*)bar + offsetof(bar_t, bar));
    
    return bar;
}


int main(void) {
    bump_t bump = bump_alloc_init(BUMP_SZ);


    // create a bunch of foos and bars and connect them up
    {
        foo_t *foo0 = new_foo(&bump);
        foo0->dummy = 10;

        foo_t *foo1 = new_foo(&bump);
        foo1->dummy = 20;

        bar_t *bar0 = new_bar(&bump);
        bar0->dummy = 30;

        bar_t *bar1 = new_bar(&bump);
        bar1->dummy = 40;
        
        foo0->bar = bar0;
        foo1->bar = bar0;

        bar0->foo = NULL; /* test null edge case */
        bar0->bar = bar1; /* test self reference edge case */

        bar1->foo = foo1;
        bar1->bar = bar0;

        bump_write_to_disk(&bump);

        // clear the relocation table and free the bump allocator just so
        // we show we have no tricks up our sleeve
        ftg_bzero(&g_relocation_table, sizeof(g_relocation_table));
        free(bump.start);
    }

    // read the allocations from disk
    {
        bump_t bump = bump_init_and_read_from_disk();

        // get a pointer to the first thing that was allocated into bump
        foo_t *foo0 = bump_get_first_alloc(&bump);
        FTG_ASSERT_ALWAYS(foo0->dummy == 10);

        bar_t *bar0 = foo0->bar;
        FTG_ASSERT_ALWAYS(bar0->dummy == 30);
        FTG_ASSERT_ALWAYS(bar0->foo == NULL);

        bar_t *bar1 = bar0->bar;
        FTG_ASSERT_ALWAYS(bar1->dummy == 40);
        FTG_ASSERT_ALWAYS(bar0->bar == bar1);

        foo_t *foo1 = bar1->foo;
        FTG_ASSERT_ALWAYS(foo1->dummy == 20);
        FTG_ASSERT_ALWAYS(foo1->bar == bar0);
    }

    puts("success");
    
    return 0;
}