Add Fp to extension field implementation (#2427)

In our PIL std lib, we have an `Ext` that abstracts over `Fp2` and
`Fp4`. But in cases where we don't need any extension field (e.g.
BN254), we wrapped base field elements as `Fp2` anyway. This is a bit
weird and can lead to issues when the optimizer removes the unused
dimension (e.g., it might still be referenced by a prover function or
witgen annotation). I had this problem in #2426.

Now, there is an `Ext<T>::Fp` variant that simply wraps `T`.

Author: georgwiese

std/math/extension_field.asm

@@ -17,66 +17,75 @@ let required_extension_size: -> int = || match known_field() {
     None => panic("The permutation/lookup argument is not implemented for the current field!")
 };
 
-/// Wrapper around Fp2 and Fp4 to abstract which extension field is used.
+/// Wrapper around T, Fp2<T> and Fp4<T> to abstract which extension field is used (if any).
 /// Once PIL supports traits, we can remove this type and the functions below.
 enum Ext<T> {
+    Fp(T),
     Fp2(std::math::fp2::Fp2<T>),
     Fp4(std::math::fp4::Fp4<T>)
 }
 
 let<T: Add> add_ext: Ext<T>, Ext<T> -> Ext<T> = |a, b| match (a, b) {
+    (Ext::Fp(aa), Ext::Fp(bb)) => Ext::Fp(aa + bb),
     (Ext::Fp2(aa), Ext::Fp2(bb)) => Ext::Fp2(std::math::fp2::add_ext(aa, bb)),
     (Ext::Fp4(aa), Ext::Fp4(bb)) => Ext::Fp4(std::math::fp4::add_ext(aa, bb)),
     _ => panic("Operands have different types")
 };
 
 let<T: Sub> sub_ext: Ext<T>, Ext<T> -> Ext<T> = |a, b| match (a, b) {
+    (Ext::Fp(aa), Ext::Fp(bb)) => Ext::Fp(aa - bb),
     (Ext::Fp2(aa), Ext::Fp2(bb)) => Ext::Fp2(std::math::fp2::sub_ext(aa, bb)),
     (Ext::Fp4(aa), Ext::Fp4(bb)) => Ext::Fp4(std::math::fp4::sub_ext(aa, bb)),
     _ => panic("Operands have different types")
 };
 
 let<T: Add + FromLiteral + Mul> mul_ext: Ext<T>, Ext<T> -> Ext<T> = |a, b| match (a, b) {
+    (Ext::Fp(aa), Ext::Fp(bb)) => Ext::Fp(aa * bb),
     (Ext::Fp2(aa), Ext::Fp2(bb)) => Ext::Fp2(std::math::fp2::mul_ext(aa, bb)),
     (Ext::Fp4(aa), Ext::Fp4(bb)) => Ext::Fp4(std::math::fp4::mul_ext(aa, bb)),
     _ => panic("Operands have different types")
 };
 
 let eval_ext: Ext<expr> -> Ext<fe> = query |a| match a {
+    Ext::Fp(aa) => Ext::Fp(std::prover::eval(aa)),
     Ext::Fp2(aa) => Ext::Fp2(std::math::fp2::eval_ext(aa)),
     Ext::Fp4(aa) => Ext::Fp4(std::math::fp4::eval_ext(aa)),
 };
 
 let inv_ext: Ext<fe> -> Ext<fe> = query |a| match a {
+    Ext::Fp(aa) => Ext::Fp(std::math::ff::inv_field(aa)),
     Ext::Fp2(aa) => Ext::Fp2(std::math::fp2::inv_ext(aa)),
     Ext::Fp4(aa) => Ext::Fp4(std::math::fp4::inv_ext(aa)),
 };
 
 let<T> unpack_ext_array: Ext<T> -> T[] = |a| match a {
+    Ext::Fp(aa) => [aa],
     Ext::Fp2(aa) => std::math::fp2::unpack_ext_array(aa),
     Ext::Fp4(aa) => std::math::fp4::unpack_ext_array(aa),
 };
 
 let next_ext: Ext<expr> -> Ext<expr> = |a| match a {
+    Ext::Fp(aa) => Ext::Fp(aa'),
     Ext::Fp2(aa) => Ext::Fp2(std::math::fp2::next_ext(aa)),
     Ext::Fp4(aa) => Ext::Fp4(std::math::fp4::next_ext(aa)),
 };
 
 let<T: FromLiteral> from_base: T -> Ext<T> = |x| match required_extension_size() {
-    1 => Ext::Fp2(std::math::fp2::from_base(x)),
+    1 => Ext::Fp(x),
     2 => Ext::Fp2(std::math::fp2::from_base(x)),
     4 => Ext::Fp4(std::math::fp4::from_base(x)),
     _ => panic("Expected 1, 2, or 4")
 };
 
-let<T: FromLiteral> from_array: T[] -> Ext<T> = |arr| match len(arr) {
-    1 => Ext::Fp2(std::math::fp2::from_array(arr)),
+let<T> from_array: T[] -> Ext<T> = |arr| match len(arr) {
+    1 => Ext::Fp(arr[0]),
     2 => Ext::Fp2(std::math::fp2::from_array(arr)),
     4 => Ext::Fp4(std::math::fp4::Fp4::Fp4(arr[0], arr[1], arr[2], arr[3])),
     _ => panic("Expected 1, 2, or 4")
 };
 
 let constrain_eq_ext: Ext<expr>, Ext<expr> -> Constr[] = |a, b| match (a, b) {
+    (Ext::Fp(aa), Ext::Fp(bb)) => [aa = bb],
     (Ext::Fp2(aa), Ext::Fp2(bb)) => std::math::fp2::constrain_eq_ext(aa, bb),
     (Ext::Fp4(aa), Ext::Fp4(bb)) => std::math::fp4::constrain_eq_ext(aa, bb),
     _ => panic("Operands have different types")

std/math/fp2.asm

@@ -141,13 +141,9 @@ let<T> unpack_ext_array: Fp2<T> -> T[] = |a| match a {
 };
 
 /// Constructs an extension field element `a0 + a1 * X` from either `[a0, a1]` or `[a0]` (setting `a1`to zero in that case)
-let<T: FromLiteral> from_array: T[] -> Fp2<T> = |arr| match len(arr) {
-    1 => {
-        let _ = assert(!needs_extension(), || "The field is too small and needs to move to the extension field. Pass two elements instead!");
-        from_base(arr[0])
-    },
+let<T> from_array: T[] -> Fp2<T> = |arr| match len(arr) {
     2 => Fp2::Fp2(arr[0], arr[1]),
-    _ => panic("Expected array of length 1 or 2")
+    _ => panic("Expected array of length 2")
 };
 
 mod test {

std/protocols/fingerprint.asm

@@ -37,6 +37,10 @@ let fingerprint_inter: expr[], Ext<expr> -> Ext<expr> = |expr_array, alpha| if l
 
     // Recursively compute the fingerprint as expr_array[0] + alpha * fingerprint(expr_array[1:], alpha)
     let intermediate_fingerprint = match fingerprint_inter(array::sub_array(expr_array, 1, len(expr_array) - 1), alpha) {
+        Ext::Fp(a) => {
+            let intermediate_fingerprint: inter = a;
+            Ext::Fp(intermediate_fingerprint)
+        },
         Ext::Fp2(std::math::fp2::Fp2::Fp2(a0, a1)) => {
             let intermediate_fingerprint_0: inter = a0;
             let intermediate_fingerprint_1: inter = a1;
@@ -70,6 +74,7 @@ mod test {
     let assert_fingerprint_equal: fe[], expr, fe -> () = query |tuple, challenge, expected| {
         let result = fingerprint(tuple, from_base(challenge));
         match result {
+            Ext::Fp(actual) => assert(expected == actual, || "expected != actual"),
             Ext::Fp2(std::math::fp2::Fp2::Fp2(actual, zero)) => {
                 assert(zero == 0, || "Returned an extension field element");
                 assert(expected == actual, || "expected != actual");