An arena's budget holds when a block grows in place, and a shrink is never over budget

runtime/flan_rt.c

@@ -1185,7 +1185,7 @@ static void *flan_heap_proc(flan_allocator *a, int32_t mode, void *p,
      * caller passes old_size for exactly this reason, and it is the one
      * number a wrong answer here would read off the end of. */
     void *q;
-    if (flan_over_budget(a, size - old_size)) return NULL;
+    if (size > old_size && flan_over_budget(a, size - old_size)) return NULL;
     q = flan_heap_proc(a, FLAN_ALLOC_ALLOC, NULL, 0, size, align);
     if (!q) return NULL;
     if (p && old_size > 0)
@@ -1331,6 +1331,9 @@ static void *flan_arena_proc(flan_allocator *a, int32_t mode, void *p,
      * without copying, which is the common shape. */
     if (p && (uint8_t *)p + old_size == ar->base + ar->offset) {
       int64_t end = (int64_t)((uint8_t *)p - ar->base) + size;
+      /* The budget is checked here as on every other path; a block grown in
+       * place is still more live bytes. */
+      if (size > old_size && flan_over_budget(a, size - old_size)) return NULL;
       if (end > ar->cap || end < 0) return NULL;
       ar->offset = end;
       if (end > ar->peak) ar->peak = end;

test/programs/exhausted.flan

@@ -115,6 +115,24 @@
       (println "INSERTIONSORT")))       ; INSERTIONSORT
   (println failures)                    ; 1 — failed once, retried once
 
+  ;; And an arena, whose budget is checked when a Vec grows its block in
+  ;; place as well as when it allocates a new one. A Vec that is the only thing
+  ;; pushing into an arena always grows in place, so without that check the
+  ;; ceiling would never be met.
+  (set tight (arena-new 65536))
+  (set-alloc-budget tight 64)
+  (set failures 0)
+  (handler-bind
+    [(StorageExhausted [c]
+       (set failures (+ failures 1))
+       (set-alloc-budget tight (* 2 (alloc-budget tight)))
+       (invoke-restart 'retry))]
+    (let [v (vec-new i32 tight)]
+      (dotimes [i 1000] (push v i))
+      (println (length v))              ; 1000
+      (println (at v 999))))            ; 999
+  (println (> failures 0))              ; true
+
   ;; And the restart is not once-per-program: it is established at each
   ;; allocation, so a later one offers it again.
   (set-alloc-budget tight 0)

test/test_acceptance.ml

@@ -1663,7 +1663,7 @@ let () =
        flow an optimiser would otherwise launder. *)
     let exhausted_out =
       "64\n0\n126\ntrue\ntrue\n4\ntrue\n0\n8\n7\ntrue\n\
-       13\n73\n84\n90\nINSERTIONSORT\n1\n"
+       13\n73\n84\n90\nINSERTIONSORT\n1\n1000\n999\ntrue\n"
     in
     outputs "storage exhausted, retried" "programs/exhausted.flan" exhausted_out;
     outputs ~opt:"-O0" "storage exhausted, retried, -O0" "programs/exhausted.flan"