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