commit b60bb52e780127814751964c5e90c7631ff1a5e4
Author: Mitja Horvat <mitja@plume.com>
Date:   Fri Nov 8 08:49:22 2024 +0100

    Arena initial implementation.
    
    Signed-off-by: Mitja Horvat <mitja@plume.com>

diff --git a/src/arena.c b/src/arena.c
new file mode 100644
index 0000000..f5ac3fe
--- /dev/null
+++ b/src/arena.c
@@ -0,0 +1,209 @@
+#include <sys/mman.h>
+
+#include <assert.h>
+#include <errno.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+#include <pthread.h>
+
+#include "arena.h"
+
+#if __STDC_VERSION__ < 201112L /* below C11 */
+typedef intptr_t max_align_t;
+#define thread_local __thread
+#endif
+
+static void arena_defer_cleanup(arena_t *arena);
+
+arena_t *arena_new(size_t sz)
+{
+    arena_t *arena;
+    size_t msz = sz + sizeof(arena_t);
+
+    if (msz < sysconf(_SC_PAGE_SIZE))
+    {
+        arena = malloc(msz);
+    }
+    else
+    {
+        arena = mmap(NULL, msz, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0);
+        if (arena == MAP_FAILED) return false;
+    }
+
+    arena->a_sz = sz;
+    arena->a_pos = 0;
+
+    return arena;
+}
+
+void arena_del(arena_t *arena)
+{
+    size_t msz = arena->a_sz + sizeof(arena_t);
+
+    arena_pull(arena, SIZE_MAX);
+
+    memset(arena, 0, sizeof(*arena));
+
+    if (msz < sysconf(_SC_PAGE_SIZE))
+    {
+        free(arena);
+    }
+    else if (munmap(arena, msz) != 0)
+    {
+        fprintf(stderr, "munmap() failed: %s\n", strerror(errno));
+    }
+}
+
+/*
+ * Extend the arena boundary by `sz` bytes.
+ * Can be used to grow the last allocated element.
+ */
+void *arena_push(arena_t *arena, size_t sz)
+{
+    void *ret = &arena->a_data[arena->a_pos];
+
+    if ((SIZE_MAX - arena->a_pos) < sz ||
+            arena->a_pos + sz > arena->a_sz)
+    {
+        return NULL;
+    }
+
+    arena->a_pos += sz;
+    return ret;
+}
+
+/*
+ * Reduce the arena boundary by `sz` bytes.
+ */
+void arena_pull(arena_t *arena, size_t sz)
+{
+    if (arena->a_pos < sz)
+    {
+        arena->a_pos = 0;
+    }
+    else
+    {
+        arena->a_pos -= sz;
+    }
+    arena_defer_cleanup(arena);
+}
+
+/*
+ * Save current arena boundary, calling arena_restore() afterwards will
+ * free everything that was allocated between arena_save() and arena_restore()
+ */
+arena_frame_t arena_save(arena_t *arena)
+{
+    if (arena == NULL) return (arena_frame_t){ .af_arena = NULL };
+    return (arena_frame_t){ .af_arena = arena, .af_pos = arena->a_pos };
+}
+
+/*
+ * Restore saved arena boundary (free elements between arena_save() and arena_restore())
+ */
+void arena_restore(arena_frame_t *frame)
+{
+    if (frame->af_arena == NULL) return;
+    if (frame->af_pos > frame->af_arena->a_pos) return;
+    arena_pull(frame->af_arena, frame->af_arena->a_pos - frame->af_pos);
+}
+
+/*
+ * Allocate a defer action on the current arena. If the defer action block
+ * is freed, the defer callback is called.
+ */
+bool arena_defer(arena_t *arena, arena_defer_fn_t *fn, void *data)
+{
+    struct arena_defer *defer;
+
+    defer = arena_malloc(arena, sizeof(*defer));
+    if (defer == NULL)
+    {
+        fprintf(stderr, "Unable to allocate defer buffer.");
+        fn(arena, data);
+        return false;
+    }
+
+    defer->ad_fn = fn;
+    defer->ad_data = data;
+
+    defer->ad_next = arena->a_defer;
+    defer->ad_magic = ARENA_MAGIC;
+    defer->ad_magic ^= (uintptr_t)defer;
+    defer->ad_magic ^= (uintptr_t)defer->ad_next;
+    defer->ad_magic ^= (uintptr_t)defer->ad_fn;
+    defer->ad_magic ^= (uintptr_t)defer->ad_data;
+    arena->a_defer = defer;
+
+    return true;
+}
+
+void arena_defer_cleanup(arena_t *arena)
+{
+    while (arena->a_defer != NULL)
+    {
+        if (((uint8_t *)arena->a_defer + sizeof(struct arena_defer)) <= &arena->a_data[arena->a_pos])
+            break;
+
+        uintptr_t magic = arena->a_defer->ad_magic;
+        magic ^= (uintptr_t)arena->a_defer;
+        magic ^= (uintptr_t)arena->a_defer->ad_next;
+        magic ^= (uintptr_t)arena->a_defer->ad_fn;
+        magic ^= (uintptr_t)arena->a_defer->ad_data;
+        if (magic != ARENA_MAGIC)
+        {
+            assert(!"Defer buffer corrupted.");
+        }
+
+        arena->a_defer->ad_fn(arena, arena->a_defer->ad_data);
+        arena->a_defer = arena->a_defer->ad_next;
+    }
+}
+
+
+void *arena_malign(arena_t *arena, size_t sz)
+{
+    size_t off;
+
+    size_t align = 1;
+
+    /* Calculate the alignment of sz */
+    while (align < sz)
+    {
+        align <<= 1;
+        if (align >= sizeof(max_align_t)) break;
+    }
+
+    /* Calculate the offset we need to add to the current position */
+    off = align;
+    off -= arena->a_pos & (align - 1);
+    off &= (align - 1);
+
+    if (arena_push(arena, off) == NULL) return NULL;
+
+    return &arena->a_data[arena->a_pos];
+}
+
+void *arena_malloc(arena_t *arena, size_t sz)
+{
+    if (arena_malign(arena, sz) == NULL) return NULL;
+    return arena_push(arena, sz);
+}
+
+/*
+ * Allocate a string on the arena
+ */
+char *arena_strdup(arena_t *arena, const char *src)
+{
+    char *dst;
+
+    size_t slen = strlen(src);
+
+    dst = arena_push(arena, slen);
+    if (dst == NULL) return NULL;
+
+    return strcpy(dst, src);
+}
+
diff --git a/src/arena.h b/src/arena.h
new file mode 100644
index 0000000..e2e2cc1
--- /dev/null
+++ b/src/arena.h
@@ -0,0 +1,63 @@
+#if !defined(ARENA_H_INCLUDED)
+#define ARENA_H_INCLUDED
+
+#include <stddef.h>
+#include <stdint.h>
+#include <stdbool.h>
+
+#define ARENA_MAGIC (0xAA55AA55AA55AA55)
+
+#define __V(x, ...) __VA_ARGS__
+#define ARENA_VPACK(...) ((arena_t *[]){ __V(dummy, ##__VA_ARGS__, NULL) })
+
+#if defined(__GNUC__)
+#define CLEANUP(...) __attribute__((cleanup(__VA_ARGS__)))
+#endif
+
+/*
+ * Automatically free arena when ARENA_AUTO goes out of scope, this is equivalent
+ * of calling arena_save() at the beginning of the current scope and
+ * arena_restore() at the end.
+ *
+ * This uses a GNU C specific extension.
+ */
+#define ARENA_AUTO(x) \
+        arena_frame_t CLEANUP(arena_restore) __cleanup_ ## x  = arena_save(x);
+
+typedef struct arena arena_t;
+typedef struct arena_frame arena_frame_t;
+typedef void arena_defer_fn_t(arena_t *, void *data);
+
+struct arena
+{
+    size_t              a_pos;              /* Current arena position */
+    size_t              a_sz;               /* Total arena size */
+    struct arena_defer *a_defer;            /* List of deferes */
+    uint8_t             a_data[0];          /* Arena data */
+};
+
+struct arena_frame
+{
+    arena_t            *af_arena;
+    size_t              af_pos;
+};
+
+struct arena_defer
+{
+    uintptr_t            ad_magic;
+    arena_defer_fn_t    *ad_fn;
+    void                *ad_data;
+    void                *ad_next;
+};
+
+arena_t *arena_new(size_t sz);
+void arena_del(arena_t *arena);
+void *arena_push(arena_t *arena, size_t sz);
+void arena_pull(arena_t *arena, size_t sz);
+arena_frame_t arena_save(arena_t *arena);
+void arena_restore(arena_frame_t *frame);
+bool arena_defer(arena_t *arena, arena_defer_fn_t *fn, void *data);
+void *arena_malloc(arena_t *arena, size_t sz);
+char *arena_strdup(arena_t *arena, const char *src);
+
+#endif /* ARENA_H_INCLUDED */
diff --git a/src/arena_thread.h b/src/arena_thread.h
new file mode 100644
index 0000000..ad54b2c
--- /dev/null
+++ b/src/arena_thread.h
@@ -0,0 +1,27 @@
+#include <pthread.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "scratch.h"
+
+void scratch_thread_cleanup(void *p)
+{
+    (void)p;
+    scratch_destroy();
+}
+
+pthread_key_t scratch_thread_key;
+void scratch_cleanup_once(void)
+{
+    printf("Multi-threaded mode.\n");
+    pthread_key_create(&scratch_thread_key, scratch_thread_cleanup);
+    atexit(scratch_destroy);
+}
+
+void scratch_cleanup_init(void)
+{
+    static pthread_once_t once = PTHREAD_ONCE_INIT;
+    (void)pthread_once(&once, scratch_cleanup_once);
+    pthread_setspecific(scratch_thread_key, (void *)0xdeadbeef);
+}
+
diff --git a/src/main.c b/src/main.c
new file mode 100644
index 0000000..dc57f91
--- /dev/null
+++ b/src/main.c
@@ -0,0 +1,130 @@
+#include <dlfcn.h>
+#include <stdio.h>
+#include <unistd.h>
+#include <pthread.h>
+
+#include "arena.h"
+#include "arena_thread.h"
+#include "scratch.h"
+
+#define TEST_FILE "/etc/os-release"
+
+void foo(arena_t *ar, const char *str)
+{
+    SCRATCH_GET(arena, ar);
+    arena_push(arena, 11);
+    printf("Got: %s\n", str);
+}
+
+void test(void)
+{
+    SCRATCH_GET(arena);
+    foo(arena, arena_strdup(arena, "hello world"));
+}
+
+void arena_defer_test(arena_t *arena, void *data)
+{
+    (void)data;
+    (void)arena;
+
+    printf("Cleanup!\n");
+}
+
+void arena_defer_close_fn(arena_t *arena, void *data)
+{
+    int fd = (intptr_t)data;
+    printf("Closing file descriptor: %d\n", fd);
+    close(fd);
+}
+
+bool arena_defer_close(arena_t *arena, int fd)
+{
+    return arena_defer(arena, arena_defer_close_fn, (void *)(uintptr_t)fd);
+}
+
+void arena_defer_fclose_fn(arena_t *arena, void *data)
+{
+    printf("Closing file: %p\n", data);
+    fclose((FILE *)data);
+}
+
+bool arena_defer_fclose(arena_t *arena, FILE *f)
+{
+    return arena_defer(arena, arena_defer_fclose_fn, f);
+}
+
+char *get_file(arena_t *arena, const char *path)
+{
+    FILE *f;
+    char *rc;
+    size_t nr;
+    char *buf;
+
+    if ((f = fopen(path, "r")) == NULL)
+    {
+        return NULL;
+    }
+
+    if (!arena_defer_fclose(arena, f))
+    {
+        return NULL;
+    }
+
+    buf = rc = arena_push(arena, 16);
+    while ((nr = fread(buf, 1, 16, f)) >= 16)
+    {
+        buf = arena_push(arena, 16);
+    }
+    buf[nr] = '\0';
+
+    return rc;
+}
+
+void cleanup(void *p)
+{
+    (void)p;
+    printf("THREAD CLEANUP!\n");
+}
+
+pthread_key_t key;
+
+void thr_defer_test(arena_t *a, void *t)
+{
+    (void)t;
+    printf("DEFER WORKS\n");
+}
+
+void *thr_main(void *p)
+{
+    arena_t *sc = scratch_get();
+    arena_defer(sc, thr_defer_test, NULL);
+    printf("> Going out of thread\n");
+    return NULL;
+}
+
+int main(void)
+{
+    char *file;
+    SCRATCH_GET(scratch);
+
+    file = get_file(scratch, TEST_FILE);
+    if (file == NULL)
+    {
+        printf("Error reading file: %s\n", TEST_FILE);
+    }
+
+    printf("FILE: %s\n", file);
+
+    printf("NEW thread\n");
+    pthread_t thr;
+    pthread_create(&thr, NULL, thr_main, NULL);
+    printf("Joining thread ...\n");
+    pthread_join(thr, NULL);
+    printf("THREAD EXIT!\n");
+
+    arena_t *sc = scratch_get(scratch);
+    arena_defer(sc, thr_defer_test, NULL);
+
+    return 0;
+}
+
diff --git a/src/scratch.c b/src/scratch.c
new file mode 100644
index 0000000..319bdc6
--- /dev/null
+++ b/src/scratch.c
@@ -0,0 +1,81 @@
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "scratch.h"
+
+#if __STDC_VERSION__ >= 201112L /* C11 and above */
+#define thread_local _Thread_local
+#elif defined(__GNUC__)
+#define thread_local __thread
+#else
+#error Unsupported compiler.
+#endif
+
+#if defined(__GNUC__)
+#define WEAK __attribute__((weak))
+#endif
+
+thread_local arena_t *scratch_arena[4] = { 0 };
+
+void scratch_destroy(void)
+{
+    for (int s = 0; s < sizeof(scratch_arena) / sizeof(scratch_arena[0]); s++)
+    {
+        if (scratch_arena[s] == NULL) continue;
+        arena_del(scratch_arena[s]);
+        scratch_arena[s] = NULL;
+    }
+}
+
+/*
+ * Optionally allocate and return a scratch arena. Make sure that the returned
+ * arena is not present in the NULL-terminated list `conflicts`
+ */
+arena_t *__scratch_get(arena_t **conflicts)
+{
+    int s = 0;
+
+    if (conflicts != NULL)
+    {
+        for (s = 0; s < sizeof(scratch_arena) / sizeof(scratch_arena[0]); s++)
+        {
+            arena_t **c = conflicts;
+            for (; *c != NULL; c++)
+            {
+                if (*c == scratch_arena[s]) break;
+            }
+            /* No conflcits found, break out */
+            if (*c == NULL) break;
+        }
+        if (s >= sizeof(scratch_arena) / sizeof(scratch_arena[0]))
+        {
+            fprintf(stderr, "Out of scratch arenas.\n");
+            return NULL;
+        }
+    }
+
+    if (scratch_arena[s] == NULL)
+    {
+        scratch_cleanup_init();
+        scratch_arena[s] = arena_new(SCRATCH_DEFAULT_SIZE);
+        if (scratch_arena[s] == NULL)
+        {
+            fprintf(stderr, "Error allocating scratch arena %d.", s);
+            return NULL;
+        }
+
+        printf("NEW SCRATCH = %p\n", scratch_arena[s]);
+    }
+
+    return scratch_arena[s];
+}
+
+void WEAK scratch_cleanup_init(void)
+{
+    static bool init = false;
+    for (;!init; init = true)
+    {
+        printf("Single-threaded mode.\n");
+        atexit(scratch_destroy);
+    }
+}
diff --git a/src/scratch.h b/src/scratch.h
new file mode 100644
index 0000000..8d9a2e9
--- /dev/null
+++ b/src/scratch.h
@@ -0,0 +1,25 @@
+#if !defined(SCRATCH_H_INCLUDED)
+#define SCRATCH_H_INCLUDED
+
+#include "arena.h"
+
+/* Default scratch arena size */
+#define SCRATCH_DEFAULT_SIZE (2*1024*1024)
+
+/*
+ * Get a scratch arena and automatically free it. `...` in this context is an
+ * array of conflicting arenas. scratch_get() ensures that `x` is not a
+ * previously allocated scratchpad arena that is present in the conflict list
+ */
+#define SCRATCH_GET(x, ...) \
+        arena_t *x = scratch_get(__VA_ARGS__); \
+        ARENA_AUTO(x)
+
+/* Wrapper around __scratch_get() */
+#define scratch_get(...)    __scratch_get(ARENA_VPACK(__VA_ARGS__))
+
+arena_t *__scratch_get(arena_t **conflicts);
+void scratch_destroy(void);
+void scratch_cleanup_init(void);
+
+#endif /* SCRATCH_H_INCLUDED */