binius_m3/gadgets/hash/keccak/
state.rs

1// Copyright 2025 Irreducible Inc.
2
3use std::{array, fmt};
4
5/// The state of the permutation. Consists of 5x5 lanes.
6///
7/// Indexing is defined as `mod 5`.
8///
9/// See section 1.6 of the [keccak specification] for more details.
10///
11/// [keccak specification]: https://keccak.team/files/Keccak-reference-3.0.pdf
12#[derive(Default, Clone, PartialEq, Eq)]
13pub struct StateMatrix<T> {
14	v: [T; 25],
15}
16
17impl<T> StateMatrix<T> {
18	/// Creates a new state matrix from an array of values stored in row-major order.
19	pub fn from_values(v: [T; 25]) -> Self {
20		Self { v }
21	}
22
23	/// Creates a new state matrix from a function that takes a coordinate and returns a value.
24	pub fn from_fn<F>(mut f: F) -> Self
25	where
26		F: FnMut((usize, usize)) -> T,
27	{
28		Self {
29			v: array::from_fn(|i| f((i % 5, i / 5))),
30		}
31	}
32
33	/// Creates a new state matrix from a function that takes a coordinate and returns a value or
34	/// an error.
35	pub fn try_from_fn<F, E>(mut f: F) -> Result<Self, E>
36	where
37		F: FnMut((usize, usize)) -> Result<T, E>,
38	{
39		Ok(Self {
40			v: array_util::try_from_fn(|i| f((i % 5, i / 5)))?,
41		})
42	}
43
44	/// Returns the raw array of values stored in row-major order.
45	pub fn as_inner(&self) -> &[T; 25] {
46		&self.v
47	}
48
49	/// Returns the raw array of values stored in row-major order.
50	#[cfg(test)]
51	pub fn into_inner(self) -> [T; 25] {
52		self.v
53	}
54}
55
56impl<T> std::ops::Index<(usize, usize)> for StateMatrix<T> {
57	type Output = T;
58	fn index(&self, (x, y): (usize, usize)) -> &T {
59		let x = x % 5;
60		let y = y % 5;
61		&self.v[x + y * 5]
62	}
63}
64
65impl<T> std::ops::IndexMut<(usize, usize)> for StateMatrix<T> {
66	fn index_mut(&mut self, (x, y): (usize, usize)) -> &mut T {
67		let x = x % 5;
68		let y = y % 5;
69		&mut self.v[x + y * 5]
70	}
71}
72
73impl<T: fmt::Debug + fmt::UpperHex> fmt::Debug for StateMatrix<T> {
74	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75		write!(f, "StateMatrix {{")?;
76		if f.alternate() {
77			writeln!(f)?;
78		}
79		// generate a 5x5 matrix of hex values.
80		for y in 0..5 {
81			write!(f, "[")?;
82			if f.alternate() {
83				write!(f, "  ")?;
84			}
85			for x in 0..5 {
86				write!(f, "0x{:02X} ", self[(x, y)])?;
87			}
88			write!(f, "]")?;
89
90			if f.alternate() {
91				writeln!(f)?;
92			}
93		}
94		write!(f, "}}")?;
95		Ok(())
96	}
97}
98
99/// A row of the state. This is basically a slice of the state matrix along the x coordinate.
100///
101/// Indexing is defined as `mod 5`.
102///
103/// See section 1.6 of the [keccak specification] for more details.
104///
105/// [keccak specification]: https://keccak.team/files/Keccak-reference-3.0.pdf
106#[derive(Clone)]
107pub struct StateRow<T>([T; 5]);
108
109impl<T> std::ops::Index<usize> for StateRow<T> {
110	type Output = T;
111	fn index(&self, i: usize) -> &T {
112		&self.0[i % 5]
113	}
114}
115
116impl<T> std::ops::IndexMut<usize> for StateRow<T> {
117	fn index_mut(&mut self, i: usize) -> &mut T {
118		&mut self.0[i % 5]
119	}
120}
121
122impl<T> StateRow<T> {
123	/// Creates a new state row from a function that takes a coordinate and returns a value.
124	pub fn from_fn<F>(f: F) -> Self
125	where
126		F: FnMut(usize) -> T,
127	{
128		Self(array::from_fn(f))
129	}
130}