From 9abadd3c15a5bec99f3e6e72999dd96ad54e6baf Mon Sep 17 00:00:00 2001
From: Judah Caruso <judah@tuta.io>
Date: Fri, 12 Dec 2025 17:11:22 -0700
Subject: [PATCH] add alloc package

---
 alloc/alloc.go      |  78 ++++++++++++++++
 alloc/alloc_test.go | 214 ++++++++++++++++++++++++++++++++++++++++++++
 alloc/allocators.go | 129 ++++++++++++++++++++++++++
 mem/mem.go          |   7 ++
 4 files changed, 428 insertions(+)
 create mode 100644 alloc/alloc.go
 create mode 100644 alloc/alloc_test.go
 create mode 100644 alloc/allocators.go

diff --git a/alloc/alloc.go b/alloc/alloc.go
new file mode 100644
index 0000000..be53d17
--- /dev/null
+++ b/alloc/alloc.go
@@ -0,0 +1,78 @@
+package alloc
+
+import (
+	"unsafe"
+
+	"git.brut.systems/judah/xx/mem"
+)
+
+// New returns a pointer to an Allocator allocated value of type T.
+//
+// Note: If allocation fails, New will panic.
+func New[T any](alloc Allocator) *T {
+	ptr, err := alloc(ActionAlloc, mem.SizeOf[T](), mem.AlignOf[T](), nil)
+	if err != nil {
+		panic(err)
+	}
+
+	return (*T)(ptr)
+}
+
+// Reset restores an Allocator to its initial state.
+//
+// Note: Use of memory allocated by an Allocator after calling Reset is unsafe.
+func Reset(alloc Allocator) {
+	if _, err := alloc(ActionReset, 0, 0, nil); err != nil {
+		panic(err)
+	}
+}
+
+// Save returns the current state of an Allocator.
+//
+// Note: The value returned is internal to the particular Allocator Save was called on.
+// The value should not be modified.
+func Save(alloc Allocator) (watermark uintptr) {
+	if _, err := alloc(ActionSave, 0, 0, &watermark); err != nil {
+		panic(err)
+	}
+
+	return
+}
+
+// Restore restores an Allocator to a previously saved state.
+func Restore(alloc Allocator, watermark uintptr) {
+	if _, err := alloc(ActionRestore, 0, 0, &watermark); err != nil {
+		panic(err)
+	}
+}
+
+// Allocator represents a memory allocator.
+type Allocator func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error)
+
+// Action is a list of distinct events an Allocator may respond to.
+type Action int
+
+const (
+	ActionAlloc Action = iota
+	ActionReset
+	ActionSave
+	ActionRestore
+	ActionReport
+)
+
+func (a Action) String() string {
+	switch a {
+	case ActionAlloc:
+		return "Alloc"
+	case ActionReset:
+		return "Reset"
+	case ActionSave:
+		return "Save"
+	case ActionRestore:
+		return "Restore"
+	case ActionReport:
+		return "Report"
+	default:
+		panic("unreachable")
+	}
+}
diff --git a/alloc/alloc_test.go b/alloc/alloc_test.go
new file mode 100644
index 0000000..f213f98
--- /dev/null
+++ b/alloc/alloc_test.go
@@ -0,0 +1,214 @@
+package alloc_test
+
+import (
+	"errors"
+	"fmt"
+	"runtime"
+	"testing"
+	"unsafe"
+
+	"git.brut.systems/judah/xx/alloc"
+	"git.brut.systems/judah/xx/mem"
+)
+
+func BenchmarkAlloc_New_Small(b *testing.B) {
+	var last *int
+	for i := range b.N {
+		v := new(int)
+		*v = i
+		last = v
+
+		if i%1000 == 0 {
+			last = nil
+		}
+	}
+
+	runtime.KeepAlive(last)
+}
+
+func BenchmarkAlloc_Closure_Small(b *testing.B) {
+	allocator := alloc.Linear(32 * mem.Kilobyte)
+
+	var last *int
+	for i := range b.N {
+		v := alloc.New[int](allocator)
+		*v = i
+		last = v
+
+		if i%1000 == 0 {
+			alloc.Reset(allocator)
+		}
+	}
+
+	runtime.KeepAlive(last)
+}
+
+func BenchmarkAlloc_Interface_Small(b *testing.B) {
+	allocator := NewLinear(32 * mem.Kilobyte)
+
+	var last *int
+	for i := range b.N {
+		v := New[int](&allocator)
+		*v = i
+		last = v
+
+		if i%1000 == 0 {
+			Reset(&allocator)
+		}
+	}
+
+	runtime.KeepAlive(last)
+}
+
+type large struct{ a, b, c, d, e, f, g, h, i int }
+
+func BenchmarkAlloc_New_Large(b *testing.B) {
+	var last *large
+	for i := range b.N {
+		v := new(large)
+		v.e = i
+		last = v
+
+		if i%1000 == 0 {
+			last = nil
+		}
+	}
+
+	runtime.KeepAlive(last)
+}
+
+func BenchmarkAlloc_Closure_Large(b *testing.B) {
+	allocator := alloc.Linear(128 * mem.Kilobyte)
+
+	var last *large
+	for i := range b.N {
+		v := alloc.New[large](allocator)
+		v.e = i
+		last = v
+
+		if i%1000 == 0 {
+			alloc.Reset(allocator)
+		}
+	}
+
+	runtime.KeepAlive(last)
+}
+
+func BenchmarkAlloc_Interface_Large(b *testing.B) {
+	allocator := NewLinear(128 * mem.Kilobyte)
+
+	var last *large
+	for i := range b.N {
+		v := New[large](&allocator)
+		v.e = i
+		last = v
+
+		if i%1000 == 0 {
+			Reset(&allocator)
+		}
+	}
+
+	runtime.KeepAlive(last)
+}
+
+func BenchmarkAlloc_Closure_HotPath(b *testing.B) {
+	allocator := alloc.Linear(8 * mem.Kilobyte)
+
+	var (
+		lastlarge *large
+		lastsmall *int
+	)
+	for i := range b.N {
+		if i%2 == 0 {
+			lastsmall = alloc.New[int](allocator)
+		} else {
+			lastlarge = alloc.New[large](allocator)
+		}
+
+		alloc.Reset(allocator)
+	}
+
+	runtime.KeepAlive(lastlarge)
+	runtime.KeepAlive(lastsmall)
+}
+
+func BenchmarkAlloc_Interface_HotPath(b *testing.B) {
+	allocator := NewLinear(8 * mem.Kilobyte)
+
+	var (
+		lastlarge *large
+		lastsmall *int
+	)
+	for i := range b.N {
+		if i%2 == 0 {
+			lastsmall = New[int](&allocator)
+		} else {
+			lastlarge = New[large](&allocator)
+		}
+
+		Reset(&allocator)
+	}
+
+	runtime.KeepAlive(lastlarge)
+	runtime.KeepAlive(lastsmall)
+}
+
+type Allocator interface {
+	Proc(a alloc.Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error)
+}
+
+func New[T any](a Allocator) *T {
+	ptr, err := a.Proc(alloc.ActionAlloc, mem.SizeOf[T](), mem.AlignOf[T](), nil)
+	if err != nil {
+		panic(err)
+	}
+
+	return (*T)(ptr)
+}
+
+func Reset(a Allocator) {
+	if _, err := a.Proc(alloc.ActionReset, 0, 0, nil); err != nil {
+		panic(err)
+	}
+}
+
+type Linear struct {
+	data    []byte
+	maxsize uintptr
+	offset  uintptr
+}
+
+func NewLinear(maxsize uintptr) Linear {
+	return Linear{
+		data:    make([]byte, maxsize),
+		maxsize: maxsize,
+	}
+}
+
+func (l *Linear) Proc(a alloc.Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
+	switch a {
+	case alloc.ActionAlloc:
+		aligned := mem.AlignForward(size, align)
+		if l.offset+aligned > l.maxsize {
+			return nil, errors.New(fmt.Sprintf("Linear: out of memory - %d bytes requested, (%d/%d) bytes available", size, l.maxsize-l.offset, l.maxsize))
+		}
+
+		ptr := &l.data[l.offset]
+		l.offset += aligned
+		return unsafe.Pointer(ptr), nil
+
+	case alloc.ActionReset:
+		clear(l.data)
+		l.offset = 0
+
+	case alloc.ActionSave:
+		*watermark = l.offset
+	case alloc.ActionRestore:
+		l.offset = *watermark
+
+	default:
+		panic("unimplemented action: " + a.String())
+	}
+
+	return nil, nil
+}
diff --git a/alloc/allocators.go b/alloc/allocators.go
new file mode 100644
index 0000000..8e36132
--- /dev/null
+++ b/alloc/allocators.go
@@ -0,0 +1,129 @@
+package alloc
+
+import (
+	"errors"
+	"fmt"
+	"log"
+	"runtime"
+	"sync"
+	"unsafe"
+
+	"git.brut.systems/judah/xx/mem"
+)
+
+// Linear is a simple bump allocator with a fixed amount of backing memory.
+func Linear(maxsize uintptr) Allocator {
+	var (
+		data   = make([]byte, maxsize)
+		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)
+			}
+
+			ptr := &data[offset]
+			offset += aligned
+			return unsafe.Pointer(ptr), nil
+
+		case ActionReset:
+			clear(data)
+			offset = 0
+
+		case ActionSave:
+			*watermark = offset
+		case ActionRestore:
+			offset = *watermark
+
+		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
+func Nil() Allocator {
+	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
+		return nil, errors.New("use of nil allocator")
+	}
+}
+
+// Temporary wraps an Allocator, restoring it to its previous state when Reset is called.
+func Temporary(alloc Allocator) Allocator {
+	watermark := new(uintptr)
+	*watermark = Save(alloc)
+	return func(a Action, size, align uintptr, wm *uintptr) (unsafe.Pointer, error) {
+		if a == ActionReset {
+			Restore(alloc, *watermark)
+			return nil, nil
+		}
+
+		return alloc(a, size, align, wm)
+	}
+}
+
+// Split wraps two [[Allocator]]s, dispatching actions based on the size of the allocation.
+func Split(split_size uintptr, smaller, larger Allocator) Allocator {
+	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
+		if size <= split_size {
+			return smaller(a, size, align, watermark)
+		}
+		return larger(a, size, align, watermark)
+	}
+}
+
+// Logger wraps an Allocator, logging its usage locations.
+func Logger(alloc Allocator) Allocator {
+	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
+		// We expect allocators to be used via the high-level API, so we grab the caller location relative to that.
+		// @todo(judah): can we determine this dynamically?
+		_, file, line, ok := runtime.Caller(2)
+		if !ok {
+			file = "<unknown>"
+			line = 0
+		}
+
+		log.Printf("%s:%d - %s (size: %d, align: %d, watermark: %p)", file, line, a, size, align, watermark)
+		return alloc(a, size, align, watermark)
+	}
+}
+
+// Concurrent wraps an Allocator, ensuring it is safe for concurrent use.
+func Concurrent(alloc Allocator) Allocator {
+	mtx := new(sync.Mutex)
+	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
+		mtx.Lock()
+		ptr, err := alloc(a, size, align, watermark)
+		mtx.Unlock()
+		return ptr, err
+	}
+}
+
+// Pinned wraps an Allocator, ensuring the memory returned is stable until Reset is called.
+//
+// The memory returned by Pinned is safe to pass over cgo boundaries.
+func Pinned(alloc Allocator) Allocator {
+	var pinner runtime.Pinner
+	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
+		ptr, err := alloc(a, size, align, watermark)
+		if err != nil {
+			return ptr, err
+		}
+
+		if a == ActionReset {
+			pinner.Unpin()
+		} else {
+			pinner.Pin(ptr)
+		}
+
+		return ptr, err
+	}
+}
diff --git a/mem/mem.go b/mem/mem.go
index 128c917..8c43319 100644
--- a/mem/mem.go
+++ b/mem/mem.go
@@ -4,6 +4,13 @@ import (
 	"unsafe"
 )
 
+const (
+	Kilobyte uintptr = 1 << (10 * (iota + 1))
+	Megabyte
+	Gigabyte
+	Terabyte
+)
+
 // SizeOf returns the size (in bytes) of the given type.
 //
 // Not to be confused with [unsafe.Sizeof] which returns the size of a type via an expression.