alloc: add Pool and Chunked

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")
+	}
+}

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.Pool[int](16)
+
+	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(16 * 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.Chunked(1 * 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
+}

alloc/allocators.go

@@ -0,0 +1,228 @@
+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(max_size uintptr) Allocator {
+	if max_size == 0 {
+		panic("linear: max_size must be greater than zero")
+	}
+
+	var (
+		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 > max_size {
+				return nil, fmt.Errorf("linear: out of memory - %d bytes requested, %d bytes free", size, max_size-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
+	}
+}
+
+// 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.
+//
+// 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")
+	}
+}
+
+// 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
+	}
+}

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.