Pre-size arena to avoid unnecessary growth

Author: kddnewton

include/prism/util/pm_arena.h

@@ -44,6 +44,16 @@ typedef struct {
     size_t block_count;
 } pm_arena_t;
 
+/**
+ * Ensure the arena has at least `capacity` bytes available in its current
+ * block, allocating a new block if necessary. This allows callers to
+ * pre-size the arena to avoid repeated small block allocations.
+ *
+ * @param arena The arena to pre-size.
+ * @param capacity The minimum number of bytes to ensure are available.
+ */
+void pm_arena_reserve(pm_arena_t *arena, size_t capacity);
+
 /**
  * Allocate memory from the arena. The returned memory is NOT zeroed. This
  * function is infallible — it aborts on allocation failure.

src/prism.c

@@ -22173,6 +22173,14 @@ pm_parser_init(pm_arena_t *arena, pm_parser_t *parser, const uint8_t *source, si
         .warn_mismatched_indentation = true
     };
 
+    // Pre-size the arenas based on input size to reduce the number of block
+    // allocations (and the kernel page zeroing they trigger). The ratios were
+    // measured empirically: AST arena ~3.3x input, metadata arena ~1.1x input.
+    // The reserve call is a no-op when the capacity is at or below the default
+    // arena block size, so small inputs don't waste an extra allocation.
+    if (size <= SIZE_MAX / 4) pm_arena_reserve(arena, size * 4);
+    if (size <= SIZE_MAX / 5 * 4) pm_arena_reserve(&parser->metadata_arena, size + size / 4);
+
     // Initialize the constant pool. We're going to completely guess as to the
     // number of constants that we'll need based on the size of the input. The
     // ratio we chose here is actually less arbitrary than you might think.

src/util/pm_arena.c

@@ -1,5 +1,7 @@
 #include "prism/util/pm_arena.h"
 
+#include <assert.h>
+
 /**
  * Compute the block allocation size using offsetof so it is correct regardless
  * of PM_FLEX_ARY_LEN.
@@ -29,6 +31,42 @@ pm_arena_next_block_size(const pm_arena_t *arena, size_t min_size) {
     return size > min_size ? size : min_size;
 }
 
+/**
+ * Allocate a new block with the given data capacity and initial usage, link it
+ * into the arena, and return it. Aborts on allocation failure.
+ */
+static pm_arena_block_t *
+pm_arena_new_block(pm_arena_t *arena, size_t data_size, size_t initial_used) {
+    assert(initial_used <= data_size);
+    pm_arena_block_t *block = (pm_arena_block_t *) xmalloc(PM_ARENA_BLOCK_SIZE(data_size));
+
+    if (block == NULL) {
+        fprintf(stderr, "prism: out of memory; aborting\n");
+        abort();
+    }
+
+    block->capacity = data_size;
+    block->used = initial_used;
+    block->prev = arena->current;
+    arena->current = block;
+    arena->block_count++;
+
+    return block;
+}
+
+/**
+ * Ensure the arena has at least `capacity` bytes available in its current
+ * block, allocating a new block if necessary. This allows callers to
+ * pre-size the arena to avoid repeated small block allocations.
+ */
+void
+pm_arena_reserve(pm_arena_t *arena, size_t capacity) {
+    if (capacity <= PM_ARENA_INITIAL_SIZE) return;
+    if (arena->current != NULL && (arena->current->capacity - arena->current->used) >= capacity) return;
+
+    pm_arena_new_block(arena, capacity, 0);
+}
+
 /**
  * Allocate memory from the arena. The returned memory is NOT zeroed. This
  * function is infallible — it aborts on allocation failure.
@@ -51,18 +89,7 @@ pm_arena_alloc(pm_arena_t *arena, size_t size, size_t alignment) {
     // New blocks from xmalloc are max-aligned, so data[] starts aligned for
     // any C type. No padding needed at the start.
     size_t block_data_size = pm_arena_next_block_size(arena, size);
-    pm_arena_block_t *block = (pm_arena_block_t *) xmalloc(PM_ARENA_BLOCK_SIZE(block_data_size));
-
-    if (block == NULL) {
-        fprintf(stderr, "prism: out of memory; aborting\n");
-        abort();
-    }
-
-    block->capacity = block_data_size;
-    block->used = size;
-    block->prev = arena->current;
-    arena->current = block;
-    arena->block_count++;
+    pm_arena_block_t *block = pm_arena_new_block(arena, block_data_size, size);
 
     return block->data;
 }