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, frontend.Variable) {
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
139	aPowRInv, hasInv := p.Inverse(aPowR[0])
140	return p.ScalarMulExtension(aPowRMinus1, aPowRInv), hasInv
141}
142
143// Divides two quadratic extension variables in the Goldilocks field.
144func (p *Chip) DivExtension(a, b QuadraticExtensionVariable) (QuadraticExtensionVariable, frontend.Variable) {
145	bInv, hasInv := p.InverseExtension(b)
146	return p.MulExtension(a, bInv), hasInv
147}
148
149// Exponentiates a quadratic extension variable to some exponent in the Golidlocks field.
150func (p *Chip) ExpExtension(
151	a QuadraticExtensionVariable,
152	exponent uint64,
153) QuadraticExtensionVariable {
154	switch exponent {
155	case 0:
156		return OneExtension()
157	case 1:
158		return a
159	case 2:
160		return p.MulExtension(a, a)
161	default:
162	}
163
164	current := a
165	product := OneExtension()
166
167	for i := 0; i < bits.Len64(exponent); i++ {
168		if i != 0 {
169			current = p.MulExtension(current, current)
170		}
171		if (exponent >> i & 1) != 0 {
172			product = p.MulExtension(product, current)
173		}
174	}
175
176	return product
177}
178
179func (p *Chip) ReduceExtension(x QuadraticExtensionVariable) QuadraticExtensionVariable {
180	return NewQuadraticExtensionVariable(p.Reduce(x[0]), p.Reduce(x[1]))
181}
182
183// Reduces a list of extension field terms with a scalar power in the Goldilocks field.
184func (p *Chip) ReduceWithPowers(
185	terms []QuadraticExtensionVariable,
186	scalar QuadraticExtensionVariable,
187) QuadraticExtensionVariable {
188	sum := ZeroExtension()
189	for i := len(terms) - 1; i >= 0; i-- {
190		sum = p.AddExtensionNoReduce(
191			p.MulExtensionNoReduce(
192				sum,
193				scalar,
194			),
195			terms[i],
196		)
197		sum = p.ReduceExtension(sum)
198	}
199	return sum
200}
201
202// Outputs whether the quadratic extension variable is zero.
203func (p *Chip) IsZero(x QuadraticExtensionVariable) frontend.Variable {
204	x0IsZero := p.api.IsZero(x[0].Limb)
205	x1IsZero := p.api.IsZero(x[1].Limb)
206	return p.api.Mul(x0IsZero, x1IsZero)
207}
208
209// Lookup is similar to select, but returns the first variable if the bit is zero and vice-versa.
210func (p *Chip) Lookup(
211	b frontend.Variable,
212	x, y QuadraticExtensionVariable,
213) QuadraticExtensionVariable {
214	c0 := p.api.Select(b, y[0].Limb, x[0].Limb)
215	c1 := p.api.Select(b, y[1].Limb, x[1].Limb)
216	return NewQuadraticExtensionVariable(NewVariable(c0), NewVariable(c1))
217}
218
219// Lookup2 is similar to select2, but returns the first variable if the bit is zero and vice-versa.
220func (p *Chip) Lookup2(
221	b0 frontend.Variable,
222	b1 frontend.Variable,
223	qe0, qe1, qe2, qe3 QuadraticExtensionVariable,
224) QuadraticExtensionVariable {
225	c0 := p.Lookup(b0, qe0, qe1)
226	c1 := p.Lookup(b0, qe2, qe3)
227	return p.Lookup(b1, c0, c1)
228}
229
230// Asserts that two quadratic extension variables are equal.
231func (p *Chip) AssertIsEqualExtension(
232	a QuadraticExtensionVariable,
233	b QuadraticExtensionVariable,
234) {
235	p.AssertIsEqual(a[0], b[0])
236	p.AssertIsEqual(a[1], b[1])
237}
238
239func (p *Chip) RangeCheckQE(a QuadraticExtensionVariable) {
240	p.RangeCheck(a[0])
241	p.RangeCheck(a[1])
242}