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	return p.ReduceExtension(product)
 62}
 63
 64// Multiplies quadratic extension variable in the Goldilocks field without reducing.
 65func (p *Chip) MulExtensionNoReduce(a, b QuadraticExtensionVariable) QuadraticExtensionVariable {
 66	c0o0 := p.MulNoReduce(a[0], b[0])
 67	c0o1 := p.MulNoReduce(p.MulNoReduce(NewVariable(W), a[1]), b[1])
 68	c0 := p.AddNoReduce(c0o0, c0o1)
 69	c1 := p.AddNoReduce(p.MulNoReduce(a[0], b[1]), p.MulNoReduce(a[1], b[0]))
 70	return NewQuadraticExtensionVariable(c0, c1)
 71}
 72
 73// Multiplies two operands a and b and adds to c in the Goldilocks extension field. a * b + c must
 74// be less than RANGE_CHECK_NB_BITS bits.
 75func (p *Chip) MulAddExtension(a, b, c QuadraticExtensionVariable) QuadraticExtensionVariable {
 76	product := p.MulExtensionNoReduce(a, b)
 77	sum := p.AddExtensionNoReduce(product, c)
 78	return p.ReduceExtension(sum)
 79}
 80
 81func (p *Chip) MulAddExtensionNoReduce(a, b, c QuadraticExtensionVariable) QuadraticExtensionVariable {
 82	product := p.MulExtensionNoReduce(a, b)
 83	sum := p.AddExtensionNoReduce(product, c)
 84	return sum
 85}
 86
 87// Multiplies two operands a and b and subtracts to c in the Goldilocks extension field. a * b - c must
 88// be less than RANGE_CHECK_NB_BITS bits.
 89func (p *Chip) SubMulExtension(a, b, c QuadraticExtensionVariable) QuadraticExtensionVariable {
 90	difference := p.SubExtensionNoReduce(a, b)
 91	product := p.MulExtensionNoReduce(difference, c)
 92	return p.ReduceExtension(product)
 93}
 94
 95// Multiplies quadratic extension variable in the Goldilocks field by a scalar.
 96func (p *Chip) ScalarMulExtension(
 97	a QuadraticExtensionVariable,
 98	b Variable,
 99) QuadraticExtensionVariable {
100	return NewQuadraticExtensionVariable(
101		p.Mul(a[0], b),
102		p.Mul(a[1], b),
103	)
104}
105
106// Computes an inner product over quadratic extension variable vectors in the Goldilocks field.
107func (p *Chip) InnerProductExtension(
108	constant Variable,
109	startingAcc QuadraticExtensionVariable,
110	pairs [][2]QuadraticExtensionVariable,
111) QuadraticExtensionVariable {
112	acc := startingAcc
113	for i := 0; i < len(pairs); i++ {
114		a := pairs[i][0]
115		b := pairs[i][1]
116		mul := p.ScalarMulExtension(a, constant)
117		acc = p.MulAddExtensionNoReduce(mul, b, acc)
118	}
119	return p.ReduceExtension(acc)
120}
121
122// Computes the inverse of a quadratic extension variable in the Goldilocks field.
123func (p *Chip) InverseExtension(a QuadraticExtensionVariable) (QuadraticExtensionVariable, frontend.Variable) {
124	aIsZero := p.IsZero(a)
125	p.api.AssertIsEqual(aIsZero, frontend.Variable(0))
126	aPowRMinus1 := QuadraticExtensionVariable{
127		a[0],
128		p.Mul(a[1], NewVariable(DTH_ROOT)),
129	}
130	aPowR := p.MulExtension(aPowRMinus1, a)
131
132	aPowRInv, hasInv := p.Inverse(aPowR[0])
133	return p.ScalarMulExtension(aPowRMinus1, aPowRInv), hasInv
134}
135
136// Divides two quadratic extension variables in the Goldilocks field.
137func (p *Chip) DivExtension(a, b QuadraticExtensionVariable) (QuadraticExtensionVariable, frontend.Variable) {
138	bInv, hasInv := p.InverseExtension(b)
139	return p.MulExtension(a, bInv), hasInv
140}
141
142// Exponentiates a quadratic extension variable to some exponent in the Golidlocks field.
143func (p *Chip) ExpExtension(
144	a QuadraticExtensionVariable,
145	exponent uint64,
146) QuadraticExtensionVariable {
147	switch exponent {
148	case 0:
149		return OneExtension()
150	case 1:
151		return a
152	case 2:
153		return p.MulExtension(a, a)
154	default:
155	}
156
157	current := a
158	product := OneExtension()
159
160	for i := 0; i < bits.Len64(exponent); i++ {
161		if i != 0 {
162			current = p.MulExtension(current, current)
163		}
164		if (exponent >> i & 1) != 0 {
165			product = p.MulExtension(product, current)
166		}
167	}
168
169	return product
170}
171
172func (p *Chip) ReduceExtension(x QuadraticExtensionVariable) QuadraticExtensionVariable {
173	return NewQuadraticExtensionVariable(p.Reduce(x[0]), p.Reduce(x[1]))
174}
175
176// Reduces a list of extension field terms with a scalar power in the Goldilocks field.
177func (p *Chip) ReduceWithPowers(
178	terms []QuadraticExtensionVariable,
179	scalar QuadraticExtensionVariable,
180) QuadraticExtensionVariable {
181	sum := ZeroExtension()
182	for i := len(terms) - 1; i >= 0; i-- {
183		sum = p.AddExtensionNoReduce(
184			p.MulExtensionNoReduce(
185				sum,
186				scalar,
187			),
188			terms[i],
189		)
190		sum = p.ReduceExtension(sum)
191	}
192	return sum
193}
194
195// Outputs whether the quadratic extension variable is zero.
196func (p *Chip) IsZero(x QuadraticExtensionVariable) frontend.Variable {
197	x0IsZero := p.api.IsZero(x[0].Limb)
198	x1IsZero := p.api.IsZero(x[1].Limb)
199	return p.api.Mul(x0IsZero, x1IsZero)
200}
201
202// Lookup is similar to select, but returns the first variable if the bit is zero and vice-versa.
203func (p *Chip) Lookup(
204	b frontend.Variable,
205	x, y QuadraticExtensionVariable,
206) QuadraticExtensionVariable {
207	c0 := p.api.Select(b, y[0].Limb, x[0].Limb)
208	c1 := p.api.Select(b, y[1].Limb, x[1].Limb)
209	return NewQuadraticExtensionVariable(NewVariable(c0), NewVariable(c1))
210}
211
212// Lookup2 is similar to select2, but returns the first variable if the bit is zero and vice-versa.
213func (p *Chip) Lookup2(
214	b0 frontend.Variable,
215	b1 frontend.Variable,
216	qe0, qe1, qe2, qe3 QuadraticExtensionVariable,
217) QuadraticExtensionVariable {
218	c0 := p.Lookup(b0, qe0, qe1)
219	c1 := p.Lookup(b0, qe2, qe3)
220	return p.Lookup(b1, c0, c1)
221}
222
223// Asserts that two quadratic extension variables are equal.
224func (p *Chip) AssertIsEqualExtension(
225	a QuadraticExtensionVariable,
226	b QuadraticExtensionVariable,
227) {
228	p.AssertIsEqual(a[0], b[0])
229	p.AssertIsEqual(a[1], b[1])
230}
231
232func (p *Chip) RangeCheckQE(a QuadraticExtensionVariable) {
233	p.RangeCheck(a[0])
234	p.RangeCheck(a[1])
235}