memory/arena: add snapshot and rollback APIs

src/memory/arena.c

@@ -193,3 +193,35 @@ void *marena_realloc_aligned(MemArena *restrict arena, void *restrict p, size_t
 
 	return _arena_realloc(arena, p, old_size, new_size, align);
 }
+
+MemArenaSnapshot marena_snapshot(MemArena *arena) {
+	return (MemArenaSnapshot) {
+		.page = _arena_active_page(arena),
+		.page_offset = arena->page_offset,
+	};
+}
+
+bool marena_rollback(MemArena *restrict arena, const MemArenaSnapshot *restrict snapshot) {
+	auto active_page = _arena_active_page(arena);
+
+	if(active_page == snapshot->page) {
+		if(UNLIKELY(snapshot->page_offset > arena->page_offset)) {
+			return false;
+		}
+
+		size_t mem_diff = arena->page_offset - snapshot->page_offset;
+		arena->page_offset = snapshot->page_offset;
+		assert(arena->total_used >= mem_diff);
+		arena->total_used -= mem_diff;
+		return true;
+	}
+
+	// New page(s) have been allocated after the snapshot was taken.
+	// We won't try to undo that, but we can at least reset the active page.
+
+	assert(arena->total_used >= arena->page_offset);
+	arena->total_used -= arena->page_offset;
+	arena->page_offset = 0;
+
+	return false;
+}

src/memory/arena.h

@@ -13,6 +13,7 @@
 
 typedef struct MemArena MemArena;
 typedef struct MemArenaPage MemArenaPage;
+typedef struct MemArenaSnapshot MemArenaSnapshot;
 
 struct MemArena {
 	LIST_ANCHOR(MemArenaPage) pages;
@@ -28,6 +29,11 @@ struct MemArenaPage {
 	alignas(alignof(max_align_t)) char data[];
 };
 
+typedef struct MemArenaSnapshot {
+	MemArenaPage *page;
+	size_t page_offset;
+} MemArenaSnapshot;
+
 void marena_init(MemArena *arena, size_t min_size)
 	attr_nonnull_all;
 
@@ -82,6 +88,12 @@ void *marena_realloc_aligned(MemArena *restrict arena, void *restrict p, size_t
 	attr_nonnull(1)
 	attr_returns_allocated;
 
+MemArenaSnapshot marena_snapshot(MemArena *arena)
+	attr_nonnull_all;
+
+bool marena_rollback(MemArena *restrict arena, const MemArenaSnapshot *restrict snapshot)
+	attr_nonnull_all;
+
 INLINE void *marena_memdup(MemArena *arena, const void *buf, size_t size) {
 	return memcpy(marena_alloc(arena, size), buf, size);
 }