4 files changed, 192 insertions, 21 deletions

diff --git a/Justfile b/Justfile
index 8343372..e1f0346 100644
--- a/Justfile
+++ b/Justfile
@@ -50,6 +50,14 @@ test-all:
     {{cc}} {{cflags}} {{test_dir}}/*.c {{src_dir}}/*.c -o {{build_dir}}/test -lcriterion {{lflags}}
     ./{{build_dir}}/test
 
+coverage:
+    mkdir -p {{build_dir}}/cov-src {{build_dir}}/cov-tests
+    cd {{build_dir}}/cov-src && {{cc}} {{cflags}} --coverage -c ../../{{src_dir}}/*.c
+    cd {{build_dir}}/cov-tests && {{cc}} {{cflags}} --coverage -c ../../{{test_dir}}/*.c
+    {{cc}} {{cflags}} --coverage {{build_dir}}/cov-src/*.o {{build_dir}}/cov-tests/*.o -o {{build_dir}}/test-coverage -lcriterion {{lflags}}
+    ./{{build_dir}}/test-coverage
+    gcovr --root . --filter '{{src_dir}}/' {{build_dir}}/cov-src --html-details {{build_dir}}/coverage.html --print-summary
+
 clean:
     rm -rf {{build_dir}}
 
diff --git a/src/smrt_arena.c b/src/smrt_arena.c
index 0fb8ab7..6862212 100644
--- a/src/smrt_arena.c
+++ b/src/smrt_arena.c
@@ -7,6 +7,9 @@
 
 #include <string.h>
 
+
+static __thread smrt_arena_t *_scratch_pool[SCRATCH_POOL_SIZE] = { 0 };
+
 smrt_arena_t *smrt_arena_create(u64 reserve_size, u64 commit_size, b32 auto_decommit) {
     u32 pagesize = plat_get_pagesize();
 
@@ -135,6 +138,52 @@ void smrt_arena_destroy(smrt_arena_t *arena) {
     plat_mem_release(arena, arena->reserve_size);
 }
 
+smrta_temp_t smrta_temp_start(smrt_arena_t *arena) {
+    return (smrta_temp_t){
+            .arena=arena,
+        .start_pos=arena->pos
+    };
+}
+
+void smrta_temp_end(smrta_temp_t temp) {
+    smrt_arena_pop_to(temp.arena, temp.start_pos);
+}
+
+smrta_temp_t smrta_scratch_start(smrt_arena_t **conflicts, u32 num_conflicts) {
+    i32 candidate_idx = -1;
+
+    for (i32 i = 0; i < SCRATCH_POOL_SIZE; i++) {
+        b32 conflict_found = false;
+
+        for (u32 j = 0; j < num_conflicts; j++) {
+            if (_scratch_pool[i] == conflicts[j]) {
+                conflict_found = true;
+            }
+        }
+
+        if (!conflict_found) {
+            candidate_idx = i;
+            break;
+        }
+    }
+
+    if (candidate_idx != -1) {
+        smrt_arena_t **selected = &_scratch_pool[candidate_idx];
+
+        if (!*selected) {
+            *selected = smrt_arena_create(MiB(64), plat_get_pagesize(), true);
+        }
+
+        return smrta_temp_start(*selected);
+    }
+
+    return (smrta_temp_t){ 0 };
+}
+
+void smrta_scratch_end(smrta_temp_t scratch) {
+    smrta_temp_end(scratch);
+}
+
 #if defined(_WIN32)
 
 #include <windows.h>
diff --git a/src/smrt_arena.h b/src/smrt_arena.h
index 375929d..0c8d848 100644
--- a/src/smrt_arena.h
+++ b/src/smrt_arena.h
@@ -18,16 +18,33 @@ typedef struct {
 
 smrt_arena_t     *smrt_arena_create(u64 reserve_size, u64 commit_size, b32 auto_decommit);
         void       *smrt_arena_push(smrt_arena_t *arena, u64 alloc_amount, b32 zero_out);
-        void  *smrt_arena_mark_push(smrt_arena_t *arena, u64 alloc_amount, b32 zero_out);
         void         smrt_arena_pop(smrt_arena_t *arena, u64 pop_amount);
         void      smrt_arena_pop_to(smrt_arena_t *arena, u64 location);
-         b32 smrt_arena_pop_to_mark(smrt_arena_t *arena);
         void       smrt_arena_clear(smrt_arena_t *arena, b32 zero_out);
-        void        smrt_arena_mark(smrt_arena_t *arena);
         void     smrt_arena_destroy(smrt_arena_t *arena);
 
+        void        smrt_arena_mark(smrt_arena_t *arena);
+        void  *smrt_arena_mark_push(smrt_arena_t *arena, u64 alloc_amount, b32 zero_out);
+         b32 smrt_arena_pop_to_mark(smrt_arena_t *arena);
+
 #define SMRTA_ALLOC_ARRAY(arena, T, n) (T*)smrt_arena_push(arena, sizeof(T) * (n), true)
 
+typedef struct {
+    smrt_arena_t *arena;
+            u64   start_pos;
+} smrta_temp_t;
+
+smrta_temp_t smrta_temp_start(smrt_arena_t *arena);
+        void   smrta_temp_end(smrta_temp_t temp);
+
+smrta_temp_t smrta_scratch_start(smrt_arena_t **conflicts, u32 num_conflicts);
+        void   smrta_scratch_end(smrta_temp_t temp);
+
+#define SCRATCH_POOL_SIZE 2
+
+#define TEMP_ARENA_PUSH(temp, T) (T*)smrt_arena_push(temp.arena, sizeof(T), true)
+#define TEMP_ARENA_PUSH_ARRAY(temp, T, n) (T*)smrt_arena_push(temp.arena, sizeof(T) * (n), true)
+
 u32 plat_get_pagesize(void);
 
 void *plat_mem_reserve(u64 size);
diff --git a/tests/smrt_arena.c b/tests/smrt_arena.c
index 87230d3..5e3ef5c 100644
--- a/tests/smrt_arena.c
+++ b/tests/smrt_arena.c
@@ -18,7 +18,7 @@ Test(smrt_arena, create_arena) {
 }
 
 Test(smrt_arena, alloc_arena) {
-    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), true, false);
+    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), plat_get_pagesize(), false);
 
     uint32_t *i = smrt_arena_push(arena, sizeof(uint32_t), true);
 
@@ -30,7 +30,7 @@ Test(smrt_arena, alloc_arena) {
 }
 
 Test(smrt_arena, alloc_zero_arena) {
-    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), true, false);
+    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), plat_get_pagesize(), false);
 
     uint32_t *i = smrt_arena_push(arena, sizeof(uint32_t), true);
 
@@ -43,7 +43,7 @@ Test(smrt_arena, alloc_array_arena) {
     const size_t len = 24;
     cr_expect(len != 0);
 
-    smrt_arena_t *arena = smrt_arena_create(sizeof(uint8_t) * len, true, false);
+    smrt_arena_t *arena = smrt_arena_create(sizeof(uint8_t) * len, plat_get_pagesize(), false);
 
     uint8_t *arr = SMRTA_ALLOC_ARRAY(arena, uint8_t, len);
 
@@ -58,7 +58,7 @@ Test(smrt_arena, alloc_array_arena) {
 }
 
 Test(smrt_arena, alloc_clear_arena) {
-    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), true, false);
+    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), plat_get_pagesize(), false);
 
     uint32_t *i = smrt_arena_push(arena, sizeof(uint32_t), true);
 
@@ -72,7 +72,7 @@ Test(smrt_arena, alloc_clear_arena) {
 }
 
 Test(smrt_arena, alloc_clear_zero_arena) {
-    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), true, false);
+    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), plat_get_pagesize(), false);
 
     uint32_t *i = smrt_arena_push(arena, sizeof(uint32_t), true);
 
@@ -95,7 +95,7 @@ Test(smrt_arena, alloc_clear_zero_arena) {
 }
 
 Test(smrt_arena, alloc_clear_dont_zero_arena) {
-    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), true, false);
+    smrt_arena_t *arena = smrt_arena_create(sizeof(uint32_t), plat_get_pagesize(), false);
 
     uint32_t *i = smrt_arena_push(arena, sizeof(uint32_t), true);
 
@@ -109,7 +109,7 @@ Test(smrt_arena, alloc_clear_dont_zero_arena) {
 }
 
 Test(smrt_arena, pop_arena) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     uint64_t mark = arena->pos;
     uint32_t *i = smrt_arena_push(arena, sizeof(uint32_t), true);
@@ -124,7 +124,7 @@ Test(smrt_arena, pop_arena) {
 }
 
 Test(smrt_arena, pop_to_arena) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     smrt_arena_push(arena, sizeof(uint32_t), true);
     uint64_t mark = arena->pos;
@@ -139,7 +139,7 @@ Test(smrt_arena, pop_to_arena) {
 }
 
 Test(smrt_arena, pop_clamped_arena) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     smrt_arena_push(arena, sizeof(uint32_t), true);
     smrt_arena_pop(arena, KiB(1));
@@ -164,7 +164,7 @@ Test(smrt_arena, push_grows_commit_arena) {
 }
 
 Test(smrt_arena, push_exceeds_reserve_arena) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     void *p = smrt_arena_push(arena, KiB(64), true);
 
@@ -174,7 +174,7 @@ Test(smrt_arena, push_exceeds_reserve_arena) {
 }
 
 Test(smrt_arena, push_alignment_arena) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     smrt_arena_push(arena, sizeof(uint8_t), true);
     void *p = smrt_arena_push(arena, sizeof(uint64_t), true);
@@ -185,13 +185,13 @@ Test(smrt_arena, push_alignment_arena) {
 }
 
 Test(smrt_arena, create_fails_when_reserve_exceeds_address_space) {
-    smrt_arena_t *arena = smrt_arena_create((u64)1 << 48, true, false);
+    smrt_arena_t *arena = smrt_arena_create((u64)1 << 48, plat_get_pagesize(), false);
 
     cr_expect_eq(arena, NULL);
 }
 
 Test(smrt_arena, mark_push_sets_mark) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     uint64_t pre_mark_pos = arena->pos;
     smrt_arena_mark_push(arena, sizeof(uint32_t), true);
@@ -202,7 +202,7 @@ Test(smrt_arena, mark_push_sets_mark) {
 }
 
 Test(smrt_arena, mark_push_failure_leaves_mark_unset) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     void *p = smrt_arena_mark_push(arena, KiB(64), true);
 
@@ -213,7 +213,7 @@ Test(smrt_arena, mark_push_failure_leaves_mark_unset) {
 }
 
 Test(smrt_arena, pop_to_mark_restores_position) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     uint32_t *i = smrt_arena_mark_push(arena, sizeof(uint32_t), true);
     smrt_arena_push(arena, sizeof(uint32_t), true);
@@ -229,7 +229,7 @@ Test(smrt_arena, pop_to_mark_restores_position) {
 }
 
 Test(smrt_arena, mark_at_base_position_is_valid) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     cr_assert_eq(arena->pos, SMRT_ARENA_BASE_POS);
 
@@ -246,7 +246,7 @@ Test(smrt_arena, mark_at_base_position_is_valid) {
 }
 
 Test(smrt_arena, pop_to_mark_without_mark_fails) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     smrt_arena_push(arena, sizeof(uint32_t), true);
 
@@ -258,7 +258,7 @@ Test(smrt_arena, pop_to_mark_without_mark_fails) {
 }
 
 Test(smrt_arena, pop_to_ahead_of_pos_is_noop) {
-    smrt_arena_t *arena = smrt_arena_create(KiB(1), true, false);
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
 
     smrt_arena_push(arena, sizeof(uint32_t), true);
     uint64_t pos = arena->pos;
@@ -348,6 +348,103 @@ Test(smrt_arena, auto_decommit_pop_past_mark_resets_mark) {
     smrt_arena_destroy(arena);
 }
 
+Test(smrt_arena, mark_sets_current_pos) {
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
+
+    smrt_arena_push(arena, sizeof(uint32_t), true);
+    uint64_t pos = arena->pos;
+
+    smrt_arena_mark(arena);
+
+    cr_expect_eq(arena->mark_pos, pos);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(smrt_arena, temp_start_captures_arena_and_pos) {
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
+    smrt_arena_push(arena, sizeof(uint32_t), true);
+
+    smrta_temp_t temp = smrta_temp_start(arena);
+
+    cr_expect_eq(temp.arena, arena);
+    cr_expect_eq(temp.start_pos, arena->pos);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(smrt_arena, temp_start_end_restores_position) {
+    smrt_arena_t *arena = smrt_arena_create(KiB(1), plat_get_pagesize(), false);
+
+    smrt_arena_push(arena, sizeof(uint32_t), true);
+    uint64_t pos_before = arena->pos;
+
+    smrta_temp_t temp = smrta_temp_start(arena);
+    smrt_arena_push(arena, sizeof(uint32_t), true);
+    smrt_arena_push(arena, sizeof(uint32_t), true);
+
+    smrta_temp_end(temp);
+
+    cr_expect_eq(arena->pos, pos_before);
+
+    smrt_arena_destroy(arena);
+}
+
+Test(smrt_arena, scratch_start_end_provides_usable_arena) {
+    smrta_temp_t scratch = smrta_scratch_start(NULL, 0);
+
+    cr_assert_neq(scratch.arena, NULL);
+
+    uint32_t *i = TEMP_ARENA_PUSH(scratch, uint32_t);
+    *i = UINT32_MAX;
+
+    cr_expect_eq(*i, UINT32_MAX);
+
+    smrta_scratch_end(scratch);
+}
+
+Test(smrt_arena, scratch_end_pops_arena_to_start) {
+    smrta_temp_t scratch = smrta_scratch_start(NULL, 0);
+    uint64_t pos_before = scratch.arena->pos;
+
+    TEMP_ARENA_PUSH(scratch, uint32_t);
+    cr_assert_gt(scratch.arena->pos, pos_before);
+
+    smrta_scratch_end(scratch);
+
+    cr_expect_eq(scratch.arena->pos, pos_before);
+}
+
+Test(smrt_arena, scratch_start_avoids_conflicts) {
+    smrta_temp_t first = smrta_scratch_start(NULL, 0);
+    cr_assert_neq(first.arena, NULL);
+
+    smrt_arena_t *conflicts[] = { first.arena };
+    smrta_temp_t second = smrta_scratch_start(conflicts, 1);
+
+    cr_assert_neq(second.arena, NULL);
+    cr_expect_neq(second.arena, first.arena);
+
+    smrta_scratch_end(second);
+    smrta_scratch_end(first);
+}
+
+Test(smrt_arena, scratch_start_returns_zeroed_when_pool_exhausted) {
+    smrta_temp_t a = smrta_scratch_start(NULL, 0);
+
+    smrt_arena_t *a_conflict[] = { a.arena };
+    smrta_temp_t b = smrta_scratch_start(a_conflict, 1);
+    cr_assert_neq(b.arena, a.arena);
+
+    smrt_arena_t *conflicts[] = { a.arena, b.arena };
+    smrta_temp_t c = smrta_scratch_start(conflicts, 2);
+
+    cr_expect_eq(c.arena, NULL);
+
+    smrta_scratch_end(b);
+    smrta_scratch_end(a);
+}
+
 Test(smrt_arena, decommit_then_recommit_pages) {
     u32 pagesize = plat_get_pagesize();
     void *mem = plat_mem_reserve(pagesize);