feat: add TryFrom<u64> and TryFrom<u128> bounds to the FieldElement

facebook · Feb 6, 2024 · 450da89 · 450da89
1 parent aed1d72
commit 450da89
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 5 deletions.
diff --git a/air/src/proof/context.rs b/air/src/proof/context.rs
@@ -224,7 +224,7 @@ fn bytes_to_element<B: StarkField>(bytes: &[u8]) -> B {
 
     let mut buf = bytes.to_vec();
     buf.resize(B::ELEMENT_BYTES, 0);
-    let element = match B::try_from(&buf) {
+    let element = match B::try_from(buf.as_slice()) {
         Ok(element) => element,
         Err(_) => panic!("element deserialization failed"),
     };

diff --git a/math/src/field/extensions/cubic.rs b/math/src/field/extensions/cubic.rs
@@ -11,8 +11,10 @@ use core::{
     slice,
 };
 use utils::{
-    collections::Vec, string::ToString, AsBytes, ByteReader, ByteWriter, Deserializable,
-    DeserializationError, Randomizable, Serializable, SliceReader,
+    collections::Vec,
+    string::{String, ToString},
+    AsBytes, ByteReader, ByteWriter, Deserializable, DeserializationError, Randomizable,
+    Serializable, SliceReader,
 };
 
 #[cfg(feature = "serde")]
@@ -319,6 +321,34 @@ impl<B: ExtensibleField<3>> From<u8> for CubeExtension<B> {
     }
 }
 
+impl<B: ExtensibleField<3>> TryFrom<u64> for CubeExtension<B> {
+    type Error = String;
+
+    fn try_from(value: u64) -> Result<Self, Self::Error> {
+        let felt = B::try_from(value);
+        match felt {
+            Ok(elem) => Ok(Self(elem, B::ZERO, B::ZERO)),
+            Err(_) => Err(format!(
+                "invalid field element: value {value} is greater than or equal to the field modulus"
+            )),
+        }
+    }
+}
+
+impl<B: ExtensibleField<3>> TryFrom<u128> for CubeExtension<B> {
+    type Error = String;
+
+    fn try_from(value: u128) -> Result<Self, Self::Error> {
+        let felt = B::try_from(value);
+        match felt {
+            Ok(elem) => Ok(Self(elem, B::ZERO, B::ZERO)),
+            Err(_) => Err(format!(
+                "invalid field element: value {value} is greater than or equal to the field modulus"
+            )),
+        }
+    }
+}
+
 impl<'a, B: ExtensibleField<3>> TryFrom<&'a [u8]> for CubeExtension<B> {
     type Error = DeserializationError;
 

diff --git a/math/src/field/extensions/quadratic.rs b/math/src/field/extensions/quadratic.rs
@@ -11,8 +11,10 @@ use core::{
     slice,
 };
 use utils::{
-    collections::Vec, string::ToString, AsBytes, ByteReader, ByteWriter, Deserializable,
-    DeserializationError, Randomizable, Serializable, SliceReader,
+    collections::Vec,
+    string::{String, ToString},
+    AsBytes, ByteReader, ByteWriter, Deserializable, DeserializationError, Randomizable,
+    Serializable, SliceReader,
 };
 
 #[cfg(feature = "serde")]
@@ -313,6 +315,34 @@ impl<B: ExtensibleField<2>> From<u8> for QuadExtension<B> {
     }
 }
 
+impl<B: ExtensibleField<2>> TryFrom<u64> for QuadExtension<B> {
+    type Error = String;
+
+    fn try_from(value: u64) -> Result<Self, Self::Error> {
+        let felt = B::try_from(value);
+        match felt {
+            Ok(elem) => Ok(Self(elem, B::ZERO)),
+            Err(_) => Err(format!(
+                "invalid field element: value {value} is greater than or equal to the field modulus"
+            )),
+        }
+    }
+}
+
+impl<B: ExtensibleField<2>> TryFrom<u128> for QuadExtension<B> {
+    type Error = String;
+
+    fn try_from(value: u128) -> Result<Self, Self::Error> {
+        let felt = B::try_from(value);
+        match felt {
+            Ok(elem) => Ok(Self(elem, B::ZERO)),
+            Err(_) => Err(format!(
+                "invalid field element: value {value} is greater than or equal to the field modulus"
+            )),
+        }
+    }
+}
+
 impl<'a, B: ExtensibleField<2>> TryFrom<&'a [u8]> for QuadExtension<B> {
     type Error = DeserializationError;
 

diff --git a/math/src/field/f128/mod.rs b/math/src/field/f128/mod.rs
@@ -355,6 +355,20 @@ impl From<u8> for BaseElement {
     }
 }
 
+impl TryFrom<u128> for BaseElement {
+    type Error = String;
+
+    fn try_from(value: u128) -> Result<Self, Self::Error> {
+        if value >= M {
+            Err(format!(
+                "invalid field element: value {value} is greater than or equal to the field modulus"
+            ))
+        } else {
+            Ok(Self::new(value))
+        }
+    }
+}
+
 impl<'a> TryFrom<&'a [u8]> for BaseElement {
     type Error = String;
 

diff --git a/math/src/field/f62/mod.rs b/math/src/field/f62/mod.rs
@@ -455,6 +455,20 @@ impl TryFrom<u64> for BaseElement {
     }
 }
 
+impl TryFrom<u128> for BaseElement {
+    type Error = String;
+
+    fn try_from(value: u128) -> Result<Self, Self::Error> {
+        if value >= M as u128 {
+            Err(format!(
+                "invalid field element: value {value} is greater than or equal to the field modulus"
+            ))
+        } else {
+            Ok(Self::new(value as u64))
+        }
+    }
+}
+
 impl TryFrom<[u8; 8]> for BaseElement {
     type Error = String;
 

diff --git a/math/src/field/f64/mod.rs b/math/src/field/f64/mod.rs
@@ -550,6 +550,20 @@ impl TryFrom<u64> for BaseElement {
     }
 }
 
+impl TryFrom<u128> for BaseElement {
+    type Error = String;
+
+    fn try_from(value: u128) -> Result<Self, Self::Error> {
+        if value >= M as u128 {
+            Err(format!(
+                "invalid field element: value {value} is greater than or equal to the field modulus"
+            ))
+        } else {
+            Ok(Self::new(value as u64))
+        }
+    }
+}
+
 impl TryFrom<[u8; 8]> for BaseElement {
     type Error = String;
 

diff --git a/math/src/field/traits.rs b/math/src/field/traits.rs
@@ -50,6 +50,8 @@ pub trait FieldElement:
     + From<u16>
     + From<u8>
     + for<'a> TryFrom<&'a [u8]>
+    + TryFrom<u64>
+    + TryFrom<u128>
     + ExtensionOf<<Self as FieldElement>::BaseField>
     + AsBytes
     + Randomizable