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

{{ $G1TAffine := print (toUpper .G1.PointName) "Affine" }}
{{ $G1TJacobian := print (toUpper .G1.PointName) "Jac" }}
{{ $G1TJacobianExtended := print (toLower .G1.PointName) "JacExtended" }}

{{ $G2TAffine := print (toUpper .G2.PointName) "Affine" }}
{{ $G2TJacobian := print (toUpper .G2.PointName) "Jac" }}
{{ $G2TJacobianExtended := print (toLower .G2.PointName) "JacExtended" }}


import (
	"github.com/consensys/gnark-crypto/ecc/{{.Name}}/fp"
	{{- if and (ne .G1.CoordType .G2.CoordType) (ne .Name "secp256k1") }}
	"github.com/consensys/gnark-crypto/ecc/{{.Name}}/internal/fptower"
	{{- end}}
)

{{ template "multiexp" dict "CoordType" .G1.CoordType "PointName" .G1.PointName "UPointName" (toUpper .G1.PointName) "TAffine" $G1TAffine "TJacobian" $G1TJacobian "TJacobianExtended" $G1TJacobianExtended "FrNbWords" .Fr.NbWords "CRange" .G1.CRange}}
{{- if ne .Name "secp256k1"}}
{{ template "multiexp" dict "CoordType" .G2.CoordType "PointName" .G2.PointName "UPointName" (toUpper .G2.PointName) "TAffine" $G2TAffine "TJacobian" $G2TJacobian "TJacobianExtended" $G2TJacobianExtended "FrNbWords" .Fr.NbWords "CRange" .G2.CRange}}
{{- end}}


{{define "multiexp" }}

type batchOp{{ $.TAffine }} struct {
	bucketID uint16
	point {{ $.TAffine }}
}

// processChunk{{ $.UPointName }}BatchAffine process a chunk of the scalars during the msm
// using affine coordinates for the buckets. To amortize the cost of the inverse in the affine addition
// we use a batch affine addition.
//
// this is derived from a PR by 0x0ece : https://github.com/ConsenSys/gnark-crypto/pull/249
// See Section 5.3: ia.cr/2022/1396
func processChunk{{ $.UPointName }}BatchAffine[BJE ib{{ $.TJacobianExtended }},B ib{{ $.TAffine }}, BS bitSet, TP p{{ $.TAffine }}, TPP pp{{ $.TAffine }}, TQ qOps{{ $.TAffine }}, TC c{{ $.TAffine}}](
	 chunk uint64,
	 chRes chan<- {{ $.TJacobianExtended }},
	 c uint64,
	 points []{{ $.TAffine }},
	 digits []uint16,
	 sem chan struct{}) {

	if sem != nil {
		// if we are limited, wait for a token in the semaphore
		<-sem
	}

	// the batch affine addition needs independent points; in other words, for a window of batchSize
	// we want to hit independent bucketIDs when processing the digit. if there is a conflict (we're trying
	// to add 2 different points to the same bucket), then we push the conflicted point to a queue.
	// each time the batch is full, we execute it, and tentatively put the points (if not conflict)
	// from the top of the queue into the next batch.
	// if the queue is full, we "flush it"; we sequentially add the points to the buckets in
	// {{ $.TJacobianExtended }} coordinates.
	// The reasoning behind this is the following; batchSize is chosen such as, for a uniformly random
	// input, the number of conflicts is going to be low, and the element added to the queue should be immediately
	// processed in the next batch. If it's not the case, then our inputs are not random; and we fallback to
	// non-batch-affine version.

	// note that we have 2 sets of buckets
	// 1 in {{ $.TAffine }} used with the batch affine additions
	// 1 in {{ $.TJacobianExtended }} used in case the queue of conflicting points
	var buckets B // in {{ $.TAffine }} coordinates, infinity point is represented as (0,0), no need to init
	var bucketsJE BJE
	for i := 0; i < len(buckets); i++ {
		bucketsJE[i].setInfinity()
	}

	// setup for the batch affine;
	var (
		bucketIds BS // bitSet to signify presence of a bucket in current batch
		cptAdd int // count the number of bucket + point added to current batch
		R TPP // bucket references
		P TP // points to be added to R (buckets); it is beneficial to store them on the stack (ie copy)
		queue TQ // queue of points that conflict the current batch
		qID int // current position in queue
	)

	batchSize := len(P)

	isFull := func() bool {	return cptAdd == batchSize}

	executeAndReset := func ()  {
		batchAdd{{ $.TAffine }}[TP, TPP, TC](&R, &P, cptAdd)
		var tmp BS
		bucketIds = tmp
		cptAdd = 0
	}

	addFromQueue := func(op batchOp{{ $.TAffine }}) {
		// @precondition: must ensures bucket is not "used" in current batch
		// note that there is a bit of duplicate logic between add and addFromQueue
		// the reason is that as of Go 1.19.3, if we pass a pointer to the queue item (see add signature)
		// the compiler will put the queue on the heap.
		BK := &buckets[op.bucketID]

		// handle special cases with inf or -P / P
		if BK.IsInfinity()  {
			BK.Set(&op.point)
			return
		}
		if BK.X.Equal(&op.point.X) {
			if BK.Y.Equal(&op.point.Y) {
				// P + P: doubling, which should be quite rare --
				// we use the other set of buckets
				bucketsJE[op.bucketID].addMixed(&op.point)
				return
			}
			BK.setInfinity()
			return
		}

		bucketIds[op.bucketID] = true
		R[cptAdd] = BK
		P[cptAdd] = op.point
		cptAdd++
	}

	add := func(bucketID uint16, PP *{{$.TAffine}}, isAdd bool) {
		// @precondition: ensures bucket is not "used" in current batch
		BK := &buckets[bucketID]
		// handle special cases with inf or -P / P
		if BK.IsInfinity()  {
			if isAdd {
				BK.Set(PP)
			} else {
				BK.Neg(PP)
			}
			return
		}
		if BK.X.Equal(&PP.X) {
			if BK.Y.Equal(&PP.Y) {
				// P + P: doubling, which should be quite rare --
				if isAdd {
					bucketsJE[bucketID].addMixed(PP)
				} else {
					BK.setInfinity()
				}
				return
			}
			if isAdd {
				BK.setInfinity()
			} else {
				bucketsJE[bucketID].subMixed(PP)
			}
			return
		}

		bucketIds[bucketID] = true
		R[cptAdd] = BK
		if isAdd {
			P[cptAdd].Set(PP)
		} else {
			P[cptAdd].Neg(PP)
		}
		cptAdd++
	}

	flushQueue := func () {
		for i:=0; i < qID; i++ {
			bucketsJE[queue[i].bucketID].addMixed(&queue[i].point)
		}
		qID = 0
	}

	processTopQueue := func () {
		for i := qID - 1; i >= 0; i-- {
			if bucketIds[queue[i].bucketID] {
				return
			}
			addFromQueue(queue[i])
			// len(queue) < batchSize so no need to check for full batch.
			qID--
		}
	}


	for i, digit := range digits {

		if digit == 0 || points[i].IsInfinity() {
			continue
		}

		bucketID := uint16((digit>>1))
		isAdd := digit&1 == 0
		if isAdd {
			// add
			bucketID-=1
		}

		if bucketIds[bucketID] {
			// put it in queue
			queue[qID].bucketID = bucketID
			if isAdd {
				queue[qID].point.Set(&points[i])
			} else {
				queue[qID].point.Neg(&points[i])
			}
			qID++

			// queue is full, flush it.
			if qID == len(queue) - 1 {
				flushQueue()
			}
			continue
		}

		// we add the point to the batch.
		add(bucketID, &points[i], isAdd)
		if isFull() {
			executeAndReset()
			processTopQueue()
		}
	}


	// flush items in batch.
	executeAndReset()

	// empty the queue
	flushQueue()


	// reduce buckets into total
	// total =  bucket[0] + 2*bucket[1] + 3*bucket[2] ... + n*bucket[n-1]
	var runningSum, total {{ $.TJacobianExtended }}
	runningSum.setInfinity()
	total.setInfinity()
	for k := len(buckets) - 1; k >= 0; k-- {
		runningSum.addMixed(&buckets[k])
		if !bucketsJE[k].IsInfinity() {
			runningSum.add(&bucketsJE[k])
		}
		total.add(&runningSum)
	}


	if sem != nil {
		// release a token to the semaphore
		// before sending to chRes
		sem <- struct{}{}
	}

	chRes <- total

}

// we declare the buckets as fixed-size array types
// this allow us to allocate the buckets on the stack
{{- range $c :=  $.CRange}}
{{- if gt $c 9}}
type bucket{{ $.TAffine }}C{{$c}} [{{nbBuckets $c}}]{{ $.TAffine }}
{{- end}}
{{- end}}


// buckets: array of {{ $.TAffine }} points of size 1 << (c-1)
type ib{{ $.TAffine }} interface {
	{{- range $i, $c :=  $.CRange}}
	{{- if gt $c 9}}
	bucket{{ $.TAffine }}C{{$c}} {{- if not (last $i $.CRange)}} | {{- end}}
	{{- end}}
	{{- end}}
}

// array of coordinates {{ $.CoordType }}
type c{{ $.TAffine }} interface {
	{{- range $i, $c :=  $.CRange}}
	{{- if gt $c 9}}
	c{{ $.TAffine }}C{{$c}} {{- if not (last $i $.CRange)}} | {{- end}}
	{{- end}}
	{{- end}}
}

// buckets: array of {{ $.TAffine }} points (for the batch addition)
type p{{ $.TAffine }} interface {
	{{- range $i, $c :=  $.CRange}}
	{{- if gt $c 9}}
	p{{ $.TAffine }}C{{$c}} {{- if not (last $i $.CRange)}} | {{- end}}
	{{- end}}
	{{- end}}
}

// buckets: array of *{{ $.TAffine }} points (for the batch addition)
type pp{{ $.TAffine }} interface {
	{{- range $i, $c :=  $.CRange}}
	{{- if gt $c 9}}
	pp{{ $.TAffine }}C{{$c}} {{- if not (last $i $.CRange)}} | {{- end}}
	{{- end}}
	{{- end}}
}

// buckets: array of {{ $.TAffine }} queue operations (for the batch addition)
type qOps{{ $.TAffine }} interface {
	{{- range $i, $c :=  $.CRange}}
	{{- if gt $c 9}}
	q{{ $.TAffine }}C{{$c}} {{- if not (last $i $.CRange)}} | {{- end}}
	{{- end}}
	{{- end}}
}


{{- range $c :=  $.CRange}}
{{if gt $c 9}}
// batch size {{batchSize $c}} when c = {{$c}}
type c{{ $.TAffine }}C{{$c}} [{{batchSize $c}}]{{ $.CoordType }}
type p{{ $.TAffine }}C{{$c}} [{{batchSize $c}}]{{ $.TAffine }}
type pp{{ $.TAffine }}C{{$c}} [{{batchSize $c}}]*{{ $.TAffine }}
type q{{ $.TAffine }}C{{$c}} [{{batchSize $c}}]batchOp{{ $.TAffine }}
{{- end}}
{{- end}}


{{end }}

{{- range $c :=  $.G1.CRange}}
type bitSetC{{$c}} [{{nbBuckets $c}}]bool
{{- end}}

type bitSet interface {
	{{- range $i, $c :=  $.G1.CRange}}
	bitSetC{{$c}} {{- if not (last $i $.G1.CRange)}} | {{- end}}
	{{- end}}
}