Prover/feat/small field exploratory—Smartvectors that support both ba…

…se field elements and field extensions. (#469)

* degree-2 extension poc

* refactor

* refactor

* templating smartvectors

* smartvectors wip

* helper functions, wip

* utilfunctions wip

* remove templating from RandVect

* REVERT templating in smartvectors

* api changes in existing smartvector types

* filling gaps in the base API

* intermediary Get function

* bug fixes

* mempool for fext

* fext vectors, wip

* code for ext types of smartvector

* bug fix in isZero

* arithmetic_test for extended smartvectors, wip

* rotated ext test

* window ext test

* setInt and setUint touch only one coordinate

* removed field amd files

* another intemediary function

* linter errors, wip

* prunning fext, attempting to fix linter errors

* cosmetic change

* Update vectorext.go

* changing build flags

* build flags + removed fieldElement interface

---------

Co-authored-by: AlexandreBelling <[email protected]>

prover/maths/common/mempoolext/debug_pool.go

@@ -0,0 +1,146 @@
+package mempoolext
+
+import (
+	"errors"
+	"fmt"
+	"github.com/consensys/linea-monorepo/prover/maths/field/fext"
+	"runtime"
+	"strconv"
+	"unsafe"
+
+	"github.com/consensys/linea-monorepo/prover/utils"
+)
+
+type DebuggeableCall struct {
+	Parent MemPool
+	Logs   map[uintptr]*[]Record
+}
+
+func NewDebugPool(p MemPool) *DebuggeableCall {
+	return &DebuggeableCall{
+		Parent: p,
+		Logs:   make(map[uintptr]*[]Record),
+	}
+}
+
+type Record struct {
+	Where string
+	What  recordType
+}
+
+func (m *DebuggeableCall) Prewarm(nbPrewarm int) MemPool {
+	m.Parent.Prewarm(nbPrewarm)
+	return m
+}
+
+type recordType string
+
+const (
+	AllocRecord recordType = "alloc"
+	FreeRecord  recordType = "free"
+)
+
+func (m *DebuggeableCall) Alloc() *[]fext.Element {
+
+	var (
+		v                = m.Parent.Alloc()
+		uptr             = uintptr(unsafe.Pointer(v))
+		logs             *[]Record
+		_, file, line, _ = runtime.Caller(2)
+	)
+
+	logs, found := m.Logs[uptr]
+
+	if !found {
+		logs = &[]Record{}
+		m.Logs[uptr] = logs
+	}
+
+	*logs = append(*logs, Record{
+		Where: file + ":" + strconv.Itoa(line),
+		What:  AllocRecord,
+	})
+
+	return v
+}
+
+func (m *DebuggeableCall) Free(v *[]fext.Element) error {
+
+	var (
+		uptr             = uintptr(unsafe.Pointer(v))
+		logs             *[]Record
+		_, file, line, _ = runtime.Caller(2)
+	)
+
+	logs, found := m.Logs[uptr]
+
+	if !found {
+		logs = &[]Record{}
+		m.Logs[uptr] = logs
+	}
+
+	*logs = append(*logs, Record{
+		Where: file + ":" + strconv.Itoa(line),
+		What:  FreeRecord,
+	})
+
+	return m.Parent.Free(v)
+}
+
+func (m *DebuggeableCall) Size() int {
+	return m.Parent.Size()
+}
+
+func (m *DebuggeableCall) TearDown() {
+	if p, ok := m.Parent.(*SliceArena); ok {
+		p.TearDown()
+	}
+}
+
+func (m *DebuggeableCall) Errors() error {
+
+	var err error
+
+	for _, logs_ := range m.Logs {
+
+		if logs_ == nil || len(*logs_) == 0 {
+			utils.Panic("got a nil entry")
+		}
+
+		logs := *logs_
+
+		for i := range logs {
+			if i == 0 && logs[i].What == FreeRecord {
+				err = errors.Join(err, fmt.Errorf("freed a vector that was not from the pool: where=%v", logs[i].Where))
+			}
+
+			if i == len(logs)-1 && logs[i].What == AllocRecord {
+				err = errors.Join(err, fmt.Errorf("leaked a vector out of the pool: where=%v", logs[i].Where))
+			}
+
+			if i == 0 {
+				continue
+			}
+
+			if logs[i-1].What == AllocRecord && logs[i].What == AllocRecord {
+				wheres := []string{logs[i-1].Where, logs[i].Where}
+				for k := i + 1; k < len(logs) && logs[k].What == AllocRecord; k++ {
+					wheres = append(wheres, logs[k].Where)
+				}
+
+				err = errors.Join(err, fmt.Errorf("vector was allocated multiple times concurrently where=%v", wheres))
+			}
+
+			if logs[i-1].What == FreeRecord && logs[i].What == FreeRecord {
+				wheres := []string{logs[i-1].Where, logs[i].Where}
+				for k := i + 1; k < len(logs) && logs[k].What == FreeRecord; k++ {
+					wheres = append(wheres, logs[k].Where)
+				}
+
+				err = errors.Join(err, fmt.Errorf("vector was freed multiple times concurrently where=%v", wheres))
+			}
+		}
+	}
+
+	return err
+}

prover/maths/common/mempoolext/debug_pool_test.go

@@ -0,0 +1,51 @@
+package mempoolext
+
+import (
+	"github.com/consensys/linea-monorepo/prover/maths/field/fext"
+	"github.com/stretchr/testify/assert"
+	"strings"
+	"testing"
+)
+
+func TestDebugPool(t *testing.T) {
+
+	t.Run("leak-detection", func(t *testing.T) {
+
+		pool := NewDebugPool(CreateFromSyncPool(32))
+
+		for i := 0; i < 16; i++ {
+			func() {
+				_ = pool.Alloc()
+			}()
+		}
+
+		err := pool.Errors().Error()
+		assert.True(t, strings.HasPrefix(err, "leaked a vector out of the pool"))
+	})
+
+	t.Run("double-free", func(t *testing.T) {
+
+		pool := NewDebugPool(CreateFromSyncPool(32))
+
+		v := pool.Alloc()
+
+		for i := 0; i < 16; i++ {
+			pool.Free(v)
+		}
+
+		err := pool.Errors().Error()
+		assert.Truef(t, strings.HasPrefix(err, "vector was freed multiple times concurrently"), err)
+	})
+
+	t.Run("foreign-free", func(t *testing.T) {
+
+		pool := NewDebugPool(CreateFromSyncPool(32))
+
+		v := make([]fext.Element, 32)
+		pool.Free(&v)
+
+		err := pool.Errors().Error()
+		assert.Truef(t, strings.HasPrefix(err, "freed a vector that was not from the pool"), err)
+	})
+
+}

prover/maths/common/mempoolext/from_sync_pool.go

@@ -0,0 +1,72 @@
+package mempoolext
+
+import (
+	"github.com/consensys/linea-monorepo/prover/maths/field/fext"
+	"github.com/consensys/linea-monorepo/prover/utils"
+	"github.com/consensys/linea-monorepo/prover/utils/parallel"
+	"sync"
+)
+
+// FromSyncPool pools the allocation for slices of [fext.Element] of size `Size`.
+// It should be used with great caution and every slice allocated via this pool
+// must be manually freed and only once.
+//
+// FromSyncPool is used to reduce the number of allocation which can be significant
+// when doing operations over field elements.
+type FromSyncPool struct {
+	size int
+	P    sync.Pool
+}
+
+// CreateFromSyncPool initializes the Pool with the given number of elements in it.
+func CreateFromSyncPool(size int) *FromSyncPool {
+	// Initializes the pool
+	return &FromSyncPool{
+		size: size,
+		P: sync.Pool{
+			New: func() any {
+				res := make([]fext.Element, size)
+				return &res
+			},
+		},
+	}
+}
+
+// Prewarm the Pool by preallocating `nbPrewarm` in it.
+func (p *FromSyncPool) Prewarm(nbPrewarm int) MemPool {
+	prewarmed := make([]fext.Element, p.size*nbPrewarm)
+	parallel.Execute(nbPrewarm, func(start, stop int) {
+		for i := start; i < stop; i++ {
+			vec := prewarmed[i*p.size : (i+1)*p.size]
+			p.P.Put(&vec)
+		}
+	})
+	return p
+}
+
+// Alloc returns a vector allocated from the pool. Vector allocated via the
+// pool should ideally be returned to the pool. If not, they are still going to
+// be picked up by the GC.
+func (p *FromSyncPool) Alloc() *[]fext.Element {
+	res := p.P.Get().(*[]fext.Element)
+	return res
+}
+
+// Free returns an object to the pool. It must never be called twice over
+// the same object or undefined behaviours are going to arise. It is fine to
+// pass objects allocated to outside of the pool as long as they have the right
+// dimension.
+func (p *FromSyncPool) Free(vec *[]fext.Element) error {
+	// Check the vector has the right size
+	if len(*vec) != p.size {
+		utils.Panic("expected size %v, expected %v", len(*vec), p.Size())
+	}
+
+	p.P.Put(vec)
+
+	return nil
+}
+
+func (p *FromSyncPool) Size() int {
+	return p.size
+}

prover/maths/common/mempoolext/mempool.go

@@ -0,0 +1,60 @@
+package mempoolext
+
+import (
+	"github.com/consensys/linea-monorepo/prover/maths/field/fext"
+	"github.com/consensys/linea-monorepo/prover/utils"
+)
+
+type MemPool interface {
+	Prewarm(nbPrewarm int) MemPool
+	Alloc() *[]fext.Element
+	Free(vec *[]fext.Element) error
+	Size() int
+}
+
+// ExtractCheckOptionalStrict returns
+//   - p[0], true if the expectedSize matches the one of the provided pool
+//   - nil, false if no `p` is provided
+//   - panic if the assigned size of the pool does not match
+//   - panic if the caller provides `nil` as argument for `p`
+//
+// This is used to unwrap a [FromSyncPool] that is commonly passed to functions as an
+// optional variadic parameter and at the same time validating that the pool
+// object has the right size.
+func ExtractCheckOptionalStrict(expectedSize int, p ...MemPool) (pool MemPool, ok bool) {
+	// Checks if there is a pool
+	hasPool := len(p) > 0 && p[0] != nil
+	if hasPool {
+		pool = p[0]
+	}
+
+	// Sanity-check that the size of the pool is actually what we expected
+	if hasPool && pool.Size() != expectedSize {
+		utils.Panic("pooled vector size are %v, but required %v", pool.Size(), expectedSize)
+	}
+
+	return pool, hasPool
+}
+
+// ExtractCheckOptionalSoft returns
+//   - p[0], true if the expectedSize matches the one of the provided pool
+//   - nil, false if no `p` is provided
+//   - nil, false if the length of the vector does not match the one of the pool
+//   - panic if the caller provides `nil` as argument for `p`
+//
+// This is used to unwrap a [FromSyncPool] that is commonly passed to functions as an
+// optional variadic parameter.
+func ExtractCheckOptionalSoft(expectedSize int, p ...MemPool) (pool MemPool, ok bool) {
+	// Checks if there is a pool
+	hasPool := len(p) > 0
+	if hasPool {
+		pool = p[0]
+	}
+
+	// Sanity-check that the size of the pool is actually what we expected
+	if hasPool && pool.Size() != expectedSize {
+		return nil, false
+	}
+
+	return pool, hasPool
+}

prover/maths/common/mempoolext/slice_arena.go

@@ -0,0 +1,51 @@
+package mempoolext
+
+import (
+	"github.com/consensys/linea-monorepo/prover/maths/field/fext"
+)
+
+// SliceArena is a simple not-threadsafe arena implementation that uses a
+// mempool to carry its allocation. It will only put back free memory in the
+// the parent pool when TearDown is called.
+type SliceArena struct {
+	frees  []*[]fext.Element
+	parent MemPool
+}
+
+func WrapsWithMemCache(pool MemPool) *SliceArena {
+	return &SliceArena{
+		frees:  make([]*[]fext.Element, 0, 1<<7),
+		parent: pool,
+	}
+}
+
+func (m *SliceArena) Prewarm(nbPrewarm int) MemPool {
+	m.parent.Prewarm(nbPrewarm)
+	return m
+}
+
+func (m *SliceArena) Alloc() *[]fext.Element {
+
+	if len(m.frees) == 0 {
+		return m.parent.Alloc()
+	}
+
+	last := m.frees[len(m.frees)-1]
+	m.frees = m.frees[:len(m.frees)-1]
+	return last
+}
+
+func (m *SliceArena) Free(v *[]fext.Element) error {
+	m.frees = append(m.frees, v)
+	return nil
+}
+
+func (m *SliceArena) Size() int {
+	return m.parent.Size()
+}
+
+func (m *SliceArena) TearDown() {
+	for i := range m.frees {
+		m.parent.Free(m.frees[i])
+	}
+}