Support FrMapPool

This uses a map to pool `fr.Element` instances, long with appropriate
locks to ensure thread-safety.

pkg/util/field_array.go

@@ -47,7 +47,9 @@ func NewFrArray(height uint, bitWidth uint) FrArray {
 		var pool FrIndexPool[uint16] = NewFrIndexPool[uint16]()
 		return NewFrPoolArray[uint16](height, bitWidth, pool)
 	default:
-		return NewFrElementArray(height, bitWidth)
+		// return NewFrElementArray(height, bitWidth)
+		var pool FrMapPool = NewFrMapPool(bitWidth)
+		return NewFrPoolArray[uint32](height, bitWidth, pool)
 	}
 }
 
@@ -75,14 +77,14 @@ func FrArrayFromBigInts(bitWidth uint, ints []*big.Int) FrArray {
 // etc.
 type FrElementArray struct {
 	// The data stored in this column (as bytes).
-	elements []*fr.Element
+	elements []fr.Element
 	// Maximum number of bits required to store an element of this array.
 	bitwidth uint
 }
 
 // NewFrElementArray constructs a new field array with a given capacity.
 func NewFrElementArray(height uint, bitwidth uint) *FrElementArray {
-	elements := make([]*fr.Element, height)
+	elements := make([]fr.Element, height)
 	return &FrElementArray{elements, bitwidth}
 }
 
@@ -98,19 +100,19 @@ func (p *FrElementArray) BitWidth() uint {
 
 // Get returns the field element at the given index in this array.
 func (p *FrElementArray) Get(index uint) *fr.Element {
-	return p.elements[index]
+	return &p.elements[index]
 }
 
 // Set sets the field element at the given index in this array, overwriting the
 // original value.
 func (p *FrElementArray) Set(index uint, element *fr.Element) {
-	p.elements[index] = element
+	p.elements[index] = *element
 }
 
 // Clone makes clones of this array producing an otherwise identical copy.
 func (p *FrElementArray) Clone() Array[*fr.Element] {
 	// Allocate sufficient memory
-	ndata := make([]*fr.Element, uint(len(p.elements)))
+	ndata := make([]fr.Element, uint(len(p.elements)))
 	// Copy over the data
 	copy(ndata, p.elements)
 	//
@@ -120,12 +122,12 @@ func (p *FrElementArray) Clone() Array[*fr.Element] {
 // PadFront (i.e. insert at the beginning) this array with n copies of the given padding value.
 func (p *FrElementArray) PadFront(n uint, padding *fr.Element) Array[*fr.Element] {
 	// Allocate sufficient memory
-	ndata := make([]*fr.Element, uint(len(p.elements))+n)
+	ndata := make([]fr.Element, uint(len(p.elements))+n)
 	// Copy over the data
 	copy(ndata[n:], p.elements)
 	// Go padding!
 	for i := uint(0); i < n; i++ {
-		ndata[i] = padding
+		ndata[i] = *padding
 	}
 	// Copy over
 	return &FrElementArray{ndata, p.bitwidth}
@@ -190,6 +192,7 @@ func (p *FrPoolArray[K, P]) Set(index uint, element *fr.Element) {
 }
 
 // Clone makes clones of this array producing an otherwise identical copy.
+// nolint: revive
 func (p *FrPoolArray[K, P]) Clone() Array[*fr.Element] {
 	// Allocate sufficient memory
 	ndata := make([]K, len(p.elements))

pkg/util/field_pool.go

@@ -16,10 +16,70 @@ type FrPool[K any] interface {
 	Get(K) *fr.Element
 }
 
-// FrSmallPoolIndex provides a simple charactirisation of types which can be
-// used to lookup into indexed pools.
-type FrSmallPoolIndex interface {
-	uint8 | uint16
+// ----------------------------------------------------------------------------
+
+// FrMapPool is a pool implementation indexed using a map and backed using a
+// dynamically sized array.
+type FrMapPool struct {
+	bitwidth uint
+}
+
+// NewFrMapPool constructs a new pool which uses map to index items.
+func NewFrMapPool(bitwidth uint) FrMapPool {
+	initMapPool()
+	return FrMapPool{bitwidth}
+}
+
+// Get looks up the given item in the pool.
+func (p FrMapPool) Get(index uint32) *fr.Element {
+	poolMapLock.RLock()
+	item := &poolMapArray[index]
+	poolMapLock.RUnlock()
+
+	return item
+}
+
+// Put allocates an item into the pool, returning its index.
+func (p FrMapPool) Put(element *fr.Element) uint32 {
+	// Lock items
+	poolMapLock.Lock()
+	index, ok := poolMapIndex[*element]
+	//
+	if !ok {
+		len := uint32(len(poolMapArray))
+		index = poolMapSize
+		// Update index
+		poolMapIndex[*element] = index
+		//
+		if index == len {
+			// capacity reached, so double it.
+			tmp := make([]fr.Element, len*5)
+			copy(tmp, poolMapArray)
+			poolMapArray = tmp
+		}
+		//
+		poolMapArray[index] = *element
+		poolMapSize++
+	}
+	// Done
+	poolMapLock.Unlock()
+	//
+	return index
+}
+
+var poolMapLock sync.RWMutex
+var poolMapIndex map[[4]uint64]uint32
+var poolMapArray []fr.Element
+var poolMapSize uint32
+
+func initMapPool() {
+	poolMapLock.Lock()
+	if poolMapIndex == nil {
+		// Initial capacity for 500 elements
+		poolMapIndex = make(map[[4]uint64]uint32, 1000)
+		poolMapArray = make([]fr.Element, 1000)
+	}
+	poolMapLock.Unlock()
 }
 
 // ----------------------------------------------------------------------------
@@ -31,7 +91,7 @@ type FrBitPool struct{}
 
 // NewFrBitPool constructs a new pool which uses a single bit for indexing.
 func NewFrBitPool() FrBitPool {
-	initPool()
+	initPool16()
 	//
 	return FrBitPool{}
 }
@@ -63,12 +123,12 @@ func (p FrBitPool) Put(element *fr.Element) bool {
 
 // FrIndexPool is a pool implementation which is backed by an array of a fixed
 // size.
-type FrIndexPool[K FrSmallPoolIndex] struct{}
+type FrIndexPool[K uint8 | uint16] struct{}
 
 // NewFrIndexPool constructs a new pool which uses a given key type for
 // indexing.
-func NewFrIndexPool[K FrSmallPoolIndex]() FrIndexPool[K] {
-	initPool()
+func NewFrIndexPool[K uint8 | uint16]() FrIndexPool[K] {
+	initPool16()
 	//
 	return FrIndexPool[K]{}
 }
@@ -96,7 +156,7 @@ var pool16init sync.Once
 var pool16bit []fr.Element
 
 // Initialise the index pool.
-func initPool() {
+func initPool16() {
 	// Singleton pattern for initialisation.
 	pool16init.Do(func() {
 		// Construct empty array