reaper.go

// Package reaper implements a deterministic, bounded rewrite loop over a
// priority structure.
//
// A Reaper repeatedly removes a single element from a heap, allows a callback
// to inspect it and optionally emit replacement elements, and then reinserts
// those elements back into the heap. Each step is strictly bounded: a callback
// may emit at most a fixed number of elements ("degree") per visit.
//
// The design enforces three core invariants:
//
//  1. Only one element is popped from the heap at a time.
//  2. Reinsertion is explicitly controlled and bounded.
//  3. Callbacks are invoked without holding heap locks.
//
// This makes Reaper suitable for schedulers, rewrite systems, planners,
// or any algorithm that requires controlled feedback into a priority queue.
package reaper

import (
	"errors"
	"sync"
)

// Heap is a thread-safe priority structure supporting push and pop operations.
//
// Reaper assumes that all heap mutations are protected by the provided Lock
// and Unlock methods. Callbacks are always invoked without holding the heap lock.
type Heap[T any] interface {
	Push(elem T)
	Pop() T
	Empty() bool
	Lock()
	Unlock()
}

// Callback is invoked for each element popped from the heap during a Reap
// operation.
//
// The callback receives the popped element ("front") and an emit function.
// The emit function may be used to buffer elements for reinsertion into the
// heap after the callback returns.
//
// Returning stop == true signals that reaping should terminate early.
// In that case, the popped element is reinserted into the heap unchanged,
// and no buffered emissions are committed.
type Callback[T any] interface {
	Visit(front T, emit func(...T) error) (stop bool)
}

// CallbackFunc allows a function to be used as a Callback.
type CallbackFunc[T any] func(front T, emit func(...T) error) (stop bool)

func (f CallbackFunc[T]) Visit(front T, emit func(...T) error) (stop bool) {
	return f(front, emit)
}

// Reaper administrates the controlled feedback loop over a Heap.
type Reaper[T any] struct {
	degree int
	pool   sync.Pool
}

// New creates a new Reaper with the specified degree.
//
// Degree specifies the maximum number of elements that can be emitted per
// Visit call. A degree of zero disables emission entirely, turning the Reaper
// into a pure consumer.
func New[T any](degree int) *Reaper[T] {
	return &Reaper[T]{
		degree: degree,
		pool: sync.Pool{
			New: func() any {
				s := make([]T, 0, degree)
				return &s
			},
		},
	}
}

// ErrEmitBufferOverflow is returned by the emit function when the number of
// emitted elements in a single Visit exceeds the configured degree.
var ErrEmitBufferOverflow = errors.New("emit buffer overflow")

// Result indicates the outcome of a Reap operation.
//
//go:generate stringer -type=Result
type Result int

const (
	// Exhausted indicates that the heap became empty and no further
	// elements were available for processing.
	Exhausted Result = iota

	// Stopped indicates that the callback requested early termination.
	// In this case, the last popped element was reinserted into the heap.
	Stopped
)

// Reap executes the rewrite loop over the provided heap using the given
// callback.
//
// Reap repeatedly pops a single element from the heap and invokes the callback
// on it. If the callback emits elements, they are reinserted atomically after
// the callback returns. If the callback requests terminations, the popped
// element is restored, the emitted elments are discarded, and Reap returns
// Stopped.
//
// Reap returns Exhausted if the heap becomes empty.
func (r *Reaper[T]) Reap(h Heap[T], cb Callback[T]) Result {
	buf := r.pool.Get().(*[]T)

	defer func() {
		*buf = (*buf)[:0]
		r.pool.Put(buf)
	}()

	emit := func(elems ...T) error {
		if len(*buf)+len(elems) > r.degree {
			return ErrEmitBufferOverflow
		}

		*buf = append(*buf, elems...)
		return nil
	}

	for {
		var front T

		h.Lock()
		if h.Empty() {
			h.Unlock()
			return Exhausted
		}
		front = h.Pop()
		h.Unlock()

		stop := cb.Visit(front, emit)
		if stop {
			h.Lock()
			h.Push(front)
			h.Unlock()
			return Stopped
		}

		h.Lock()
		for _, elem := range *buf {
			h.Push(elem)
		}
		h.Unlock()
		*buf = (*buf)[:0]
	}
}