1package goldilocks
  2
  3import (
  4	"math/bits"
  5
  6	"github.com/consensys/gnark/frontend"
  7)
  8
  9const W uint64 = 7
 10const DTH_ROOT uint64 = 18446744069414584320
 11
 12type QuadraticExtensionVariable [2]Variable
 13
 14func NewQuadraticExtensionVariable(x Variable, y Variable) QuadraticExtensionVariable {
 15	return QuadraticExtensionVariable{x, y}
 16}
 17
 18func (p Variable) ToQuadraticExtension() QuadraticExtensionVariable {
 19	return NewQuadraticExtensionVariable(p, Zero())
 20}
 21
 22func ZeroExtension() QuadraticExtensionVariable {
 23	return Zero().ToQuadraticExtension()
 24}
 25
 26func OneExtension() QuadraticExtensionVariable {
 27	return One().ToQuadraticExtension()
 28}
 29
 30// Adds two quadratic extension variables in the Goldilocks field.
 31func (p *Chip) AddExtension(a, b QuadraticExtensionVariable) QuadraticExtensionVariable {
 32	c0 := p.Add(a[0], b[0])
 33	c1 := p.Add(a[1], b[1])
 34	return NewQuadraticExtensionVariable(c0, c1)
 35}
 36
 37// Adds two quadratic extension variables in the Goldilocks field without reducing.
 38func (p *Chip) AddExtensionNoReduce(a, b QuadraticExtensionVariable) QuadraticExtensionVariable {
 39	c0 := p.AddNoReduce(a[0], b[0])
 40	c1 := p.AddNoReduce(a[1], b[1])
 41	return NewQuadraticExtensionVariable(c0, c1)
 42}
 43
 44// Subtracts two quadratic extension variables in the Goldilocks field.
 45func (p *Chip) SubExtension(a, b QuadraticExtensionVariable) QuadraticExtensionVariable {
 46	c0 := p.Sub(a[0], b[0])
 47	c1 := p.Sub(a[1], b[1])
 48	return NewQuadraticExtensionVariable(c0, c1)
 49}
 50
 51// Subtracts two quadratic extension variables in the Goldilocks field without reducing.
 52func (p *Chip) SubExtensionNoReduce(a, b QuadraticExtensionVariable) QuadraticExtensionVariable {
 53	c0 := p.SubNoReduce(a[0], b[0])
 54	c1 := p.SubNoReduce(a[1], b[1])
 55	return NewQuadraticExtensionVariable(c0, c1)
 56}
 57
 58// Multiplies quadratic extension variable in the Goldilocks field.
 59func (p *Chip) MulExtension(a, b QuadraticExtensionVariable) QuadraticExtensionVariable {
 60	product := p.MulExtensionNoReduce(a, b)
 61	product[0] = p.Reduce(product[0])
 62	product[1] = p.Reduce(product[1])
 63	return product
 64}
 65
 66// Multiplies quadratic extension variable in the Goldilocks field without reducing.
 67func (p *Chip) MulExtensionNoReduce(a, b QuadraticExtensionVariable) QuadraticExtensionVariable {
 68	c0o0 := p.MulNoReduce(a[0], b[0])
 69	c0o1 := p.MulNoReduce(p.MulNoReduce(NewVariable(7), a[1]), b[1])
 70	c0 := p.AddNoReduce(c0o0, c0o1)
 71	c1 := p.AddNoReduce(p.MulNoReduce(a[0], b[1]), p.MulNoReduce(a[1], b[0]))
 72	return NewQuadraticExtensionVariable(c0, c1)
 73}
 74
 75// Multiplies two operands a and b and adds to c in the Goldilocks extension field. a * b + c must
 76// be less than RANGE_CHECK_NB_BITS bits.
 77func (p *Chip) MulAddExtension(a, b, c QuadraticExtensionVariable) QuadraticExtensionVariable {
 78	product := p.MulExtensionNoReduce(a, b)
 79	sum := p.AddExtensionNoReduce(product, c)
 80	sum[0] = p.Reduce(sum[0])
 81	sum[1] = p.Reduce(sum[1])
 82	return sum
 83}
 84
 85func (p *Chip) MulAddExtensionNoReduce(a, b, c QuadraticExtensionVariable) QuadraticExtensionVariable {
 86	product := p.MulExtensionNoReduce(a, b)
 87	sum := p.AddExtensionNoReduce(product, c)
 88	return sum
 89}
 90
 91// Multiplies two operands a and b and subtracts to c in the Goldilocks extension field. a * b - c must
 92// be less than RANGE_CHECK_NB_BITS bits.
 93func (p *Chip) SubMulExtension(a, b, c QuadraticExtensionVariable) QuadraticExtensionVariable {
 94	difference := p.SubExtensionNoReduce(a, b)
 95	product := p.MulExtensionNoReduce(difference, c)
 96	product[0] = p.Reduce(product[0])
 97	product[1] = p.Reduce(product[1])
 98	return product
 99}
100
101// Multiplies quadratic extension variable in the Goldilocks field by a scalar.
102func (p *Chip) ScalarMulExtension(
103	a QuadraticExtensionVariable,
104	b Variable,
105) QuadraticExtensionVariable {
106	return NewQuadraticExtensionVariable(
107		p.Mul(a[0], b),
108		p.Mul(a[1], b),
109	)
110}
111
112// Computes an inner product over quadratic extension variable vectors in the Goldilocks field.
113func (p *Chip) InnerProductExtension(
114	constant Variable,
115	startingAcc QuadraticExtensionVariable,
116	pairs [][2]QuadraticExtensionVariable,
117) QuadraticExtensionVariable {
118	acc := startingAcc
119	for i := 0; i < len(pairs); i++ {
120		a := pairs[i][0]
121		b := pairs[i][1]
122		mul := p.ScalarMulExtension(a, constant)
123		acc = p.MulAddExtensionNoReduce(mul, b, acc)
124	}
125	return p.ReduceExtension(acc)
126}
127
128// Computes the inverse of a quadratic extension variable in the Goldilocks field.
129func (p *Chip) InverseExtension(a QuadraticExtensionVariable) QuadraticExtensionVariable {
130	a0IsZero := p.api.IsZero(a[0].Limb)
131	a1IsZero := p.api.IsZero(a[1].Limb)
132	p.api.AssertIsEqual(p.api.Mul(a0IsZero, a1IsZero), frontend.Variable(0))
133	aPowRMinus1 := QuadraticExtensionVariable{
134		a[0],
135		p.Mul(a[1], NewVariable(DTH_ROOT)),
136	}
137	aPowR := p.MulExtension(aPowRMinus1, a)
138	return p.ScalarMulExtension(aPowRMinus1, p.Inverse(aPowR[0]))
139}
140
141// Divides two quadratic extension variables in the Goldilocks field.
142func (p *Chip) DivExtension(a, b QuadraticExtensionVariable) QuadraticExtensionVariable {
143	return p.MulExtension(a, p.InverseExtension(b))
144}
145
146// Exponentiates a quadratic extension variable to some exponent in the Golidlocks field.
147func (p *Chip) ExpExtension(
148	a QuadraticExtensionVariable,
149	exponent uint64,
150) QuadraticExtensionVariable {
151	switch exponent {
152	case 0:
153		return OneExtension()
154	case 1:
155		return a
156	case 2:
157		return p.MulExtension(a, a)
158	default:
159	}
160
161	current := a
162	product := OneExtension()
163
164	for i := 0; i < bits.Len64(exponent); i++ {
165		if i != 0 {
166			current = p.MulExtension(current, current)
167		}
168		if (exponent >> i & 1) != 0 {
169			product = p.MulExtension(product, current)
170		}
171	}
172
173	return product
174}
175
176func (p *Chip) ReduceExtension(x QuadraticExtensionVariable) QuadraticExtensionVariable {
177	return NewQuadraticExtensionVariable(p.Reduce(x[0]), p.Reduce(x[1]))
178}
179
180// Reduces a list of extension field terms with a scalar power in the Goldilocks field.
181func (p *Chip) ReduceWithPowers(
182	terms []QuadraticExtensionVariable,
183	scalar QuadraticExtensionVariable,
184) QuadraticExtensionVariable {
185	sum := ZeroExtension()
186	for i := len(terms) - 1; i >= 0; i-- {
187		sum = p.AddExtensionNoReduce(
188			p.MulExtensionNoReduce(
189				sum,
190				scalar,
191			),
192			terms[i],
193		)
194		sum = p.ReduceExtension(sum)
195	}
196	return sum
197}
198
199// Outputs whether the quadratic extension variable is zero.
200func (p *Chip) IsZero(x QuadraticExtensionVariable) frontend.Variable {
201	x0IsZero := p.api.IsZero(x[0].Limb)
202	x1IsZero := p.api.IsZero(x[1].Limb)
203	return p.api.Mul(x0IsZero, x1IsZero)
204}
205
206// Lookup is similar to select, but returns the first variable if the bit is zero and vice-versa.
207func (p *Chip) Lookup(
208	b frontend.Variable,
209	x, y QuadraticExtensionVariable,
210) QuadraticExtensionVariable {
211	c0 := p.api.Select(b, y[0].Limb, x[0].Limb)
212	c1 := p.api.Select(b, y[1].Limb, x[1].Limb)
213	return NewQuadraticExtensionVariable(NewVariable(c0), NewVariable(c1))
214}
215
216// Lookup2 is similar to select2, but returns the first variable if the bit is zero and vice-versa.
217func (p *Chip) Lookup2(
218	b0 frontend.Variable,
219	b1 frontend.Variable,
220	qe0, qe1, qe2, qe3 QuadraticExtensionVariable,
221) QuadraticExtensionVariable {
222	c0 := p.Lookup(b0, qe0, qe1)
223	c1 := p.Lookup(b0, qe2, qe3)
224	return p.Lookup(b1, c0, c1)
225}
226
227// Asserts that two quadratic extension variables are equal.
228func (p *Chip) AssertIsEqualExtension(
229	a QuadraticExtensionVariable,
230	b QuadraticExtensionVariable,
231) {
232	p.AssertIsEqual(a[0], b[0])
233	p.AssertIsEqual(a[1], b[1])
234}
235
236func (p *Chip) RangeCheckQE(a QuadraticExtensionVariable) {
237	p.RangeCheck(a[0])
238	p.RangeCheck(a[1])
239}