Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

feat: replace stdlib poseidon implementation with optimized version #5122

Merged
merged 9 commits into from
May 28, 2024
Merged

feat: replace stdlib poseidon implementation with optimized version #5122

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
c82ba2b
feat: replace stdlib poseidon implementation with optimized version
TomAFrench May 28, 2024
41fb722
chore: fix test
TomAFrench May 28, 2024
2c43892
Merge branch 'master' into tf/zac-poseidon
TomAFrench May 28, 2024
2a0866a
Apply suggestions from code review
TomAFrench May 28, 2024
599713a
chore: add documentation, remove unnecessary generic and refactor
TomAFrench May 28, 2024
950d3f1
chore: align generic for state length with filecoin notation
TomAFrench May 28, 2024
7b60214
chore: refactor
TomAFrench May 28, 2024
18f5169
chore: refactor
TomAFrench May 28, 2024
074da71
chore: fmt
TomAFrench May 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
198 changes: 145 additions & 53 deletions noir_stdlib/src/hash/poseidon.nr
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,66 +3,134 @@ use crate::field::modulus_num_bits;
use crate::hash::Hasher;
use crate::default::Default;

struct PoseidonConfig<M,N> {
t: Field, // Width, i.e. state size
rf: u8, // Number of full rounds; should be even
rp: u8, // Number of partial rounds
alpha: Field, // S-box power; depends on the underlying field
ark: [Field; M], // Additive round keys
mds: [Field; N] // MDS Matrix in row-major order
// A config struct defining the parameters of the Poseidon instance to use.
//
// A thorough writeup of this method (along with an unoptimized method) can be found at: https://spec.filecoin.io/algorithms/crypto/poseidon/
struct PoseidonConfig<T, N, X> {
// State width, should be equal to `T`
t: Field,
// Number of full rounds. should be even
rf: u8,
// Number of partial rounds
rp: u8,
// S-box power; depends on the underlying field
alpha: Field,
// The round constants for the
round_constants: [Field; N],
// The MDS matrix for the Poseidon instance
mds: [[Field; T]; T],
// An MDS matrix which also applies a basis transformation which allows
// sparse matrices to be used for the partial rounds.
//
// This should be applied instead of `mds` in the final full round.
presparse_mds: [[Field; T]; T],
// A set of sparse matrices used for MDS mixing for the partial rounds.
sparse_mds: [Field; X],
}

pub fn config<M, N>(
pub fn config<T, N, X>(
t: Field,
rf: u8,
rp: u8,
alpha: Field,
ark: [Field; M],
mds: [Field; N]
) -> PoseidonConfig<M,N> {
round_constants: [Field; N],
mds: [[Field; T]; T],
presparse_mds: [[Field; T]; T],
sparse_mds: [Field; X]
) -> PoseidonConfig<T, N, X> {
// Input checks
let mul = crate::wrapping_mul(t as u8, (rf + rp));
assert(mul == ark.len() as u8);
assert(t * t == mds.len() as Field);
assert_eq(rf & 1, 0);
assert_eq((t as u8) * rf + rp, N);
assert_eq(t, T);
assert(alpha != 0);

PoseidonConfig { t, rf, rp, alpha, ark, mds }
PoseidonConfig { t, rf, rp, alpha, round_constants, mds, presparse_mds, sparse_mds }
}
// General Poseidon permutation on elements of type Field
fn permute<M, N, O>(pos_conf: PoseidonConfig<M, N>, mut state: [Field; O]) -> [Field; O] {
let PoseidonConfig {t, rf, rp, alpha, ark, mds} = pos_conf;

assert(t == state.len() as Field);
pub fn permute<T, N, X>(pos_conf: PoseidonConfig<T, N, X>, mut state: [Field; T]) -> [Field; T] {
let PoseidonConfig {t, rf, rp, alpha, round_constants, mds, presparse_mds, sparse_mds } = pos_conf;

let mut count = 0;
// for r in 0..rf + rp
for r in 0..(ark.len() / state.len()) {
for i in 0..state.len() {
state[i] = state[i] + ark[count + i];
} // Shift by round constants
for i in 0..state.len() {
state[i] += round_constants[i];
}

for _r in 0..rf / 2 - 1 {
state = sigma(state);
for i in 0..T {
state[i] += round_constants[T * (_r + 1) as u64 + i];
}
state = apply_matrix(mds, state);
}

state = sigma(state);
for i in 0..T {
state[i] += round_constants[T * (rf / 2) as u64 + i];
}
state = apply_matrix(presparse_mds, state);

for i in 0..T {
crate::as_witness(state[i]);
}

for _r in 0..rp {
state[0] = state[0].pow_32(alpha);
// Check whether we are in a full round
if (r as u8 < rf / 2) | (r as u8 >= rf / 2 + rp) {
for i in 1..state.len() {
state[i] = state[i].pow_32(alpha);
state[0] += round_constants[(rf/2 + 1) as u64 * T + _r as u64];
crate::as_witness(state[0]);
{
let mut newState0 = 0;
for j in 0..T {
newState0 += sparse_mds[(T * 2 - 1) * _r as u64 + j as u64] * state[j];
}
for k in 1..T {
state[k] += state[0] * sparse_mds[(t * 2 - 1) as u64 * _r as u64 + T + k - 1];
}
state[0] = newState0;

if (_r & 1 == 0) {
for k in 1..T {
crate::as_witness(state[k]);
}
}
}
}

state = apply_matrix(mds, state); // Apply MDS matrix
count = count + t as u64;
for _r in 0..rf / 2 - 1 {
state = sigma(state);
for i in 0..state.len() {
state[i] += round_constants[(rf/2+1) as u64 * T + rp as u64 + (_r as u64) * T + i];
}
state = apply_matrix(mds, state);
}

state = sigma(state);
state = apply_matrix(mds, state);

state
}
// Absorption. Fully absorbs input message.
fn absorb<M, N, O, P>(
pos_conf: PoseidonConfig<M, N>,
mut state: [Field; O], // Initial state; usually [0; O]
rate: Field, // Rate
capacity: Field, // Capacity; usually 1
msg: [Field; P]
) -> [Field; O] {
assert(pos_conf.t == rate + capacity);

// Performs matrix multiplication on a vector
fn apply_matrix<N>(matrix: [[Field; N]; N], vec: [Field; N]) -> [Field; N] {
let mut out = [0; N];

for i in 0..N {
for j in 0..N {
out[i] += vec[j] * matrix[j][i];
}
}

out
}

// Corresponding absorption.
fn absorb<T, N, X, O>(
pos_conf: PoseidonConfig<T, N, X>,
// Initial state; usually [0; O]
mut state: [Field; T],
rate: Field,
capacity: Field,
msg: [Field; O] // Arbitrary length message
) -> [Field; T] {
assert_eq(pos_conf.t, rate + capacity);

let mut i = 0;

Expand All @@ -83,25 +151,24 @@ fn absorb<M, N, O, P>(

state
}

fn sigma<O>(x: [Field; O]) -> [Field; O] {
let mut y = x;
for i in 0..O {
let t = y[i];
let tt = t * t;
let tttt = tt * tt;
y[i] *= tttt;
}
y
}

// Check security of sponge instantiation
fn check_security(rate: Field, width: Field, security: Field) -> bool {
let n = modulus_num_bits();

((n - 1) as Field * (width - rate) / 2) as u8 > security as u8
}
// A*x where A is an n x n matrix in row-major order and x an n-vector
fn apply_matrix<N, M>(a: [Field; M], x: [Field; N]) -> [Field; N] {
let mut y = x;

for i in 0..x.len() {
y[i] = 0;
for j in 0..x.len() {
y[i] = y[i] + a[x.len()*i + j]* x[j];
}
}

y
}

struct PoseidonHasher{
_state: [Field],
Expand Down Expand Up @@ -174,3 +241,28 @@ impl Default for PoseidonHasher{
}
}
}

mod poseidon_tests {
use crate::hash::poseidon;

#[test]
fn reference_impl_test_vectors() {
// hardcoded test vectors from https://extgit.iaik.tugraz.at/krypto/hadeshash/-/blob/master/code/test_vectors.txt
{
let mut state = [0, 1, 2];
let mut expected = [
0x115cc0f5e7d690413df64c6b9662e9cf2a3617f2743245519e19607a4417189a, 0x0fca49b798923ab0239de1c9e7a4a9a2210312b6a2f616d18b5a87f9b628ae29, 0x0e7ae82e40091e63cbd4f16a6d16310b3729d4b6e138fcf54110e2867045a30c
];
assert_eq(expected, poseidon::bn254::perm::x5_3(state), "Failed to reproduce output for [0, 1, 2]");
}
{
let mut state = [0, 1, 2, 3, 4];
let mut expected = [
0x299c867db6c1fdd79dcefa40e4510b9837e60ebb1ce0663dbaa525df65250465, 0x1148aaef609aa338b27dafd89bb98862d8bb2b429aceac47d86206154ffe053d, 0x24febb87fed7462e23f6665ff9a0111f4044c38ee1672c1ac6b0637d34f24907, 0x0eb08f6d809668a981c186beaf6110060707059576406b248e5d9cf6e78b3d3e, 0x07748bc6877c9b82c8b98666ee9d0626ec7f5be4205f79ee8528ef1c4a376fc7
];
assert_eq(
expected, poseidon::bn254::perm::x5_5(state), "Failed to reproduce output for [0, 1, 2, 3, 4]"
);
}
}
}
85 changes: 3 additions & 82 deletions noir_stdlib/src/hash/poseidon/bn254.nr
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,93 +2,14 @@
mod perm;
mod consts;

use crate::hash::poseidon::PoseidonConfig;
use crate::hash::poseidon::apply_matrix;
// Optimised permutation for this particular field; uses hardcoded rf and rp values,
// which should agree with those in pos_conf.
#[field(bn254)]
pub fn permute<M, N, O>(pos_conf: PoseidonConfig<M, N>, mut state: [Field; O]) -> [Field; O] {
vezenovm marked this conversation as resolved.
Show resolved Hide resolved
let PoseidonConfig {t, rf: config_rf, rp: config_rp, alpha, ark, mds} = pos_conf;
let rf: u8 = 8;
let rp: u8 = [56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68][state.len() - 2];

assert(t == state.len() as Field);
assert(rf == config_rf);
assert(rp == config_rp);

let mut count = 0;
// First half of full rounds
for _r in 0..rf / 2 {
for i in 0..state.len() {
state[i] = state[i] + ark[count + i];
} // Shift by round constants
for i in 0..state.len() {
state[i] = state[i].pow_32(alpha);
}

state = apply_matrix(mds, state); // Apply MDS matrix
count = count + t as u64;
}
// Partial rounds
for _r in 0..rp {
for i in 0..state.len() {
state[i] = state[i] + ark[count + i];
} // Shift by round constants
state[0] = state[0].pow_32(alpha);

state = apply_matrix(mds, state); // Apply MDS matrix
count = count + t as u64;
}
// Second half of full rounds
for _r in 0..rf / 2 {
for i in 0..state.len() {
state[i] = state[i] + ark[count + i];
} // Shift by round constants
for i in 0..state.len() {
state[i] = state[i].pow_32(alpha);
}

state = apply_matrix(mds, state); // Apply MDS matrix
count = count + t as u64;
}

state
}
// Corresponding absorption.
#[field(bn254)]
fn absorb<M, N, O, P>(
pos_conf: PoseidonConfig<M, N>,
mut state: [Field; O], // Initial state; usually [0; O]
rate: Field, // Rate
capacity: Field, // Capacity; usually 1
msg: [Field; P] // Arbitrary length message
) -> [Field; O] {
assert(pos_conf.t == rate + capacity);

let mut i = 0;

for k in 0..msg.len() {
// Add current block to state
state[capacity + i] += msg[k];
i = i+1;
// Enough to absorb
if i == rate {
state = permute(pos_conf, state);
i = 0;
}
}
// If we have one more block to permute
if i != 0 {
state = permute(pos_conf, state);
}

state
}
use crate::hash::poseidon::{PoseidonConfig, absorb};

// Variable-length Poseidon-128 sponge as suggested in second bullet point of §3 of https://eprint.iacr.org/2019/458.pdf
#[field(bn254)]
pub fn sponge<N>(msg: [Field; N]) -> Field {
absorb(consts::x5_5_config(), [0; 5], 4, 1, msg)[1]
}

// Various instances of the Poseidon hash function
// Consistent with Circom's implementation
pub fn hash_1(input: [Field; 1]) -> Field {
Expand Down
Loading
Loading