Write trace and trace asserts (#542)

This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/stwo/542)
starkware-libs · Mar 27, 2024 · 4e331fc · 4e331fc
2 parents a157d76 + 8a95e86
commit 4e331fc
Show file tree

Hide file tree

Showing 3 changed files with 420 additions and 2 deletions.
diff --git a/src/examples/wide_fibonacci/mod.rs b/src/examples/wide_fibonacci/mod.rs
@@ -1,3 +1,36 @@
 pub mod constraint_eval;
 pub mod structs;
+pub mod trace_asserts;
 pub mod trace_gen;
+
+#[cfg(test)]
+mod tests {
+    use itertools::Itertools;
+    use num_traits::Zero;
+
+    use super::structs::Input;
+    use super::trace_asserts::assert_constraints_on_row;
+    use super::trace_gen::write_trace_row;
+    use crate::core::fields::m31::BaseField;
+
+    fn fill_trace(private_input: &[Input]) -> Vec<Vec<BaseField>> {
+        let zero_vec = vec![BaseField::zero(); private_input.len()];
+        let mut dst = vec![zero_vec; 64];
+        for (offset, input) in private_input.iter().enumerate() {
+            write_trace_row(&mut dst, input, offset);
+        }
+        dst
+    }
+
+    #[test]
+    fn test_wide_fib_trace() {
+        let input = Input {
+            a: BaseField::from_u32_unchecked(0x76),
+            b: BaseField::from_u32_unchecked(0x483),
+        };
+
+        let trace = fill_trace(&[input]);
+        let flat_trace = trace.into_iter().flatten().collect_vec();
+        assert_constraints_on_row(&flat_trace);
+    }
+}
diff --git a/src/examples/wide_fibonacci/trace_asserts.rs b/src/examples/wide_fibonacci/trace_asserts.rs
@@ -0,0 +1,254 @@
+use num_traits::Zero;
+
+use crate::core::fields::m31::BaseField;
+
+pub fn assert_constraints_on_row(row: &[BaseField]) {
+    assert_eq!(
+        (row[2] - ((row[0] * row[0]) + (row[1] * row[1]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[3] - ((row[1] * row[1]) + (row[2] * row[2]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[4] - ((row[2] * row[2]) + (row[3] * row[3]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[5] - ((row[3] * row[3]) + (row[4] * row[4]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[6] - ((row[4] * row[4]) + (row[5] * row[5]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[7] - ((row[5] * row[5]) + (row[6] * row[6]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[8] - ((row[6] * row[6]) + (row[7] * row[7]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[9] - ((row[7] * row[7]) + (row[8] * row[8]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[10] - ((row[8] * row[8]) + (row[9] * row[9]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[11] - ((row[9] * row[9]) + (row[10] * row[10]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[12] - ((row[10] * row[10]) + (row[11] * row[11]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[13] - ((row[11] * row[11]) + (row[12] * row[12]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[14] - ((row[12] * row[12]) + (row[13] * row[13]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[15] - ((row[13] * row[13]) + (row[14] * row[14]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[16] - ((row[14] * row[14]) + (row[15] * row[15]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[17] - ((row[15] * row[15]) + (row[16] * row[16]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[18] - ((row[16] * row[16]) + (row[17] * row[17]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[19] - ((row[17] * row[17]) + (row[18] * row[18]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[20] - ((row[18] * row[18]) + (row[19] * row[19]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[21] - ((row[19] * row[19]) + (row[20] * row[20]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[22] - ((row[20] * row[20]) + (row[21] * row[21]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[23] - ((row[21] * row[21]) + (row[22] * row[22]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[24] - ((row[22] * row[22]) + (row[23] * row[23]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[25] - ((row[23] * row[23]) + (row[24] * row[24]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[26] - ((row[24] * row[24]) + (row[25] * row[25]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[27] - ((row[25] * row[25]) + (row[26] * row[26]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[28] - ((row[26] * row[26]) + (row[27] * row[27]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[29] - ((row[27] * row[27]) + (row[28] * row[28]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[30] - ((row[28] * row[28]) + (row[29] * row[29]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[31] - ((row[29] * row[29]) + (row[30] * row[30]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[32] - ((row[30] * row[30]) + (row[31] * row[31]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[33] - ((row[31] * row[31]) + (row[32] * row[32]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[34] - ((row[32] * row[32]) + (row[33] * row[33]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[35] - ((row[33] * row[33]) + (row[34] * row[34]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[36] - ((row[34] * row[34]) + (row[35] * row[35]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[37] - ((row[35] * row[35]) + (row[36] * row[36]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[38] - ((row[36] * row[36]) + (row[37] * row[37]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[39] - ((row[37] * row[37]) + (row[38] * row[38]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[40] - ((row[38] * row[38]) + (row[39] * row[39]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[41] - ((row[39] * row[39]) + (row[40] * row[40]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[42] - ((row[40] * row[40]) + (row[41] * row[41]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[43] - ((row[41] * row[41]) + (row[42] * row[42]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[44] - ((row[42] * row[42]) + (row[43] * row[43]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[45] - ((row[43] * row[43]) + (row[44] * row[44]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[46] - ((row[44] * row[44]) + (row[45] * row[45]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[47] - ((row[45] * row[45]) + (row[46] * row[46]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[48] - ((row[46] * row[46]) + (row[47] * row[47]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[49] - ((row[47] * row[47]) + (row[48] * row[48]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[50] - ((row[48] * row[48]) + (row[49] * row[49]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[51] - ((row[49] * row[49]) + (row[50] * row[50]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[52] - ((row[50] * row[50]) + (row[51] * row[51]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[53] - ((row[51] * row[51]) + (row[52] * row[52]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[54] - ((row[52] * row[52]) + (row[53] * row[53]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[55] - ((row[53] * row[53]) + (row[54] * row[54]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[56] - ((row[54] * row[54]) + (row[55] * row[55]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[57] - ((row[55] * row[55]) + (row[56] * row[56]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[58] - ((row[56] * row[56]) + (row[57] * row[57]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[59] - ((row[57] * row[57]) + (row[58] * row[58]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[60] - ((row[58] * row[58]) + (row[59] * row[59]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[61] - ((row[59] * row[59]) + (row[60] * row[60]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[62] - ((row[60] * row[60]) + (row[61] * row[61]))),
+        BaseField::zero()
+    );
+    assert_eq!(
+        (row[63] - ((row[61] * row[61]) + (row[62] * row[62]))),
+        BaseField::zero()
+    );
+}