From c42cdcbaa63a6c620f762444de9eb9432e1cff96 Mon Sep 17 00:00:00 2001
From: Judah Caruso <judah@tuta.io>
Date: Fri, 12 Dec 2025 18:24:11 -0700
Subject: [PATCH] alloc: add Pool and Chunked

---
 alloc/alloc_test.go |   6 +--
 alloc/allocators.go | 111 +++++++++++++++++++++++++++++++++++++++++---
 2 files changed, 108 insertions(+), 9 deletions(-)

diff --git a/alloc/alloc_test.go b/alloc/alloc_test.go
index f213f98..7e55be1 100644
--- a/alloc/alloc_test.go
+++ b/alloc/alloc_test.go
@@ -27,7 +27,7 @@ func BenchmarkAlloc_New_Small(b *testing.B) {
 }
 
 func BenchmarkAlloc_Closure_Small(b *testing.B) {
-	allocator := alloc.Linear(32 * mem.Kilobyte)
+	allocator := alloc.Pool[int](16)
 
 	var last *int
 	for i := range b.N {
@@ -44,7 +44,7 @@ func BenchmarkAlloc_Closure_Small(b *testing.B) {
 }
 
 func BenchmarkAlloc_Interface_Small(b *testing.B) {
-	allocator := NewLinear(32 * mem.Kilobyte)
+	allocator := NewLinear(16 * mem.Kilobyte)
 
 	var last *int
 	for i := range b.N {
@@ -112,7 +112,7 @@ func BenchmarkAlloc_Interface_Large(b *testing.B) {
 }
 
 func BenchmarkAlloc_Closure_HotPath(b *testing.B) {
-	allocator := alloc.Linear(8 * mem.Kilobyte)
+	allocator := alloc.Chunked(1 * mem.Kilobyte)
 
 	var (
 		lastlarge *large
diff --git a/alloc/allocators.go b/alloc/allocators.go
index 8e36132..d1e4498 100644
--- a/alloc/allocators.go
+++ b/alloc/allocators.go
@@ -12,18 +12,21 @@ import (
 )
 
 // Linear is a simple bump allocator with a fixed amount of backing memory.
-func Linear(maxsize uintptr) Allocator {
+func Linear(max_size uintptr) Allocator {
+	if max_size == 0 {
+		panic("linear: max_size must be greater than zero")
+	}
+
 	var (
-		data   = make([]byte, maxsize)
+		data   = make([]byte, max_size)
 		offset uintptr
 	)
-
 	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
 		switch a {
 		case ActionAlloc:
 			aligned := mem.AlignForward(size, align)
-			if offset+aligned > maxsize {
-				return nil, fmt.Errorf("linear: out of memory - %d bytes requested, %d bytes free", size, maxsize-offset)
+			if offset+aligned > max_size {
+				return nil, fmt.Errorf("linear: out of memory - %d bytes requested, %d bytes free", size, max_size-offset)
 			}
 
 			ptr := &data[offset]
@@ -47,9 +50,105 @@ func Linear(maxsize uintptr) Allocator {
 	}
 }
 
+// Pool is an Allocator that only allocates values of a single type.
+//
+// Note: Allocating different types from the same Pool is unsafe and may cause memory corruption.
+func Pool[T any](base_capacity uintptr) Allocator {
+	if base_capacity == 0 {
+		panic("pool: base_capacity must be greater than zero")
+	}
+
+	pointers := make([]T, 0, base_capacity)
+	return func(a Action, _, _ uintptr, watermark *uintptr) (unsafe.Pointer, error) {
+		switch a {
+		case ActionAlloc:
+			pointers = append(pointers, mem.ZeroValue[T]())
+			return unsafe.Pointer(&pointers[len(pointers)-1]), nil
+
+		case ActionReset:
+			clear(pointers)
+			pointers = pointers[:0]
+
+		case ActionSave:
+			*watermark = uintptr(len(pointers))
+
+		case ActionRestore:
+			clear(pointers[*watermark:])
+			pointers = pointers[:*watermark]
+
+		default:
+		}
+
+		return nil, nil
+	}
+}
+
+// Chunked is an Allocator that groups allocations by size.
+func Chunked(chunk_size uintptr) Allocator {
+	if chunk_size == 0 {
+		panic("chunked: chunk_size must be greater than zero")
+	}
+
+	type chunk struct {
+		data   []byte
+		offset uintptr
+	}
+
+	groups := make(map[uintptr][]chunk)
+	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
+		switch a {
+		case ActionAlloc:
+			aligned := mem.AlignForward(size, align)
+			group, ok := groups[aligned]
+			if !ok {
+				group = make([]chunk, 0, 16)
+				group = append(group, chunk{
+					data:   make([]byte, chunk_size),
+					offset: 0,
+				})
+
+				groups[aligned] = group
+			}
+
+			c := &group[len(group)-1]
+			if c.offset+aligned > chunk_size {
+				group = append(group, chunk{
+					data:   make([]byte, chunk_size),
+					offset: 0,
+				})
+
+				c = &group[len(group)-1]
+				groups[aligned] = group
+			}
+
+			ptr := &c.data[c.offset]
+			c.offset += aligned
+
+			return unsafe.Pointer(ptr), nil
+
+		case ActionReset:
+			for _, g := range groups {
+				for i := range len(g) {
+					c := &g[i]
+					c.offset = 0
+					clear(c.data)
+				}
+			}
+
+		case ActionSave:
+		case ActionRestore:
+
+		default:
+			panic("unimplemented action: " + a.String())
+		}
+
+		return nil, nil
+	}
+}
+
 // Nil is an Allocator that always returns an error.
 //
-// This is useful for tracking usage locations
+// Note: This is useful for tracking usage locations
 func Nil() Allocator {
 	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
 		return nil, errors.New("use of nil allocator")