github.com/consensys/gnark-crypto@v0.14.0/internal/generator/polynomial/template/pool.go.tmpl

{{ $sham := eq .ElementType "small_rational.SmallRational"}}
import (
"{{.FieldPackagePath}}"
{{- if not $sham}}
	"encoding/json"
	"fmt"
	"runtime"
	"sort"
	"sync"
	"unsafe"
{{- end}}
)

{{ if $sham}}
// Do as little as possible to instantiate the interface
type Pool struct {
}

func NewPool(...int) (pool Pool) {
	return Pool{}
}

func (p *Pool) Make(n int) []{{.ElementType}} {
	return make([]{{.ElementType}}, n)
}

func (p *Pool) Dump(...[]{{.ElementType}}) {
}

func (p *Pool) Clone(slice []{{.ElementType}}) []{{.ElementType}} {
	res := p.Make(len(slice))
	copy(res, slice)
	return res
}
{{ else}}
// Memory management for polynomials
// WARNING: This is not thread safe TODO: Make sure that is not a problem
// TODO: There is a lot of "unsafe" memory management here and needs to be vetted thoroughly

type sizedPool struct {
	maxN  int
	pool  sync.Pool
	stats poolStats
}

type inUseData struct {
	allocatedFor []uintptr
	pool         *sizedPool
}

type Pool struct {
	//lock     sync.Mutex
	inUse    sync.Map
	subPools []sizedPool
}

func (p *sizedPool) get(n int) *{{.ElementType}} {
	p.stats.make(n)
	return p.pool.Get().(*{{.ElementType}})
}

func (p *sizedPool) put(ptr *{{.ElementType}}) {
	p.stats.dump()
	p.pool.Put(ptr)
}

func NewPool(maxN ...int) (pool Pool) {

	sort.Ints(maxN)
	pool = Pool{
		subPools: make([]sizedPool, len(maxN)),
	}

	for i := range pool.subPools {
		subPool := &pool.subPools[i]
		subPool.maxN = maxN[i]
		subPool.pool = sync.Pool{
			New: func() interface{} {
				subPool.stats.Allocated++
				return getDataPointer(make([]{{.ElementType}}, 0, subPool.maxN))
			},
		}
	}
	return
}

func (p *Pool) findCorrespondingPool(n int) *sizedPool {
	poolI := 0
	for poolI < len(p.subPools) && n > p.subPools[poolI].maxN {
		poolI++
	}
	return &p.subPools[poolI] // out of bounds error here would mean that n is too large
}

func (p *Pool) Make(n int) []{{.ElementType}} {
	pool := p.findCorrespondingPool(n)
	ptr := pool.get(n)
	p.addInUse(ptr, pool)
	return unsafe.Slice(ptr,n)
}

// Dump dumps a set of polynomials into the pool
func (p *Pool) Dump(slices ...[]{{.ElementType}}) {
	for _, slice := range slices {
		ptr := getDataPointer(slice)
		if metadata, ok := p.inUse.Load(ptr); ok {
			p.inUse.Delete(ptr)
			metadata.(inUseData).pool.put(ptr)
		} else {
			panic("attempting to dump a slice not created by the pool")
		}
	}
}

func (p *Pool) addInUse(ptr *{{.ElementType}}, pool *sizedPool) {
	pcs := make([]uintptr, 2)
	n := runtime.Callers(3, pcs)

	if prevPcs, ok := p.inUse.Load(ptr); ok { // TODO: remove if unnecessary for security
		panic(fmt.Errorf("re-allocated non-dumped slice, previously allocated at %v", runtime.CallersFrames(prevPcs.(inUseData).allocatedFor)))
	}
	p.inUse.Store(ptr, inUseData{
		allocatedFor: pcs[:n],
		pool:         pool,
	})
}

func printFrame(frame runtime.Frame) {
	fmt.Printf("\t%s line %d, function %s\n", frame.File, frame.Line, frame.Function)
}

func (p *Pool) printInUse() {
	fmt.Println("slices never dumped allocated at:")
	p.inUse.Range(func(_, pcs any) bool {
		fmt.Println("-------------------------")

		var frame runtime.Frame
		frames := runtime.CallersFrames(pcs.(inUseData).allocatedFor)
		more := true
		for more {
			frame, more = frames.Next()
			printFrame(frame)
		}
		return true
	})
}

type poolStats struct {
	Used          int
	Allocated     int
	ReuseRate     float64
	InUse         int
	GreatestNUsed int
	SmallestNUsed int
}

type poolsStats struct {
	SubPools []poolStats
	InUse    int
}

func (s *poolStats) make(n int) {
	s.Used++
	s.InUse++
	if n > s.GreatestNUsed {
		s.GreatestNUsed = n
	}
	if s.SmallestNUsed == 0 || s.SmallestNUsed > n {
		s.SmallestNUsed = n
	}
}

func (s *poolStats) dump() {
	s.InUse--
}

func (s *poolStats) finalize() {
	s.ReuseRate = float64(s.Used) / float64(s.Allocated)
}

func getDataPointer(slice []{{.ElementType}}) *{{.ElementType}} {
	return (*{{.ElementType}})(unsafe.SliceData(slice))
}

func (p *Pool) PrintPoolStats() {
	InUse := 0
	subStats := make([]poolStats, len(p.subPools))
	for i := range p.subPools {
		subPool := &p.subPools[i]
		subPool.stats.finalize()
		subStats[i] = subPool.stats
		InUse += subPool.stats.InUse
	}

	stats := poolsStats{
		SubPools: subStats,
		InUse:    InUse,
	}
	serialized, _ := json.MarshalIndent(stats, "", "  ")
	fmt.Println(string(serialized))
	p.printInUse()
}

func (p *Pool) Clone(slice []{{.ElementType}}) []{{.ElementType}} {
	res := p.Make(len(slice))
	copy(res, slice)
	return res
}
{{end}}