Fix Poseidon2 arity and visibility (#3)

* fix: apply fixes to allow for new poseidon arities

* fix: expand visibility of apply_mat4

* fix: support all arities % 4 == 0

poseidon2/src/lib.rs

@@ -21,8 +21,6 @@ use rand::distributions::{Distribution, Standard};
 use rand::Rng;
 pub use round_numbers::poseidon2_round_numbers_128;
 
-const SUPPORTED_WIDTHS: [usize; 8] = [2, 3, 4, 8, 12, 16, 20, 24];
-
 /// The Poseidon2 permutation.
 #[derive(Clone, Debug)]
 pub struct Poseidon2<F, MdsLight, Diffusion, const WIDTH: usize, const D: u64> {
@@ -60,7 +58,7 @@ where
         internal_constants: Vec<F>,
         internal_linear_layer: Diffusion,
     ) -> Self {
-        assert!(SUPPORTED_WIDTHS.contains(&WIDTH));
+        assert!(WIDTH == 2 || WIDTH == 3 || WIDTH % 4 == 0);
         Self {
             rounds_f,
             external_constants,

poseidon2/src/matrix.rs

@@ -39,7 +39,7 @@ where
 // [ 1 1 2 3 ]
 // [ 3 1 1 2 ].
 // This is more efficient than the previous matrix.
-fn apply_mat4<AF>(x: &mut [AF; 4])
+pub fn apply_mat4<AF>(x: &mut [AF; 4])
 where
     AF: AbstractField,
 {
@@ -106,7 +106,8 @@ fn mds_light_permutation<AF: AbstractField, MdsPerm4: MdsPermutation<AF, 4>, con
             state[2] += sum;
         }
 
-        4 | 8 | 12 | 16 | 20 | 24 => {
+        _ => {
+            assert!(WIDTH % 4 == 0, "Unsupported width");
             // First, we apply M_4 to each consecutive four elements of the state.
             // In Appendix B's terminology, this replaces each x_i with x_i'.
             for i in (0..WIDTH).step_by(4) {
@@ -136,10 +137,6 @@ fn mds_light_permutation<AF: AbstractField, MdsPerm4: MdsPermutation<AF, 4>, con
                 state[i] += sums[i % 4].clone();
             }
         }
-
-        _ => {
-            panic!("Unsupported width");
-        }
     }
 }
 

symmetric/src/lib.rs

@@ -7,7 +7,7 @@ extern crate alloc;
 mod compression;
 mod hash;
 mod hasher;
-mod permutation;
+pub mod permutation;
 mod serializing_hasher;
 mod sponge;