xx/arena/arenas.go

package arena

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) Arena {
	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 ACTION_ALLOC:
			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 ACTION_RESET:
			clear(data)
			offset = 0
		case ACTION_SAVE:
			*watermark = offset
		case ACTION_RESTORE:
			offset = *watermark
		default:
			panic("linear: unimplemented action - " + a.String())
		}

		return nil, nil
	}
}

// Pool is an Arena that only allocates values of the given type.
//
// Note: Allocating different types from the same Pool is unsafe and may cause memory corruption.
func Pool[T any](base_capacity uintptr) Arena {
	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 ACTION_ALLOC:
			pointers = append(pointers, mem.ZeroValue[T]())
			return unsafe.Pointer(&pointers[len(pointers)-1]), nil
		case ACTION_RESET:
			clear(pointers)
			pointers = pointers[:0]
		case ACTION_SAVE:
			*watermark = uintptr(len(pointers))
		case ACTION_RESTORE:
			clear(pointers[*watermark:])
			pointers = pointers[:*watermark]
		default:
			panic("pool: unimplemented action - " + a.String())
		}

		return nil, nil
	}
}

// Chunked is an Arena that groups allocations by size.
func Chunked(chunk_size uintptr) Arena {
	if chunk_size <= 0 {
		panic("chunked: chunk_size must be greater than zero")
	}

	// @todo(judah): this can be drastically improved.

	type chunk struct {
		data   []byte
		offset uintptr
		saved  uintptr
	}

	groups := make(map[uintptr][]chunk)
	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
		switch a {
		case ACTION_ALLOC:
			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),
				})

				groups[aligned] = group
			}

			c := &group[len(group)-1]
			if c.offset+aligned > chunk_size {
				group = append(group, chunk{
					data: make([]byte, chunk_size),
				})

				c = &group[len(group)-1]
				groups[aligned] = group
			}

			ptr := &c.data[c.offset]
			c.offset += aligned

			return unsafe.Pointer(ptr), nil
		case ACTION_RESET:
			for _, g := range groups {
				for i := range len(g) {
					c := &g[i]
					c.offset = 0
					c.saved = 0
					clear(c.data)
				}
			}
		case ACTION_SAVE:
			for _, g := range groups {
				for i := range len(g) {
					c := &g[i]
					c.saved = c.offset
				}
			}
		case ACTION_RESTORE:
			for _, g := range groups {
				for i := range len(g) {
					c := &g[i]
					c.offset = c.saved
				}
			}
		default:
			panic("chunked: unimplemented action - " + a.String())
		}

		return nil, nil
	}
}

// Nil is an Arena that always returns an error.
//
// Note: This is useful for tracking usage locations
func Nil() Arena {
	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
		return nil, errors.New("use of nil allocator")
	}
}

// Region wraps an Arena, restoring it to its previous state when Reset is called.
func Region(arena Arena) Arena {
	watermark := Save(arena)
	return func(a Action, size, align uintptr, wm *uintptr) (unsafe.Pointer, error) {
		if a == ACTION_RESET {
			Restore(arena, watermark)
			return nil, nil
		}

		return arena(a, size, align, wm)
	}
}

// Split wraps two [[Arena]]s, dispatching allocations to a particular one based on the requested size.
func Split(split_size uintptr, smaller, larger Arena) Arena {
	var watermarks [2]uintptr
	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
		switch a {
		case ACTION_ALLOC:
			if size <= split_size {
				return smaller(a, size, align, watermark)
			}
			return larger(a, size, align, watermark)
		case ACTION_RESET:
			Reset(smaller)
			Reset(larger)
		case ACTION_SAVE:
			watermarks[0] = Save(smaller)
			watermarks[1] = Save(larger)
		case ACTION_RESTORE:
			Restore(smaller, watermarks[0])
			Restore(larger, watermarks[1])
		default:
			panic("split: unimplemented action - " + a.String())
		}

		return nil, nil
	}
}

// Logger wraps an Arena, logging its usage locations.
func Logger(arena Arena) Arena {
	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 arena(a, size, align, watermark)
	}
}

// Concurrent wraps an Arena, ensuring it is safe for concurrent use.
func Concurrent(arena Arena) Arena {
	mtx := new(sync.Mutex)
	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
		mtx.Lock()
		ptr, err := arena(a, size, align, watermark)
		mtx.Unlock()
		return ptr, err
	}
}

// Pinned wraps an Arena, ensuring the memory returned is stable until Reset is called.
//
// The memory returned by Pinned is safe to pass over cgo boundaries.
func Pinned(arena Arena) Arena {
	var pinner runtime.Pinner
	return func(a Action, size, align uintptr, watermark *uintptr) (unsafe.Pointer, error) {
		ptr, err := arena(a, size, align, watermark)
		if err != nil {
			return ptr, err
		}

		if a == ACTION_RESET {
			pinner.Unpin()
		} else {
			pinner.Pin(ptr)
		}

		return ptr, err
	}
}