1+ use std::check::assert;
2+ use std::check::panic;
3+ use std::convert::fe;
4+ use std::convert::int;
5+ use std::convert::expr;
6+ use std::math::ff::inv_field;
7+ use std::prover::eval;
8+
9+ /// An element of the extension field over the fields (BabyBear, Goldilocks, and BN254)
10+ /// relative to the polynomial X^4 - 11 , which is irreducible all of the above fields
11+ /// where Fp4(a0 , a1 , a2 , a3 ) is interpreted as a0 + a1 * X + a2 * X^2 + a3 * X^3
12+ /// T is assumed to either be fe, expr or any other object whose algebraic operations
13+ /// are compatible with fe.
14+ enum Fp4<T> {
15+ Fp4(T, T, T, T)
16+ }
17+
18+ /// Converts a base field element to the extension field
19+ let <T: FromLiteral> from_base: T -> Fp4<T> = |x| Fp4::Fp4(x, 0 , 0 , 0 );
20+
21+ /// Addition for extension field
22+ let <T: Add > add_ext: Fp4<T>, Fp4<T> -> Fp4<T> = |a, b| match (a, b) {
23+ (Fp4::Fp4(a0 , a1 , a2 , a3 ), Fp4::Fp4(b0, b1, b2, b3)) => Fp4::Fp4(
24+ a0 + b0,
25+ a1 + b1,
26+ a2 + b2,
27+ a3 + b3
28+ )
29+ };
30+
31+ /// Subtraction for extension field
32+ let <T: Sub > sub_ext: Fp4<T>, Fp4<T> -> Fp4<T> = |a, b| match (a, b) {
33+ (Fp4::Fp4(a0 , a1 , a2 , a3 ), Fp4::Fp4(b0, b1, b2, b3)) => Fp4::Fp4(
34+ a0 - b0,
35+ a1 - b1,
36+ a2 - b2,
37+ a3 - b3
38+ )
39+ };
40+
41+ /// Multiplication for the extension field:
42+ /// Multiply out the polynomial representations, and then reduce modulo
43+ /// `x^4 - B`, which means powers >= 4 get shifted back 4 and
44+ /// multiplied by `beta`.
45+ ///
46+ /// Multiplication modulo the polynomial x^4 - 11 . We'll use the fact
47+ /// that x^4 == 11 (mod x^4 - 11 ), so:
48+ /// (a0 + a1 * x + a2 * x^2 + a3 * x^3) * (b0 + b1 * x + b2 * x^2 + b3 * x^3) =
49+ /// a0 * b0 + NBETA * (a1 * b3 + a2 * b2 + a3 * b1)
50+ /// + (a0 * b1 + a1 * b0 + NBETA * (a2 * b3 + a3 * b2)) * X
51+ /// + (a0 * b2 + a1 * b1 + a2 * b0 + NBETA * (a3 * b3)) * X^2
52+ /// + (a0 * b3 + a1 * b2 + a2 * b1 + a3 * b0) * X^3
53+ let <T: Add + FromLiteral + Mul> mul_ext: Fp4<T>, Fp4<T> -> Fp4<T> = |a, b| match (a, b) {
54+ (Fp4::Fp4(a0 , a1 , a2 , a3 ), Fp4::Fp4(b0, b1, b2, b3)) => Fp4::Fp4(
55+ a0 * b0 + 11 * (a1 * b3 + a2 * b2 + a3 * b1),
56+ a0 * b1 + a1 * b0 + 11 * (a2 * b3 + a3 * b2),
57+ a0 * b2 + a1 * b1 + a2 * b0 + 11 * (a3 * b3),
58+ a0 * b3 + a1 * b2 + a2 * b1 + a3 * b0
59+ )
60+ };
61+
62+ /// Inversion for an Fp4 element
63+ /// The inverse of (a0 , a1 , a2 , a3 ) is a point (b0, b1, b2, b3) such that:
64+ /// (a0 + a1 * x + a2 * x^2 + a3 * x^3) (b0 + b1 * x + b2 * x^2 + b3 * x^3) = 1 (mod x^4 - 11)
65+ /// Calculating inverse of z as following
66+ /// a * z = 1, where z is the inverse of a
67+ /// z = 1 / a
68+ /// Multiply each side with a' where a' = a0 - a1 * x + a2 * x^2 - a3 * x^3
69+ /// z = a' / (a * a')
70+ /// Substitute (a * a') = b after multipliying (a * a')
71+ /// b = b0 + b * x^2 (Since the multiplication of a * a' doesn't result x and x^3 parts)
72+ /// By substituting x^4 = 11 , we have
73+ /// b0 = a0 * a0 - 11 * (a1 * (a3 + a3) - a2 * a2);
74+ /// b2 = a0 * (a2 + a2) - a1 * a1 - 11 * (a3 * a3);
75+ /// z = a' / b
76+ /// Multiply each side with b' where b' = b0 - b0 * x^2
77+ /// z = (a' * b') / (b * b')
78+ /// Multiplying (b * b') results c = b0^2 * b2^2 * 11
79+ /// z = (a' * b') / c
80+ /// z = (a' * b') * ('inverse' of c)
81+ /// z = a' * (b0^2 * ic - 11 b2^2 * ic)
82+ /// z = (a0 * b0 - 11 * a2 * b2) * 1
83+ /// + (-1 * a1 * b0 + 11 * a3 * b2) * x
84+ /// + (-1 * a0 * b2 + a2 * b0) * x^2
85+ /// + (a1 * b2 - a3 * b0) * x^3
86+ let inv_ext: Fp4<fe> -> Fp4<fe> = |a| match a {
87+ Fp4::Fp4(a0 , a1 , a2 , a3 ) => {
88+ let b0 = a0 * a0 - 11 * (a1 * (a3 + a3) - a2 * a2);
89+ let b2 = a0 * (a2 + a2) - a1 * a1 - 11 * (a3 * a3);
90+ let c = b0 * b0 - 11 * b2 * b2;
91+ let ic = inv_field(c);
92+ let b_0 = b0 * ic;
93+ let b_2 = b2 * ic;
94+ Fp4::Fp4(
95+ a0 * b_0 - 11 * a2 * b_2,
96+ -1 * a1 * b_0 + 11 * a3 * b_2,
97+ -1 * a0 * b_2 + a2 * b_0,
98+ a1 * b_2 - a3 * b_0
99+ )
100+ }
101+ };
102+
103+ /// Converts an Fp4<expr> into an Fp4<fe>
104+ let eval_ext: Fp4<expr> -> Fp4<fe> = query |a| match a {
105+ Fp4::Fp4(a0 , a1 , a2 , a3 ) => Fp4::Fp4(eval(a0 ), eval(a1 ), eval(a2 ), eval(a3 ))
106+ };
107+
108+ /// Converts an Fp4<fe> into an Fp4<expr>
109+ let expr_ext: Fp4<fe> -> Fp4<expr> = |a| match a {
110+ Fp4::Fp4(a0 , a1 , a2 , a3 ) => Fp4::Fp4(expr(a0 ), expr(a1 ), expr(a2 ), expr(a3 ))
111+ };
112+
113+ /// Checks the equality of two Fp4 elements
114+ let eq_ext: Fp4<fe>, Fp4<fe> -> bool = |a, b| match (a, b) {
115+ (Fp4::Fp4(a0 , a1 , a2 , a3 ), Fp4::Fp4(b0, b1, b2, b3)) => (a0 == b0) && (a1 == b1) && (a2 == b2) && (a3 == b3)
116+ };
117+
118+ /// Returns constraints that two Fp4 elements are equal
119+ let constrain_eq_ext: Fp4<expr>, Fp4<expr> -> Constr[] = |a, b| match (a, b) {
120+ (Fp4::Fp4(a0 , a1 , a2 , a3 ), Fp4::Fp4(b0, b1, b2, b3)) => [a0 = b0, a1 = b1, a2 = b2, a3 = b3]
121+ };
122+
123+ /// Applies the next operator to both components of the extension field element
124+ let next_ext: Fp4<expr> -> Fp4<expr> = |a| match a {
125+ Fp4::Fp4(a0 , a1 , a2 , a3 ) => Fp4::Fp4(a0 ', a1 ', a2 ', a3 ')
126+ };
127+
128+ /// Returns the two components of the extension field element as a tuple
129+ let <T> unpack_ext: Fp4<T> -> (T, T, T, T) = |a| match a {
130+ Fp4::Fp4(a0 , a1 , a2 , a3 ) => (a0 , a1 , a2 , a3 )
131+ };
132+
133+ /// Returns the two components of the extension field element as an array
134+ let <T> unpack_ext_array: Fp4<T> -> T[] = |a| match a {
135+ Fp4::Fp4(a0 , a1 , a2 , a3 ) => [a0 , a1 , a2 , a3 ]
136+ };
137+
138+ mod test {
139+ use super::Fp4;
140+ use super::from_base;
141+ use super::add_ext;
142+ use super::sub_ext;
143+ use super::mul_ext;
144+ use super::inv_ext;
145+ use super::eq_ext;
146+ use std::check::assert;
147+ use std::array::map;
148+
149+ let add = || {
150+ let test_add = |a, b, c| assert(eq_ext(add_ext(a, b), c), || "Wrong addition result" );
151+
152+ // Test adding 0
153+ let _ = test_add(from_base(0 ), from_base(0 ), from_base(0 ));
154+ let _ = test_add(Fp4::Fp4(123 , 1234 , 1 , 2 ), from_base(0 ), Fp4::Fp4(123 , 1234 , 1 , 2 ));
155+ let _ = test_add(from_base(0 ), Fp4::Fp4(123 , 1234 , 123 , 1334 ), Fp4::Fp4(123 , 1234 , 123 , 1334 ));
156+
157+ // Add arbitrary elements
158+ let _ = test_add(Fp4::Fp4(123 , 1234 , 123 , 122 ), Fp4::Fp4(567 , 5678 , 250 , 678 ), Fp4::Fp4(690 , 6912 , 373 , 800 ));
159+ test_add(Fp4::Fp4(-1 , -1 , -1 , -1 ), Fp4::Fp4(3 , 4 , 5 , 6 ), Fp4::Fp4(2 , 3 , 4 , 5 ));
160+
161+ // Add to the modulo
162+ test_add(Fp4::Fp4(-11 , -11 , -11 , -11 ), Fp4::Fp4(0 , 0 , 0 , 0 ), Fp4::Fp4(-11 , -11 , -11 , -11 ));
163+
164+ // p - 1 + 1 = 0
165+ test_add(Fp4::Fp4(-1 , 0 , 0 , 0 ), Fp4::Fp4(1 , 0 , 0 , 0 ), from_base(0 ));
166+ };
167+
168+ let sub = || {
169+ let test_sub = |a, b, c| assert(eq_ext(sub_ext(a, b), c), || "Wrong subtraction result" );
170+
171+ // Test subtracting 0
172+ let _ = test_sub(from_base(0 ), from_base(0 ), from_base(0 ));
173+ let _ = test_sub(Fp4::Fp4(123 , 1234 , 124 , 1235 ), from_base(0 ), Fp4::Fp4(123 , 1234 , 124 , 1235 ));
174+
175+ // Subtract arbitrary elements
176+ let _ = test_sub(Fp4::Fp4(123 , 1234 , 248 , 5000 ), Fp4::Fp4(567 , 5678 , 300 , 2380 ), Fp4::Fp4(123 - 567 , 1234 - 5678 , 248 - 300 , 5000 - 2380 ));
177+ test_sub(Fp4::Fp4(-1 , -1 , 0 , 0 ), Fp4::Fp4(0x78000000, 1 , 0 , 0 ), Fp4::Fp4(-0x78000000 - 1 , -2 , 0 , 0 ))
178+ };
179+
180+ let mul = || {
181+ let test_mul = |a, b, c| assert(eq_ext(mul_ext(a, b), c), || "Wrong multiplication result" );
182+
183+ // Test multiplication by 1
184+ let _ = test_mul(from_base(1 ), from_base(1 ), from_base(1 ));
185+ let _ = test_mul(Fp4::Fp4(123 , 1234 , 280 , 400 ), from_base(1 ), Fp4::Fp4(123 , 1234 , 280 , 400 ));
186+ let _ = test_mul(from_base(1 ), Fp4::Fp4(123 , 1234 , 12 , 15 ), Fp4::Fp4(123 , 1234 , 12 , 15 ));
187+
188+ // Test multiplication by 0
189+ let _ = test_mul(Fp4::Fp4(123 , 1234 , 234 , 500 ), from_base(0 ), from_base(0 ));
190+ let _ = test_mul(from_base(0 ), Fp4::Fp4(123 , 1234 , 33 , 200 ), from_base(0 ));
191+
192+ // Multiply arbitrary elements
193+ test_mul(Fp4::Fp4(1 , 2 , 3 , 4 ), Fp4::Fp4(5 , 6 , 7 , 8 ), Fp4::Fp4(676 , 588 , 386 , 60 ));
194+
195+ // Multiplication with field overflow
196+ //2013265921
197+ test_mul(Fp4::Fp4(-1 , -2 , -3 , -4 ), Fp4::Fp4(-3 , 4 , 4 , 5 ), Fp4::Fp4(-415 , -339 , -223 , -13 ));
198+ };
199+
200+ let inverse = || {
201+ let test_elements = [
202+ from_base(1 ),
203+ Fp4::Fp4(123 , 1234 , 1 , 2 ),
204+ Fp4::Fp4(-1 , 500 , 3 , 5 )
205+ ];
206+
207+ map(test_elements, |x| {
208+ let mul_with_inverse = mul_ext(x, inv_ext(x));
209+
210+ assert(eq_ext(mul_with_inverse, from_base(1 )), || "Should be 1" )
211+ })
212+ };
213+ }
0 commit comments