1package field
  2
  3import (
  4	"fmt"
  5	"math/bits"
  6
  7	"github.com/consensys/gnark-crypto/field/goldilocks"
  8	"github.com/consensys/gnark/frontend"
  9)
 10
 11const D = 2
 12
 13type QuadraticExtension = [2]F
 14type QEAlgebra = [D]QuadraticExtension
 15
 16type QuadraticExtensionAPI struct {
 17	fieldAPI frontend.API
 18
 19	W        F
 20	DTH_ROOT F
 21
 22	ONE_QE  QuadraticExtension
 23	ZERO_QE QuadraticExtension
 24
 25	ZERO_QE_ALGEBRA QEAlgebra
 26}
 27
 28func NewQuadraticExtensionAPI(fieldAPI frontend.API, degreeBits uint64) *QuadraticExtensionAPI {
 29	// TODO:  Should degreeBits be verified that it fits within the field and that degree is within uint64?
 30
 31	var ZERO_QE = QuadraticExtension{ZERO_F, ZERO_F}
 32
 33	var ZERO_QE_ALGEBRA QEAlgebra
 34
 35	for i := 0; i < D; i++ {
 36		ZERO_QE_ALGEBRA[i] = ZERO_QE
 37	}
 38
 39	return &QuadraticExtensionAPI{
 40		fieldAPI: fieldAPI,
 41
 42		W:        NewFieldElement(7),
 43		DTH_ROOT: NewFieldElement(18446744069414584320),
 44
 45		ONE_QE:  QuadraticExtension{ONE_F, ZERO_F},
 46		ZERO_QE: ZERO_QE,
 47
 48		ZERO_QE_ALGEBRA: ZERO_QE_ALGEBRA,
 49	}
 50}
 51
 52func (c *QuadraticExtensionAPI) SquareExtension(a QuadraticExtension) QuadraticExtension {
 53	return c.MulExtension(a, a)
 54}
 55
 56func (c *QuadraticExtensionAPI) MulExtension(a QuadraticExtension, b QuadraticExtension) QuadraticExtension {
 57	c_0 := c.fieldAPI.Add(c.fieldAPI.Mul(a[0], b[0]).(F), c.fieldAPI.Mul(c.W, a[1], b[1])).(F)
 58	c_1 := c.fieldAPI.Add(c.fieldAPI.Mul(a[0], b[1]).(F), c.fieldAPI.Mul(a[1], b[0])).(F)
 59	return QuadraticExtension{c_0, c_1}
 60}
 61
 62func (c *QuadraticExtensionAPI) AddExtension(a QuadraticExtension, b QuadraticExtension) QuadraticExtension {
 63	c_0 := c.fieldAPI.Add(a[0], b[0]).(F)
 64	c_1 := c.fieldAPI.Add(a[1], b[1]).(F)
 65	return QuadraticExtension{c_0, c_1}
 66}
 67
 68func (c *QuadraticExtensionAPI) SubExtension(a QuadraticExtension, b QuadraticExtension) QuadraticExtension {
 69	c_0 := c.fieldAPI.Sub(a[0], b[0]).(F)
 70	c_1 := c.fieldAPI.Sub(a[1], b[1]).(F)
 71	return QuadraticExtension{c_0, c_1}
 72}
 73
 74func (c *QuadraticExtensionAPI) DivExtension(a QuadraticExtension, b QuadraticExtension) QuadraticExtension {
 75	return c.MulExtension(a, c.InverseExtension(b))
 76}
 77
 78func (c *QuadraticExtensionAPI) IsZero(a QuadraticExtension) frontend.Variable {
 79	return c.fieldAPI.Mul(c.fieldAPI.IsZero(a[0]), c.fieldAPI.IsZero(a[1]))
 80}
 81
 82// TODO: Instead of calculating the inverse within the circuit, can witness the
 83// inverse and assert that a_inverse * a = 1.  Should reduce # of constraints.
 84func (c *QuadraticExtensionAPI) InverseExtension(a QuadraticExtension) QuadraticExtension {
 85	// First assert that a doesn't have 0 value coefficients
 86	a0_is_zero := c.fieldAPI.IsZero(a[0])
 87	a1_is_zero := c.fieldAPI.IsZero(a[1])
 88
 89	// assert that a0_is_zero OR a1_is_zero == false
 90	c.fieldAPI.AssertIsEqual(c.fieldAPI.Mul(a0_is_zero, a1_is_zero).(F), ZERO_F)
 91
 92	a_pow_r_minus_1 := QuadraticExtension{a[0], c.fieldAPI.Mul(a[1], c.DTH_ROOT).(F)}
 93	a_pow_r := c.MulExtension(a_pow_r_minus_1, a)
 94	return c.ScalarMulExtension(a_pow_r_minus_1, c.fieldAPI.Inverse(a_pow_r[0]).(F))
 95}
 96
 97func (c *QuadraticExtensionAPI) ScalarMulExtension(a QuadraticExtension, scalar F) QuadraticExtension {
 98	return QuadraticExtension{c.fieldAPI.Mul(a[0], scalar).(F), c.fieldAPI.Mul(a[1], scalar).(F)}
 99}
100
101func (c *QuadraticExtensionAPI) FieldToQE(a F) QuadraticExtension {
102	return QuadraticExtension{a, ZERO_F}
103}
104
105// / Exponentiate `base` to the power of a known `exponent`.
106func (c *QuadraticExtensionAPI) ExpU64Extension(a QuadraticExtension, exponent uint64) QuadraticExtension {
107	switch exponent {
108	case 0:
109		return c.ONE_QE
110	case 1:
111		return a
112	case 2:
113		return c.SquareExtension(a)
114	default:
115	}
116
117	current := a
118	product := c.ONE_QE
119
120	for i := 0; i < bits.Len64(exponent); i++ {
121		if i != 0 {
122			current = c.SquareExtension(current)
123		}
124
125		if (exponent >> i & 1) != 0 {
126			product = c.MulExtension(product, current)
127		}
128	}
129
130	return product
131}
132
133func (c *QuadraticExtensionAPI) ReduceWithPowers(terms []QuadraticExtension, scalar QuadraticExtension) QuadraticExtension {
134	sum := c.ZERO_QE
135
136	for i := len(terms) - 1; i >= 0; i-- {
137		sum = c.AddExtension(
138			c.MulExtension(
139				sum,
140				scalar,
141			),
142			terms[i],
143		)
144	}
145
146	return sum
147}
148
149func (c *QuadraticExtensionAPI) Select(b frontend.Variable, qe0, qe1 QuadraticExtension) QuadraticExtension {
150	var retQE QuadraticExtension
151
152	for i := 0; i < 2; i++ {
153		retQE[i] = c.fieldAPI.Select(b, qe0[i], qe1[i]).(F)
154	}
155
156	return retQE
157}
158
159func (c *QuadraticExtensionAPI) Lookup2(b0 frontend.Variable, b1 frontend.Variable, qe0, qe1, qe2, qe3 QuadraticExtension) QuadraticExtension {
160	var retQE QuadraticExtension
161
162	for i := 0; i < 2; i++ {
163		retQE[i] = c.fieldAPI.Lookup2(b0, b1, qe0[i], qe1[i], qe2[i], qe3[i]).(F)
164	}
165
166	return retQE
167}
168
169func (c *QuadraticExtensionAPI) AssertIsEqual(a, b QuadraticExtension) {
170	for i := 0; i < 2; i++ {
171		c.fieldAPI.AssertIsEqual(a[i], b[i])
172	}
173}
174
175func (c *QuadraticExtensionAPI) InnerProductExtension(constant F, startingAcc QuadraticExtension, pairs [][2]QuadraticExtension) QuadraticExtension {
176	acc := startingAcc
177
178	for i := 0; i < len(pairs); i++ {
179		a := pairs[i][0]
180		b := pairs[i][1]
181		mul := c.ScalarMulExtension(a, constant)
182		mul = c.MulExtension(mul, b)
183		acc = c.AddExtension(acc, mul)
184	}
185
186	return acc
187}
188
189func (c *QuadraticExtensionAPI) Println(a QuadraticExtension) {
190	fmt.Print("Degree 0 coefficient")
191	c.fieldAPI.Println(a[0])
192
193	fmt.Print("Degree 1 coefficient")
194	c.fieldAPI.Println(a[1])
195}
196
197func (c *QuadraticExtensionAPI) MulExtensionAlgebra(a, b QEAlgebra) QEAlgebra {
198	var inner [D][][2]QuadraticExtension
199	var inner_w [D][][2]QuadraticExtension
200	for i := 0; i < D; i++ {
201		for j := 0; j < D-i; j++ {
202			idx := (i + j) % D
203			inner[idx] = append(inner[idx], [2]QuadraticExtension{a[i], b[j]})
204		}
205		for j := D - i; j < D; j++ {
206			idx := (i + j) % D
207			inner_w[idx] = append(inner_w[idx], [2]QuadraticExtension{a[i], b[j]})
208		}
209	}
210
211	var product QEAlgebra
212	for i := 0; i < D; i++ {
213		acc := c.InnerProductExtension(NewFieldElement(7), c.ZERO_QE, inner_w[i])
214		product[i] = c.InnerProductExtension(ONE_F, acc, inner[i])
215	}
216
217	return product
218}
219
220func (c *QuadraticExtensionAPI) ScalarMulExtensionAlgebra(a QuadraticExtension, b QEAlgebra) QEAlgebra {
221	var product QEAlgebra
222	for i := 0; i < D; i++ {
223		product[i] = c.MulExtension(a, b[i])
224	}
225
226	return product
227}
228
229func (c *QuadraticExtensionAPI) AddExtensionAlgebra(a, b QEAlgebra) QEAlgebra {
230	var sum QEAlgebra
231	for i := 0; i < D; i++ {
232		sum[i] = c.AddExtension(a[i], b[i])
233	}
234
235	return sum
236}
237
238func (c *QuadraticExtensionAPI) SubExtensionAlgebra(a, b QEAlgebra) QEAlgebra {
239	var diff QEAlgebra
240	for i := 0; i < D; i++ {
241		diff[i] = c.SubExtension(a[i], b[i])
242	}
243
244	return diff
245}
246
247func (c *QuadraticExtensionAPI) PartialInterpolateExtAlgebra(
248	domain []goldilocks.Element,
249	values []QEAlgebra,
250	barycentricWeights []goldilocks.Element,
251	point QEAlgebra,
252	initialEval QEAlgebra,
253	initialPartialProd QEAlgebra,
254) (QEAlgebra, QEAlgebra) {
255	n := len(values)
256	if n == 0 {
257		panic("Cannot interpolate with no values")
258	}
259	if n != len(domain) {
260		panic("Domain and values must have the same length")
261	}
262	if n != len(barycentricWeights) {
263		panic("Domain and barycentric weights must have the same length")
264	}
265
266	newEval := initialEval
267	newPartialProd := initialPartialProd
268	for i := 0; i < n; i++ {
269		val := values[i]
270		x := domain[i]
271		xField := NewFieldElement(x.Uint64())
272		xQE := QuadraticExtension{xField, ZERO_F}
273		xQEAlgebra := QEAlgebra{xQE, c.ZERO_QE}
274		weight := QuadraticExtension{NewFieldElement(barycentricWeights[i].Uint64()), ZERO_F}
275		term := c.SubExtensionAlgebra(point, xQEAlgebra)
276		weightedVal := c.ScalarMulExtensionAlgebra(weight, val)
277		newEval = c.MulExtensionAlgebra(newEval, term)
278		tmp := c.MulExtensionAlgebra(weightedVal, newPartialProd)
279		newEval = c.AddExtensionAlgebra(newEval, tmp)
280		newPartialProd = c.MulExtensionAlgebra(newPartialProd, term)
281	}
282
283	return newEval, newPartialProd
284}