From 1e606326cf2b918fdd4fc1764ff5a754d4f98ec6 Mon Sep 17 00:00:00 2001
From: Matej Penciak <matej.penciak@gmail.com>
Date: Thu, 30 May 2024 17:24:37 -0400
Subject: [PATCH] more progress

---
 src/poseidon/air.rs    |  3 +--
 src/poseidon/config.rs | 20 +++++++++++++++++---
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/poseidon/air.rs b/src/poseidon/air.rs
index d4013309..e525ddc2 100644
--- a/src/poseidon/air.rs
+++ b/src/poseidon/air.rs
@@ -160,8 +160,7 @@ where
         {
             // Use a simple matrix multiplication as the permutation.
             let mut state = sbox_result;
-            let matmul_constants: Array<<<AB as AirBuilder>::Expr as AbstractField>::F, C::WIDTH> =
-                C::matrix_diag().map(|u| <AB::Expr as AbstractField>::F::from_f(u));
+            let matmul_constants = C::matrix_diag().map(AB::F::from);
             matmul_generic(&mut state, &matmul_constants);
 
             for (state, &output_expected) in zip(state, &local.output) {
diff --git a/src/poseidon/config.rs b/src/poseidon/config.rs
index b66ebb5c..45a7ce76 100644
--- a/src/poseidon/config.rs
+++ b/src/poseidon/config.rs
@@ -1,8 +1,9 @@
 use hybrid_array::{typenum::*, Array, ArraySize};
+use itertools::Itertools;
 use p3_baby_bear::BabyBear;
-use p3_field::Field;
+use p3_field::{AbstractField, Field};
 
-use super::constants::MATRIX_DIAG_4_BABYBEAR;
+use super::constants::{FULL_RC_4_8, MATRIX_DIAG_4_BABYBEAR, PART_RC_4_21};
 // use p3_baby_bear::BabyBear;
 
 trait ConstantsProvided {}
@@ -41,7 +42,20 @@ impl PoseidonConfig for BabyBearConfig4 {
     }
 
     fn round_constants() -> impl IntoIterator<Item = &'static Array<Self::F, Self::WIDTH>> {
-        todo!()
+        let first_half = FULL_RC_4_8
+            .iter()
+            .map(|u| Array::from_slice(u).as_ref())
+            .take(4);
+        let second_half = FULL_RC_4_8
+            .iter()
+            .map(|u| Array::from_slice(u).as_ref())
+            .skip(4);
+
+        let partial_round_constants = PART_RC_4_21
+            .into_iter()
+            .map(|v| Array::from_fn(|i| if i == 0 { v } else { BabyBear::zero() }).as_ref());
+
+        first_half.chain(partial_round_constants).chain(second_half)
     }
 }