xx/alloc/allocators.go

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 := 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
	}
}