1// This package implements efficient Golidlocks arithmetic operations within Gnark. We do not use
  2// the emulated field arithmetic API, because it is too slow for our purposes. Instead, we use
  3// an efficient reduction method that leverages the fact that the modulus is a simple
  4// linear combination of powers of two.
  5package goldilocks
  6
  7// In general, methods whose name do not contain `NoReduce` can be used without any extra mental
  8// overhead. These methods act exactly as you would expect a normal field would operate.
  9//
 10// However, if you want to aggressively optimize the number of constraints in your circuit, it can
 11// be very beneficial to use the no reduction methods and keep track of the maximum number of bits
 12// your computation uses.
 13
 14// This implementation is based on the following plonky2 implementation of Goldilocks
 15// Available here: https://github.com/0xPolygonZero/plonky2/blob/main/field/src/goldilocks_field.rs#L70
 16
 17import (
 18	"fmt"
 19	"math"
 20	"math/big"
 21	"os"
 22	"strconv"
 23	"sync"
 24
 25	"github.com/consensys/gnark-crypto/field/goldilocks"
 26	"github.com/consensys/gnark/constraint/solver"
 27	"github.com/consensys/gnark/frontend"
 28	"github.com/consensys/gnark/std/math/emulated"
 29	"github.com/consensys/gnark/std/rangecheck"
 30)
 31
 32// The multiplicative group generator of the field.
 33var MULTIPLICATIVE_GROUP_GENERATOR goldilocks.Element = goldilocks.NewElement(7)
 34
 35// The two adicity of the field.
 36var TWO_ADICITY uint64 = 32
 37
 38// The power of two generator of the field.
 39var POWER_OF_TWO_GENERATOR goldilocks.Element = goldilocks.NewElement(1753635133440165772)
 40
 41// The modulus of the field.
 42var MODULUS *big.Int = emulated.Goldilocks{}.Modulus()
 43
 44// The number of bits to use for range checks on inner products of field elements.
 45// This MUST be a multiple of EXPECTED_OPTIMAL_BASEWIDTH if the commit based range checker is used.
 46// There is a bug in the pre 0.9.2 gnark range checker where it wouldn't appropriately range check a bitwidth that
 47// is misaligned from EXPECTED_OPTIMAL_BASEWIDTH:  https://github.com/Consensys/gnark/security/advisories/GHSA-rjjm-x32p-m3f7
 48var RANGE_CHECK_NB_BITS int = 144
 49
 50// The bit width size that the gnark commit based range checker should use.
 51var EXPECTED_OPTIMAL_BASEWIDTH int = 16
 52
 53// Registers the hint functions with the solver.
 54func init() {
 55	solver.RegisterHint(MulAddHint)
 56	solver.RegisterHint(ReduceHint)
 57	solver.RegisterHint(InverseHint)
 58	solver.RegisterHint(SplitLimbsHint)
 59}
 60
 61// A type alias used to represent Goldilocks field elements.
 62type Variable struct {
 63	Limb frontend.Variable
 64}
 65
 66// Creates a new Goldilocks field element from an existing variable. Assumes that the element is
 67// already reduced.
 68func NewVariable(x frontend.Variable) Variable {
 69	return Variable{Limb: x}
 70}
 71
 72// The zero element in the Golidlocks field.
 73func Zero() Variable {
 74	return NewVariable(0)
 75}
 76
 77// The one element in the Goldilocks field.
 78func One() Variable {
 79	return NewVariable(1)
 80}
 81
 82// The negative one element in the Goldilocks field.
 83func NegOne() Variable {
 84	return NewVariable(MODULUS.Uint64() - 1)
 85}
 86
 87type RangeCheckerType int
 88
 89const (
 90	NATIVE_RANGE_CHECKER RangeCheckerType = iota
 91	COMMIT_RANGE_CHECKER
 92	BIT_DECOMP_RANGE_CHECKER
 93)
 94
 95// The chip used for Goldilocks field operations.
 96type Chip struct {
 97	api frontend.API
 98
 99	rangeChecker     frontend.Rangechecker
100	rangeCheckerType RangeCheckerType
101
102	rangeCheckCollected []checkedVariable // These field are used if rangeCheckerType == commit_range_checker
103	collectedMutex      sync.Mutex
104}
105
106var (
107	poseidonChips = make(map[frontend.API]*Chip)
108	mutex         sync.Mutex
109)
110
111// Creates a new Goldilocks Chip.
112func New(api frontend.API) *Chip {
113	mutex.Lock()
114	defer mutex.Unlock()
115
116	if chip, ok := poseidonChips[api]; ok {
117		return chip
118	}
119
120	c := &Chip{api: api}
121
122	// Instantiate the range checker gadget
123	// Per Gnark's range checker gadget's New function, there are three possible range checkers:
124	// 1. The native range checker
125	// 2. The commit range checker
126	// 3. The bit decomposition range checker
127	//
128	// See https://github.com/Consensys/gnark/blob/3421eaa7d544286abf3de8c46282b8d4da6d5da0/std/rangecheck/rangecheck.go#L3
129
130	// This function will emulate gnark's range checker selection logic (within the gnarkRangeCheckSelector func).  However,
131	// if the USE_BIT_DECOMPOSITION_RANGE_CHECK env var is set, then it will explicitly use the bit decomposition range checker.
132
133	rangeCheckerType := gnarkRangeCheckerSelector(api)
134	useBitDecomp := os.Getenv("USE_BIT_DECOMPOSITION_RANGE_CHECK")
135	if useBitDecomp == "true" {
136		fmt.Println("The USE_BIT_DECOMPOSITION_RANGE_CHECK env var is set to true.  Using the bit decomposition range checker.")
137		rangeCheckerType = BIT_DECOMP_RANGE_CHECKER
138	}
139
140	c.rangeCheckerType = rangeCheckerType
141
142	// If we are using the bit decomposition range checker, then create bitDecompChecker object
143	if c.rangeCheckerType == BIT_DECOMP_RANGE_CHECKER {
144		c.rangeChecker = bitDecompChecker{api: api}
145	} else {
146		if c.rangeCheckerType == COMMIT_RANGE_CHECKER {
147			api.Compiler().Defer(c.checkCollected)
148		}
149
150		// If we are using the native or commit range checker, then have gnark's range checker gadget's New function create it.
151		// Also, note that the range checker will need to be created AFTER the c.checkCollected function is deferred.
152		// The commit range checker gadget will also call a deferred function, which needs to be called after c.checkCollected.
153		c.rangeChecker = rangecheck.New(api)
154	}
155
156	poseidonChips[api] = c
157
158	return c
159}
160
161// Adds two goldilocks field elements and returns a value within the goldilocks field.
162func (p *Chip) Add(a Variable, b Variable) Variable {
163	return p.MulAdd(a, NewVariable(1), b)
164}
165
166// Adds two goldilocks field elements and returns a value that may not be within the goldilocks field
167// (e.g. the sum is not reduced).
168func (p *Chip) AddNoReduce(a Variable, b Variable) Variable {
169	return NewVariable(p.api.Add(a.Limb, b.Limb))
170}
171
172// Subracts two goldilocks field elements and returns a value within the goldilocks field.
173func (p *Chip) Sub(a Variable, b Variable) Variable {
174	return p.MulAdd(b, NegOne(), a)
175}
176
177// Subracts two goldilocks field elements and returns a value that may not be within the goldilocks field
178// (e.g. the difference is not reduced).
179func (p *Chip) SubNoReduce(a Variable, b Variable) Variable {
180	return NewVariable(p.api.Add(a.Limb, p.api.Mul(b.Limb, NegOne().Limb)))
181}
182
183// Multiplies two goldilocks field elements and returns a value within the goldilocks field.
184func (p *Chip) Mul(a Variable, b Variable) Variable {
185	return p.MulAdd(a, b, Zero())
186}
187
188// Multiplies two goldilocks field elements and returns a value that may not be within the goldilocks field
189// (e.g. the product is not reduced).
190func (p *Chip) MulNoReduce(a Variable, b Variable) Variable {
191	return NewVariable(p.api.Mul(a.Limb, b.Limb))
192}
193
194// Multiplies two field elements and adds a field element (e.g. computes a * b + c).  The returned value
195// will be within the goldilocks field.
196func (p *Chip) MulAdd(a Variable, b Variable, c Variable) Variable {
197	result, err := p.api.Compiler().NewHint(MulAddHint, 2, a.Limb, b.Limb, c.Limb)
198	if err != nil {
199		panic(err)
200	}
201
202	quotient := NewVariable(result[0])
203	remainder := NewVariable(result[1])
204
205	cLimbCopy := p.api.Mul(c.Limb, 1)
206	lhs := p.api.MulAcc(cLimbCopy, a.Limb, b.Limb)
207	rhs := p.api.MulAcc(remainder.Limb, MODULUS, quotient.Limb)
208	p.api.AssertIsEqual(lhs, rhs)
209
210	p.RangeCheck(quotient)
211	p.RangeCheck(remainder)
212	return remainder
213}
214
215// Multiplies two field elements and adds a field element (e.g. computes a * b + c).  The returned value
216// may no be within the goldilocks field (e.g. the result is not reduced).
217func (p *Chip) MulAddNoReduce(a Variable, b Variable, c Variable) Variable {
218	cLimbCopy := p.api.Mul(c.Limb, 1)
219	return NewVariable(p.api.MulAcc(cLimbCopy, a.Limb, b.Limb))
220}
221
222// The hint used to compute MulAdd.
223func MulAddHint(_ *big.Int, inputs []*big.Int, results []*big.Int) error {
224	if len(inputs) != 3 {
225		panic("MulAddHint expects 3 input operands")
226	}
227
228	for _, operand := range inputs {
229		if operand.Cmp(MODULUS) >= 0 {
230			panic(fmt.Sprintf("%s is not in the field", operand.String()))
231		}
232	}
233
234	product := new(big.Int).Mul(inputs[0], inputs[1])
235	sum := new(big.Int).Add(product, inputs[2])
236	quotient := new(big.Int).Div(sum, MODULUS)
237	remainder := new(big.Int).Rem(sum, MODULUS)
238
239	results[0] = quotient
240	results[1] = remainder
241
242	return nil
243}
244
245// Reduces a field element x such that x % MODULUS = y.
246func (p *Chip) Reduce(x Variable) Variable {
247	// Witness a `quotient` and `remainder` such that:
248	//
249	// 		MODULUS * quotient + remainder = x
250	//
251	// Must check that offset \in [0, MODULUS) and carry \in [0, 2^RANGE_CHECK_NB_BITS) to ensure
252	// that this computation does not overflow. We use 2^RANGE_CHECK_NB_BITS to reduce the cost of the range check
253	//
254	// In other words, we assume that we at most compute a a dot product with dimension at most RANGE_CHECK_NB_BITS - 128.
255	return p.ReduceWithMaxBits(x, uint64(RANGE_CHECK_NB_BITS))
256}
257
258// Reduces a field element x such that x % MODULUS = y.
259func (p *Chip) ReduceWithMaxBits(x Variable, maxNbBits uint64) Variable {
260	// Witness a `quotient` and `remainder` such that:
261	//
262	// 		MODULUS * quotient + remainder = x
263	//
264	// Must check that remainder \in [0, MODULUS) and quotient \in [0, 2^maxNbBits) to ensure that this
265	// computation does not overflow.
266
267	result, err := p.api.Compiler().NewHint(ReduceHint, 2, x.Limb)
268	if err != nil {
269		panic(err)
270	}
271
272	quotient := result[0]
273	p.rangeCheckerCheck(quotient, int(maxNbBits))
274
275	remainder := NewVariable(result[1])
276	p.RangeCheck(remainder)
277
278	p.api.AssertIsEqual(x.Limb, p.api.Add(p.api.Mul(quotient, MODULUS), remainder.Limb))
279
280	return remainder
281}
282
283// The hint used to compute Reduce.
284func ReduceHint(_ *big.Int, inputs []*big.Int, results []*big.Int) error {
285	if len(inputs) != 1 {
286		panic("ReduceHint expects 1 input operand")
287	}
288	input := inputs[0]
289	quotient := new(big.Int).Div(input, MODULUS)
290	remainder := new(big.Int).Rem(input, MODULUS)
291	results[0] = quotient
292	results[1] = remainder
293	return nil
294}
295
296// Computes the inverse of a field element x such that x * x^-1 = 1.
297func (p *Chip) Inverse(x Variable) (Variable, frontend.Variable) {
298	result, err := p.api.Compiler().NewHint(InverseHint, 1, x.Limb)
299	if err != nil {
300		panic(err)
301	}
302
303	inverse := NewVariable(result[0])
304	isZero := p.api.IsZero(x.Limb)
305	hasInv := p.api.Sub(1, isZero)
306	p.RangeCheck(inverse)
307
308	product := p.Mul(inverse, x)
309	productToCheck := p.api.Select(hasInv, product.Limb, frontend.Variable(1))
310	p.api.AssertIsEqual(productToCheck, frontend.Variable(1))
311
312	return inverse, hasInv
313}
314
315// The hint used to compute Inverse.
316func InverseHint(_ *big.Int, inputs []*big.Int, results []*big.Int) error {
317	if len(inputs) != 1 {
318		panic("InverseHint expects 1 input operand")
319	}
320
321	input := inputs[0]
322	if input.Cmp(MODULUS) == 0 || input.Cmp(MODULUS) == 1 {
323		panic("Input is not in the field")
324	}
325
326	inputGl := goldilocks.NewElement(input.Uint64())
327	resultGl := goldilocks.NewElement(0)
328
329	// Will set resultGL if inputGL == 0
330	resultGl.Inverse(&inputGl)
331
332	result := big.NewInt(0)
333	results[0] = resultGl.BigInt(result)
334
335	return nil
336}
337
338// The hint used to split a GoldilocksVariable into 2 32 bit limbs.
339func SplitLimbsHint(_ *big.Int, inputs []*big.Int, results []*big.Int) error {
340	if len(inputs) != 1 {
341		panic("SplitLimbsHint expects 1 input operand")
342	}
343
344	// The Goldilocks field element
345	input := inputs[0]
346
347	if input.Cmp(MODULUS) == 0 || input.Cmp(MODULUS) == 1 {
348		return fmt.Errorf("input is not in the field")
349	}
350
351	two_32 := big.NewInt(int64(math.Pow(2, 32)))
352
353	// The most significant bits
354	results[0] = new(big.Int).Quo(input, two_32)
355	// The least significant bits
356	results[1] = new(big.Int).Rem(input, two_32)
357
358	return nil
359}
360
361// Range checks a field element x to be less than the Golidlocks modulus 2 ^ 64 - 2 ^ 32 + 1.
362func (p *Chip) RangeCheck(x Variable) {
363	// The Goldilocks' modulus is 2^64 - 2^32 + 1, which is:
364	//
365	// 		1111111111111111111111111111111100000000000000000000000000000001
366	//
367	// in big endian binary. This function will first verify that x is at most 64 bits wide. Then it
368	// checks that if the bits[0:31] (in big-endian) are all 1, then bits[32:64] are all zero.
369
370	result, err := p.api.Compiler().NewHint(SplitLimbsHint, 2, x.Limb)
371	if err != nil {
372		panic(err)
373	}
374
375	// We check that this is a valid decomposition of the Goldilock's element and range-check each limb.
376	mostSigLimb := result[0]
377	leastSigLimb := result[1]
378	p.api.AssertIsEqual(
379		p.api.Add(
380			p.api.Mul(mostSigLimb, uint64(math.Pow(2, 32))),
381			leastSigLimb,
382		),
383		x.Limb,
384	)
385	p.rangeCheckerCheck(mostSigLimb, 32)
386	p.rangeCheckerCheck(leastSigLimb, 32)
387
388	// If the most significant bits are all 1, then we need to check that the least significant bits are all zero
389	// in order for element to be less than the Goldilock's modulus.
390	// Otherwise, we don't need to do any checks, since we already know that the element is less than the Goldilocks modulus.
391	shouldCheck := p.api.IsZero(p.api.Sub(mostSigLimb, uint64(math.Pow(2, 32))-1))
392	p.api.AssertIsEqual(
393		p.api.Select(
394			shouldCheck,
395			leastSigLimb,
396			frontend.Variable(0),
397		),
398		frontend.Variable(0),
399	)
400}
401
402// This function will assert that the field element x is less than 2^maxNbBits.
403func (p *Chip) RangeCheckWithMaxBits(x Variable, maxNbBits uint64) {
404	p.rangeCheckerCheck(x.Limb, int(maxNbBits))
405}
406
407func (p *Chip) AssertIsEqual(x, y Variable) {
408	p.api.AssertIsEqual(x.Limb, y.Limb)
409}
410
411func (p *Chip) rangeCheckerCheck(x frontend.Variable, nbBits int) {
412	switch p.rangeCheckerType {
413	case NATIVE_RANGE_CHECKER:
414	case BIT_DECOMP_RANGE_CHECKER:
415		p.rangeChecker.Check(x, nbBits)
416	case COMMIT_RANGE_CHECKER:
417		p.collectedMutex.Lock()
418		defer p.collectedMutex.Unlock()
419		p.rangeCheckCollected = append(p.rangeCheckCollected, checkedVariable{v: x, bits: nbBits})
420	}
421}
422
423func (p *Chip) checkCollected(api frontend.API) error {
424	if p.rangeCheckerType != COMMIT_RANGE_CHECKER {
425		panic("checkCollected should only be called when using the commit range checker")
426	}
427
428	nbBits := getOptimalBasewidth(p.api, p.rangeCheckCollected)
429	if nbBits != EXPECTED_OPTIMAL_BASEWIDTH {
430		panic("nbBits should be " + strconv.Itoa(EXPECTED_OPTIMAL_BASEWIDTH))
431	}
432
433	for _, v := range p.rangeCheckCollected {
434		if v.bits%nbBits != 0 {
435			panic("v.bits is not nbBits aligned")
436		}
437
438		p.rangeChecker.Check(v.v, v.bits)
439	}
440
441	return nil
442}
443
444// Computes the n'th primitive root of unity for the Goldilocks field.
445func PrimitiveRootOfUnity(nLog uint64) goldilocks.Element {
446	if nLog > TWO_ADICITY {
447		panic("nLog is greater than TWO_ADICITY")
448	}
449	res := goldilocks.NewElement(POWER_OF_TWO_GENERATOR.Uint64())
450	for i := 0; i < int(TWO_ADICITY-nLog); i++ {
451		res.Square(&res)
452	}
453	return res
454}
455
456func TwoAdicSubgroup(nLog uint64) []goldilocks.Element {
457	if nLog > TWO_ADICITY {
458		panic("nLog is greater than GOLDILOCKS_TWO_ADICITY")
459	}
460
461	var res []goldilocks.Element
462	rootOfUnity := PrimitiveRootOfUnity(nLog)
463	res = append(res, goldilocks.NewElement(1))
464
465	for i := 0; i < (1<<nLog)-1; i++ {
466		lastElement := res[len(res)-1]
467		res = append(res, *lastElement.Mul(&lastElement, &rootOfUnity))
468	}
469
470	return res
471}