├── .gitignore ├── Cargo.toml ├── LICENSE └── src ├── main.rs ├── poseidon_chip.rs └── poseidon_wrapped.rs /.gitignore: -------------------------------------------------------------------------------- 1 | /target 2 | **/*.rs.bk 3 | Cargo.lock 4 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "halo2-merkle" 3 | version = "0.1.0" 4 | edition = "2018" 5 | 6 | [dependencies] 7 | ff = "0.9" 8 | halo2 = { git = "https://github.com/zcash/halo2", branch = "main" } 9 | lazy_static = "1.4.0" 10 | rand = "0.8" 11 | rand_chacha = "0.3.0" 12 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 DrPeterVanNostrand 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | // The constraint system matrix for an arity-2 Merkle tree of 8 leaves using a mocked hasher (one 2 | // selector/gate `s_hash` and one allocation `digest = (l + GAMMA) * (r + GAMMA)` for a random 3 | // gamma and Merkle left/right inputs `l` and `r`). 4 | 5 | // |-----||------------------|------------------|----------|---------|-------|---------|--------|--------| 6 | // | row || a_col | b_col | c_col | pub_col | s_pub | s_bool | s_swap | s_hash | 7 | // |-----||------------------|------------------|----------|---------|-------|---------|--------|--------| 8 | // | 0 || leaf | elem_1 | cbit_1 | cbit_1 | 1 | 1 | 1 | 0 | 9 | // | 1 || leaf/elem_1 | leaf/elem_1 | digest_1 | | 0 | 0 | 0 | 1 | 10 | // | 2 || digest_1* | elem_2 | cbit_2 | cbit_2 | 1 | 1 | 1 | 0 | 11 | // | 3 || digest_1/elem_2 | digest_1/elem_2 | digest_2 | | 0 | 0 | 0 | 1 | 12 | // | 4 || digest_2* | elem_3 | cbit_3 | cbit_3 | 1 | 1 | 1 | 0 | 13 | // | 5 || digest_2/elem_3 | digest_2/elem_3 | digest_3 | root | 1 | 0 | 0 | 1 | 14 | // |-----||------------------|------------------|----------|---------|-------|---------|--------|--------| 15 | // "*" = copy 16 | 17 | use ff::Field; 18 | use halo2::{ 19 | circuit::{layouter::SingleChipLayouter, Cell, Chip, Layouter}, 20 | dev::{MockProver, VerifyFailure}, 21 | pasta::Fp, 22 | plonk::{ 23 | Advice, Assignment, Circuit, Column, ConstraintSystem, Error, Expression, Instance, 24 | Permutation, Selector, 25 | }, 26 | poly::Rotation, 27 | }; 28 | use lazy_static::lazy_static; 29 | use rand::{thread_rng, Rng, SeedableRng}; 30 | use rand_chacha::ChaCha8Rng; 31 | 32 | // The number of leafs in the Merkle tree. This value can be changed to any power of two. 33 | const N_LEAFS: usize = 8; 34 | const PATH_LEN: usize = N_LEAFS.trailing_zeros() as usize; 35 | const TREE_LAYERS: usize = PATH_LEN + 1; 36 | 37 | // The number of rows used in the constraint system matrix (two rows per path element). 38 | const N_ROWS_USED: usize = 2 * PATH_LEN; 39 | const LAST_ROW: usize = N_ROWS_USED - 1; 40 | 41 | lazy_static! { 42 | static ref GAMMA: Fp = Fp::random(ChaCha8Rng::from_seed([101u8; 32])); 43 | } 44 | 45 | // This serves as a mock hash function because the Poseidon chip has not yet been implemented. 46 | fn mock_hash(a: Fp, b: Fp) -> Fp { 47 | (a + *GAMMA) * (b + *GAMMA) 48 | } 49 | 50 | struct Alloc { 51 | cell: Cell, 52 | value: Fp, 53 | } 54 | 55 | enum MaybeAlloc { 56 | Alloc(Alloc), 57 | Unalloc(Fp), 58 | } 59 | 60 | impl MaybeAlloc { 61 | fn value(&self) -> Fp { 62 | match self { 63 | MaybeAlloc::Alloc(alloc) => alloc.value.clone(), 64 | MaybeAlloc::Unalloc(value) => value.clone(), 65 | } 66 | } 67 | 68 | fn cell(&self) -> Cell { 69 | match self { 70 | MaybeAlloc::Alloc(alloc) => alloc.cell.clone(), 71 | MaybeAlloc::Unalloc(_) => unreachable!(), 72 | } 73 | } 74 | } 75 | 76 | struct MerkleChip { 77 | config: MerkleChipConfig, 78 | } 79 | 80 | #[derive(Clone, Debug)] 81 | struct MerkleChipConfig { 82 | a_col: Column, 83 | b_col: Column, 84 | c_col: Column, 85 | pub_col: Column, 86 | s_pub: Selector, 87 | s_bool: Selector, 88 | s_swap: Selector, 89 | s_hash: Selector, 90 | perm_digest: Permutation, 91 | } 92 | 93 | impl Chip for MerkleChip { 94 | type Config = MerkleChipConfig; 95 | type Loaded = (); 96 | 97 | fn config(&self) -> &Self::Config { 98 | &self.config 99 | } 100 | 101 | fn loaded(&self) -> &Self::Loaded { 102 | &() 103 | } 104 | } 105 | 106 | impl MerkleChip { 107 | fn new(config: MerkleChipConfig) -> Self { 108 | MerkleChip { config } 109 | } 110 | 111 | fn configure(cs: &mut ConstraintSystem) -> MerkleChipConfig { 112 | let a_col = cs.advice_column(); 113 | let b_col = cs.advice_column(); 114 | let c_col = cs.advice_column(); 115 | let pub_col = cs.instance_column(); 116 | 117 | let s_pub = cs.selector(); 118 | let s_bool = cs.selector(); 119 | let s_swap = cs.selector(); 120 | let s_hash = cs.selector(); 121 | 122 | cs.create_gate("public input", |cs| { 123 | let c = cs.query_advice(c_col, Rotation::cur()); 124 | let pi = cs.query_instance(pub_col, Rotation::cur()); 125 | let s_pub = cs.query_selector(s_pub, Rotation::cur()); 126 | s_pub * (c - pi) 127 | }); 128 | 129 | cs.create_gate("boolean constrain", |cs| { 130 | let c = cs.query_advice(c_col, Rotation::cur()); 131 | let s_bool = cs.query_selector(s_bool, Rotation::cur()); 132 | s_bool * c.clone() * (Expression::Constant(Fp::one()) - c) 133 | }); 134 | 135 | // |-------|-------|-------|--------| 136 | // | a_col | b_col | c_col | s_swap | 137 | // |-------|-------|-------|--------| 138 | // | a | b | bit | 1 | 139 | // | l | r | | | 140 | // |-------|-------|-------|--------| 141 | // where: 142 | // bit = 0 ==> l = a, r = b 143 | // bit = 1 ==> l = b, r = a 144 | // 145 | // Choose left gate: 146 | // logic: let l = if bit == 0 { a } else { b } 147 | // poly: bit * (b - a) - (l - a) = 0 148 | // 149 | // Choose right gate: 150 | // logic: let r = if bit == 0 { b } else { a } 151 | // poly: bit * (b - a) - (b - r) = 0 152 | // 153 | // Swap gate = choose left + choose right: 154 | // logic: let (l, r) = if bit == 0 { (a, b) } else { (b, a) } 155 | // poly: bit * (b - a) - (l - a) + bit * (b - a) - (b - r) = 0 156 | // bit * 2 * (b - a) - (l - a) - (b - r) = 0 157 | cs.create_gate("swap", |cs| { 158 | let a = cs.query_advice(a_col, Rotation::cur()); 159 | let b = cs.query_advice(b_col, Rotation::cur()); 160 | let bit = cs.query_advice(c_col, Rotation::cur()); 161 | let s_swap = cs.query_selector(s_swap, Rotation::cur()); 162 | let l = cs.query_advice(a_col, Rotation::next()); 163 | let r = cs.query_advice(b_col, Rotation::next()); 164 | s_swap * ((bit * Fp::from(2) * (b.clone() - a.clone()) - (l - a)) - (b - r)) 165 | }); 166 | 167 | // (l + gamma) * (r + gamma) = digest 168 | cs.create_gate("hash", |cs| { 169 | let l = cs.query_advice(a_col, Rotation::cur()); 170 | let r = cs.query_advice(b_col, Rotation::cur()); 171 | let digest = cs.query_advice(c_col, Rotation::cur()); 172 | let s_hash = cs.query_selector(s_hash, Rotation::cur()); 173 | s_hash * ((l + Expression::Constant(*GAMMA)) * (r + Expression::Constant(*GAMMA)) - digest) 174 | }); 175 | 176 | let perm_digest = Permutation::new(cs, &[c_col.into(), a_col.into()]); 177 | 178 | MerkleChipConfig { 179 | a_col, 180 | b_col, 181 | c_col, 182 | pub_col, 183 | s_pub, 184 | s_bool, 185 | s_swap, 186 | s_hash, 187 | perm_digest, 188 | } 189 | } 190 | 191 | fn hash_leaf_layer( 192 | &self, 193 | layouter: &mut impl Layouter, 194 | leaf: Fp, 195 | path_elem: Fp, 196 | c_bit: Fp, 197 | ) -> Result { 198 | self.hash_layer_inner(layouter, MaybeAlloc::Unalloc(leaf), path_elem, c_bit, 0) 199 | } 200 | 201 | fn hash_non_leaf_layer( 202 | &self, 203 | layouter: &mut impl Layouter, 204 | prev_digest: Alloc, 205 | path_elem: Fp, 206 | c_bit: Fp, 207 | layer: usize, 208 | ) -> Result { 209 | self.hash_layer_inner(layouter, MaybeAlloc::Alloc(prev_digest), path_elem, c_bit, layer) 210 | } 211 | 212 | fn hash_layer_inner( 213 | &self, 214 | layouter: &mut impl Layouter, 215 | leaf_or_digest: MaybeAlloc, 216 | path_elem: Fp, 217 | c_bit: Fp, 218 | layer: usize, 219 | ) -> Result { 220 | let mut digest_alloc: Option = None; 221 | 222 | layouter.assign_region( 223 | || "leaf layer", 224 | |mut region| { 225 | let mut row_offset = 0; 226 | 227 | // Allocate in `a_col` either the leaf or reallocate the previous tree layer's 228 | // calculated digest (stored in the previous row's `c_col`). 229 | let a_value = leaf_or_digest.value(); 230 | 231 | let a_cell = region.assign_advice( 232 | || format!("{} (layer {})", if layer == 0 { "leaf" } else { "a" }, layer), 233 | self.config.a_col, 234 | row_offset, 235 | || Ok(a_value), 236 | )?; 237 | 238 | if layer > 0 { 239 | let prev_digest_cell = leaf_or_digest.cell(); 240 | region.constrain_equal(&self.config.perm_digest, prev_digest_cell, a_cell)?; 241 | } 242 | 243 | // Allocate private inputs for this tree layer's path element and challenge bit (in 244 | // columns `b_col` and `c_col` respectively). Expose the challenge bit as a public 245 | // input. 246 | let _elem_cell = region.assign_advice( 247 | || format!("path elem (layer {})", layer), 248 | self.config.b_col, 249 | row_offset, 250 | || Ok(path_elem), 251 | )?; 252 | 253 | let _c_bit_cell = region.assign_advice( 254 | || format!("challenge bit (layer {})", layer), 255 | self.config.c_col, 256 | row_offset, 257 | || Ok(c_bit), 258 | )?; 259 | 260 | // Expose the challenge bit as a public input. 261 | self.config.s_pub.enable(&mut region, row_offset)?; 262 | 263 | // Boolean constrain the challenge bit. 264 | self.config.s_bool.enable(&mut region, row_offset)?; 265 | 266 | // Enable the "swap" gate to ensure the correct order of the Merkle hash inputs. 267 | self.config.s_swap.enable(&mut region, row_offset)?; 268 | 269 | // In the next row, allocate the correctly ordered Merkle hash inputs, calculated digest, and 270 | // enable the "hash" gate. If this is the last tree layer, expose the calculated 271 | // digest as a public input for the tree's root. 272 | row_offset += 1; 273 | 274 | let (preimg_l_value, preimg_r_value): (Fp, Fp) = if c_bit == Fp::zero() { 275 | (a_value, path_elem) 276 | } else { 277 | (path_elem, a_value) 278 | }; 279 | 280 | let _preimg_l_cell = region.assign_advice( 281 | || format!("preimg_l (layer {})", layer), 282 | self.config.a_col, 283 | row_offset, 284 | || Ok(preimg_l_value), 285 | )?; 286 | 287 | let _preimg_r_cell = region.assign_advice( 288 | || format!("preimage right (layer {})", layer), 289 | self.config.b_col, 290 | row_offset, 291 | || Ok(preimg_r_value), 292 | )?; 293 | 294 | let digest_value = mock_hash(preimg_l_value, preimg_r_value); 295 | 296 | let digest_cell = region.assign_advice( 297 | || format!("digest (layer {})", layer), 298 | self.config.c_col, 299 | row_offset, 300 | || Ok(digest_value), 301 | )?; 302 | 303 | digest_alloc = Some(Alloc { 304 | cell: digest_cell, 305 | value: digest_value, 306 | }); 307 | 308 | self.config.s_hash.enable(&mut region, row_offset)?; 309 | 310 | // If the calculated digest is the tree's root, expose it as a public input. 311 | let digest_is_root = layer == PATH_LEN - 1; 312 | if digest_is_root { 313 | self.config.s_pub.enable(&mut region, row_offset)?; 314 | } 315 | 316 | Ok(()) 317 | } 318 | )?; 319 | 320 | Ok(digest_alloc.unwrap()) 321 | } 322 | } 323 | 324 | #[derive(Clone)] 325 | struct MerkleCircuit { 326 | // Private inputs. 327 | leaf: Option, 328 | path: Option>, 329 | // Public inputs (from the prover). The root is also a public input, but it is calculated within 330 | // the circuit. 331 | c_bits: Option>, 332 | } 333 | 334 | impl Circuit for MerkleCircuit { 335 | type Config = MerkleChipConfig; 336 | 337 | fn configure(cs: &mut ConstraintSystem) -> Self::Config { 338 | MerkleChip::configure(cs) 339 | } 340 | 341 | fn synthesize(&self, cs: &mut impl Assignment, config: Self::Config) -> Result<(), Error> { 342 | let mut layouter = SingleChipLayouter::new(cs)?; 343 | let merkle_chip = MerkleChip::new(config); 344 | let mut layer_digest = merkle_chip.hash_leaf_layer( 345 | &mut layouter, 346 | self.leaf.as_ref().unwrap().clone(), 347 | self.path.as_ref().unwrap()[0], 348 | self.c_bits.as_ref().unwrap()[0].clone(), 349 | )?; 350 | for layer in 1..PATH_LEN { 351 | layer_digest = merkle_chip.hash_non_leaf_layer( 352 | &mut layouter, 353 | layer_digest, 354 | self.path.as_ref().unwrap()[layer].clone(), 355 | self.c_bits.as_ref().unwrap()[layer].clone(), 356 | layer, 357 | )?; 358 | } 359 | Ok(()) 360 | } 361 | } 362 | 363 | struct Tree(Vec>); 364 | 365 | impl Tree { 366 | fn rand() -> Self { 367 | let mut rng = thread_rng(); 368 | let leafs: Vec = (0..N_LEAFS).map(|_| Fp::random(&mut rng)).collect(); 369 | let mut layers = vec![leafs]; 370 | for l in 1..TREE_LAYERS { 371 | let layer: Vec = layers[l - 1] 372 | .chunks(2) 373 | .map(|pair| mock_hash(pair[0], pair[1])) 374 | .collect(); 375 | layers.push(layer) 376 | } 377 | assert_eq!(layers.last().unwrap().len(), 1); 378 | Tree(layers) 379 | } 380 | 381 | fn root(&self) -> Fp { 382 | self.0.last().unwrap()[0] 383 | } 384 | 385 | fn leafs(&self) -> &[Fp] { 386 | self.0.first().unwrap() 387 | } 388 | 389 | fn gen_path(&self, c: usize) -> Vec { 390 | let mut path = vec![]; 391 | let mut node_index = c; 392 | for layer in 0..PATH_LEN { 393 | let sib = if node_index & 1 == 0 { 394 | self.0[layer][node_index + 1].clone() 395 | } else { 396 | self.0[layer][node_index - 1].clone() 397 | }; 398 | path.push(sib); 399 | node_index /= 2; 400 | } 401 | path 402 | } 403 | } 404 | 405 | fn main() { 406 | assert!(N_LEAFS.is_power_of_two()); 407 | 408 | // Generate a Merkle tree from random data. 409 | let tree = Tree::rand(); 410 | 411 | // Generate a random challenge, i.e. a leaf index in `[0, N_LEAFS)`. 412 | let c: usize = thread_rng().gen_range(0..N_LEAFS); 413 | let c_bits: Vec = (0..PATH_LEN) 414 | .map(|i| { 415 | if (c >> i) & 1 == 0 { 416 | Fp::zero() 417 | } else { 418 | Fp::one() 419 | } 420 | }) 421 | .collect(); 422 | 423 | // Create the public inputs. Every other row in the constraint system has a public input for a 424 | // challenge bit, additionally the last row has a public input for the root. 425 | let k = (N_ROWS_USED as f32).log2().ceil() as u32; 426 | let mut pub_inputs = vec![Fp::zero(); 1 << k]; 427 | for i in 0..PATH_LEN { 428 | pub_inputs[2 * i] = c_bits[i].clone(); 429 | } 430 | pub_inputs[LAST_ROW] = tree.root(); 431 | 432 | // Assert that the constraint system is satisfied for a witness corresponding to `pub_inputs`. 433 | let circuit = MerkleCircuit { 434 | leaf: Some(tree.leafs()[c].clone()), 435 | path: Some(tree.gen_path(c)), 436 | c_bits: Some(c_bits), 437 | }; 438 | let prover = MockProver::run(k, &circuit, vec![pub_inputs.clone()]).unwrap(); 439 | assert!(prover.verify().is_ok()); 440 | 441 | // Assert that changing the public challenge causes the constraint system to become unsatisfied. 442 | let mut bad_pub_inputs = pub_inputs.clone(); 443 | bad_pub_inputs[0] = if pub_inputs[0] == Fp::zero() { 444 | Fp::one() 445 | } else { 446 | Fp::zero() 447 | }; 448 | let prover = MockProver::run(k, &circuit, vec![bad_pub_inputs]).unwrap(); 449 | let res = prover.verify(); 450 | assert!(res.is_err()); 451 | if let Err(errors) = res { 452 | assert_eq!(errors.len(), 1); 453 | if let VerifyFailure::Gate { gate_name, row, .. } = errors[0] { 454 | assert_eq!(gate_name, "public input"); 455 | assert_eq!(row, 0); 456 | } else { 457 | panic!("expected public input gate failure"); 458 | } 459 | } 460 | 461 | // Assert that changing the public root causes the constraint system to become unsatisfied. 462 | let mut bad_pub_inputs = pub_inputs.clone(); 463 | bad_pub_inputs[LAST_ROW] += Fp::one(); 464 | let prover = MockProver::run(k, &circuit, vec![bad_pub_inputs]).unwrap(); 465 | let res = prover.verify(); 466 | assert!(res.is_err()); 467 | if let Err(errors) = res { 468 | assert_eq!(errors.len(), 1); 469 | if let VerifyFailure::Gate { gate_name, row, .. } = errors[0] { 470 | assert_eq!(gate_name, "public input"); 471 | assert_eq!(row, LAST_ROW); 472 | } else { 473 | panic!("expected public input gate failure"); 474 | } 475 | } 476 | 477 | // Assert that a non-boolean challenge bit causes the boolean constrain and swap gates to fail. 478 | let mut bad_pub_inputs = pub_inputs.clone(); 479 | bad_pub_inputs[0] = Fp::from(2); 480 | let mut bad_circuit = circuit.clone(); 481 | bad_circuit.c_bits.as_mut().unwrap()[0] = Fp::from(2); 482 | let prover = MockProver::run(k, &bad_circuit, vec![bad_pub_inputs]).unwrap(); 483 | let res = prover.verify(); 484 | assert!(res.is_err()); 485 | if let Err(errors) = res { 486 | assert_eq!(errors.len(), 2); 487 | if let VerifyFailure::Gate { gate_name, row, .. } = errors[0] { 488 | assert_eq!(gate_name, "boolean constrain"); 489 | assert_eq!(row, 0); 490 | } else { 491 | panic!("expected boolean constrain gate failure"); 492 | } 493 | if let VerifyFailure::Gate { gate_name, row, .. } = errors[1] { 494 | assert_eq!(gate_name, "swap"); 495 | assert_eq!(row, 0); 496 | } else { 497 | panic!("expected swap gate failure"); 498 | } 499 | } 500 | 501 | // Assert that changing the witnessed path causes the constraint system to become unsatisfied 502 | // when checking that the calculated root is equal to the public input root. 503 | let mut bad_circuit = circuit.clone(); 504 | bad_circuit.path.as_mut().unwrap()[0] += Fp::one(); 505 | let prover = MockProver::run(k, &bad_circuit, vec![pub_inputs.clone()]).unwrap(); 506 | let res = prover.verify(); 507 | assert!(res.is_err()); 508 | if let Err(errors) = res { 509 | assert_eq!(errors.len(), 1); 510 | if let VerifyFailure::Gate { gate_name, row, .. } = errors[0] { 511 | assert_eq!(gate_name, "public input"); 512 | assert_eq!(row, LAST_ROW); 513 | } else { 514 | panic!("expected public input gate failure"); 515 | } 516 | } 517 | } 518 | -------------------------------------------------------------------------------- /src/poseidon_chip.rs: -------------------------------------------------------------------------------- 1 | // TODO: do not perform S-box and MDS mixing on first round's `state[0]`. 2 | // TODO: perform MDS mixing exclusively on `state[1]` in the last round. 3 | 4 | const WIDTH: usize = 3; 5 | #[allow(non_upper_case_globals)] 6 | const R_f: usize = 3; 7 | const R_P: usize = 2; 8 | const R: usize = 2 * R_f + R_P; 9 | 10 | use halo2::{ 11 | arithmetic::{Field, FieldExt}, 12 | circuit::{layouter::SingleChipLayouter, Cell, Chip, Layouter}, 13 | dev::{MockProver, VerifyFailure}, 14 | pasta::Fp, 15 | plonk::{ 16 | Advice, Assignment, Circuit, Column, ConstraintSystem, Error, Expression, Fixed, Instance, 17 | Permutation, Selector, 18 | }, 19 | poly::Rotation, 20 | }; 21 | use lazy_static::lazy_static; 22 | use rand::{thread_rng, Rng, SeedableRng}; 23 | use rand_chacha::ChaCha8Rng; 24 | 25 | lazy_static! { 26 | /* 27 | static ref PRE_KEYS: Vec = rand_pre_keys([1; 32]); 28 | static ref POST_KEYS: Vec> = rand_post_keys([2; 32]); 29 | static ref MDS: Vec> = rand_matrix([3; 32]); 30 | static ref PRE_SPARSE: Vec> = rand_matrix([4; 32]); 31 | static ref SPARSE: Vec>> = (0..R_P) 32 | .map(|i| rand_matrix([5 + i as u8; 32])) 33 | .collect(); 34 | */ 35 | 36 | static ref PRE_KEYS: Vec = vec![Fp::from(1), Fp::from(2), Fp::from(3)]; 37 | static ref POST_KEYS: Vec> = vec![ 38 | vec![Fp::from(1), Fp::from(2), Fp::from(3)], 39 | vec![Fp::from(4), Fp::from(5), Fp::from(6)], 40 | vec![Fp::from(7), Fp::from(8), Fp::from(9)], 41 | vec![Fp::from(1)], 42 | vec![Fp::from(2)], 43 | vec![Fp::from(3), Fp::from(4), Fp::from(5)], 44 | vec![Fp::from(6), Fp::from(7), Fp::from(8)], 45 | vec![], 46 | ]; 47 | static ref MDS: Vec> = vec![ 48 | vec![Fp::from(1), Fp::from(2), Fp::from(3)], 49 | vec![Fp::from(4), Fp::from(5), Fp::from(6)], 50 | vec![Fp::from(7), Fp::from(8), Fp::from(9)], 51 | ]; 52 | static ref PRE_SPARSE: Vec> = vec![ 53 | vec![Fp::from(3), Fp::from(4), Fp::from(5)], 54 | vec![Fp::from(6), Fp::from(7), Fp::from(8)], 55 | vec![Fp::from(9), Fp::from(1), Fp::from(2)], 56 | ]; 57 | static ref SPARSE: Vec>> = vec![ 58 | vec![ 59 | vec![Fp::from(5), Fp::from(6), Fp::from(7)], 60 | vec![Fp::from(8), Fp::from(9), Fp::from(1)], 61 | vec![Fp::from(2), Fp::from(3), Fp::from(4)], 62 | ], 63 | vec![ 64 | vec![Fp::from(7), Fp::from(8), Fp::from(9)], 65 | vec![Fp::from(1), Fp::from(2), Fp::from(3)], 66 | vec![Fp::from(4), Fp::from(5), Fp::from(6)], 67 | ], 68 | ]; 69 | } 70 | 71 | fn rand_pre_keys(seed: [u8; 32]) -> Vec { 72 | let mut rng = ChaCha8Rng::from_seed(seed); 73 | (0..WIDTH) 74 | .map(|_| { 75 | // Fp::random(&mut rng) 76 | Fp::from(rng.gen::()) 77 | }) 78 | .collect() 79 | } 80 | 81 | fn rand_post_keys(seed: [u8; 32]) -> Vec> { 82 | let mut rng = ChaCha8Rng::from_seed(seed); 83 | (0..R) 84 | .map(|round| { 85 | if is_full_round(round) && round != R - 1 { 86 | // (0..WIDTH).map(|_| Fp::random(&mut rng)).collect() 87 | (0..WIDTH).map(|_| Fp::from(rng.gen::())).collect() 88 | } else if is_partial_round(round) { 89 | // vec![Fp::random(&mut rng)] 90 | vec![Fp::from(rng.gen::())] 91 | } else { 92 | vec![] 93 | } 94 | }) 95 | .collect() 96 | } 97 | 98 | fn rand_matrix(seed: [u8; 32]) -> Vec> { 99 | let mut rng = ChaCha8Rng::from_seed(seed); 100 | (0..WIDTH) 101 | // .map(|_| (0..WIDTH).map(|_| Fp::random(&mut rng)).collect()) 102 | .map(|_| (0..WIDTH).map(|_| Fp::from(rng.gen::())).collect()) 103 | .collect() 104 | } 105 | 106 | fn is_full_round(round: usize) -> bool { 107 | round < R_f || round >= R_f + R_P 108 | } 109 | 110 | fn is_partial_round(round: usize) -> bool { 111 | round >= R_f && round < R_f + R_P 112 | } 113 | 114 | fn pow5(x: Fp) -> Fp { 115 | x.square().square() * x 116 | } 117 | 118 | fn sbox_pre_post(state: &[Fp]) -> Vec { 119 | (0..WIDTH) 120 | .map(|i| pow5(state[i] + PRE_KEYS[i]) + POST_KEYS[0][i]) 121 | .collect() 122 | } 123 | 124 | fn sbox_post(state: &[Fp], post_keys: &[Fp]) -> Vec { 125 | (0..WIDTH) 126 | .map(|i| pow5(state[i]) + post_keys[i]) 127 | .collect() 128 | } 129 | 130 | fn sbox_no_add(state: &[Fp]) -> Vec { 131 | (0..WIDTH) 132 | .map(|i| pow5(state[i])) 133 | .collect() 134 | } 135 | 136 | fn vec_matrix_mul(v: &[Fp], m: &[Vec]) -> Vec { 137 | let n = v.len(); 138 | assert_eq!(m.len(), n); 139 | (0..n) 140 | .map(|col| { 141 | let mut lc = Fp::zero(); 142 | for i in 0..n { 143 | lc += v[i] * m[i][col]; 144 | } 145 | lc 146 | }) 147 | .collect() 148 | } 149 | 150 | #[derive(Clone, Debug)] 151 | struct Alloc { 152 | cell: Cell, 153 | value: Fp, 154 | } 155 | 156 | #[derive(Clone, Debug)] 157 | enum MaybeAlloc { 158 | Allocated(Alloc), 159 | Unallocated(Fp), 160 | } 161 | 162 | impl MaybeAlloc { 163 | fn value(&self) -> Fp { 164 | match self { 165 | MaybeAlloc::Allocated(alloc) => alloc.value.clone(), 166 | MaybeAlloc::Unallocated(value) => value.clone(), 167 | } 168 | } 169 | 170 | fn cell(&self) -> Cell { 171 | match self { 172 | MaybeAlloc::Allocated(alloc) => alloc.cell.clone(), 173 | MaybeAlloc::Unallocated(_) => 174 | panic!("called `MaybeAlloc.cell()` on an unallocated value"), 175 | } 176 | } 177 | 178 | fn is_allocated(&self) -> bool { 179 | match self { 180 | MaybeAlloc::Allocated(_) => true, 181 | MaybeAlloc::Unallocated(_) => false, 182 | } 183 | } 184 | 185 | fn is_unallocated(&self) -> bool { 186 | !self.is_allocated() 187 | } 188 | } 189 | 190 | struct PoseidonChip { 191 | config: PoseidonChipConfig, 192 | } 193 | 194 | #[derive(Clone, Debug)] 195 | struct PoseidonChipConfig { 196 | a_col: Column, 197 | sbox_out_col: Column, 198 | mds_out_col: Column, 199 | pre_key_col: Column, 200 | post_key_col: Column, 201 | mds_cols: Vec>, 202 | s_sbox_pre_post: Selector, 203 | s_sbox_post: Selector, 204 | s_sbox_no_add: Selector, 205 | s_mds: Vec, 206 | perm_output_to_input: Permutation, 207 | perm_output_to_sbox_output: Permutation, 208 | } 209 | 210 | 211 | impl Chip for PoseidonChip { 212 | type Config = PoseidonChipConfig; 213 | type Loaded = (); 214 | 215 | fn config(&self) -> &Self::Config { 216 | &self.config 217 | } 218 | 219 | fn loaded(&self) -> &Self::Loaded { 220 | &() 221 | } 222 | } 223 | 224 | impl PoseidonChip { 225 | fn new(config: PoseidonChipConfig) -> Self { 226 | PoseidonChip { config } 227 | } 228 | 229 | fn configure(cs: &mut ConstraintSystem, digest_col: Column) -> PoseidonChipConfig { 230 | let a_col = cs.advice_column(); 231 | let sbox_out_col = cs.advice_column(); 232 | let mds_out_col = digest_col; 233 | let pre_key_col = cs.fixed_column(); 234 | let post_key_col = cs.fixed_column(); 235 | let mds_cols = vec![cs.fixed_column(), cs.fixed_column(), cs.fixed_column()]; 236 | 237 | let s_sbox_pre_post = cs.selector(); 238 | let s_sbox_post = cs.selector(); 239 | let s_sbox_no_add = cs.selector(); 240 | let s_mds = vec![cs.selector(), cs.selector(), cs.selector()]; 241 | 242 | cs.create_gate("s_sbox_pre_post", |cs| { 243 | let a = cs.query_advice(a_col, Rotation::cur()); 244 | let pre_key = cs.query_fixed(pre_key_col, Rotation::cur()); 245 | let post_key = cs.query_fixed(post_key_col, Rotation::cur()); 246 | let sbox_out = cs.query_advice(sbox_out_col, Rotation::cur()); 247 | let s_sbox_pre_post = cs.query_selector(s_sbox_pre_post, Rotation::cur()); 248 | // (a + pre_key)^5 + post_key = out 249 | let a_plus_pre = a + pre_key; 250 | s_sbox_pre_post * ( 251 | a_plus_pre.clone() * a_plus_pre.clone() * a_plus_pre.clone() * a_plus_pre.clone() * 252 | a_plus_pre + post_key - sbox_out 253 | ) 254 | }); 255 | 256 | cs.create_gate("s_sbox_post", |cs| { 257 | let a = cs.query_advice(a_col, Rotation::cur()); 258 | let post_key = cs.query_fixed(post_key_col, Rotation::cur()); 259 | let sbox_out = cs.query_advice(sbox_out_col, Rotation::cur()); 260 | let s_sbox_post = cs.query_selector(s_sbox_post, Rotation::cur()); 261 | // a^5 + post_key = b 262 | s_sbox_post * (a.clone() * a.clone() * a.clone() * a.clone() * a + post_key - sbox_out) 263 | }); 264 | 265 | cs.create_gate("s_sbox_no_add", |cs| { 266 | let a = cs.query_advice(a_col, Rotation::cur()); 267 | let sbox_out = cs.query_advice(sbox_out_col, Rotation::cur()); 268 | let s_sbox_no_add = cs.query_selector(s_sbox_no_add, Rotation::cur()); 269 | // a^5 = b 270 | s_sbox_no_add * (a.clone() * a.clone() * a.clone() * a.clone() * a - sbox_out) 271 | }); 272 | 273 | // Calculates the dot product of the sbox outputs with column `0` of the MDS matrix. Note 274 | // that `s_mds_0` is enabled in the first MDS row. 275 | cs.create_gate("s_mds_0", |cs| { 276 | let sbox_out_0 = cs.query_advice(sbox_out_col, Rotation::cur()); 277 | let sbox_out_1 = cs.query_advice(sbox_out_col, Rotation::next()); 278 | let sbox_out_2 = cs.query_advice(sbox_out_col, Rotation(2)); 279 | let mds_out_0 = cs.query_advice(mds_out_col, Rotation::cur()); 280 | let s_mds_0 = cs.query_selector(s_mds[0], Rotation::cur()); 281 | 282 | // The first MDS column. 283 | let m_0 = cs.query_fixed(mds_cols[0], Rotation::cur()); 284 | let m_1 = cs.query_fixed(mds_cols[0], Rotation::next()); 285 | let m_2 = cs.query_fixed(mds_cols[0], Rotation(2)); 286 | 287 | // Dot product of sbox outputs with the first MDS column. 288 | s_mds_0 * (sbox_out_0 * m_0 + sbox_out_1 * m_1 + sbox_out_2 * m_2 - mds_out_0) 289 | }); 290 | 291 | // Calculates the dot product of the sbox outputs with column `1` of the MDS matrix. Note 292 | // that `s_mds_1` is enabled in the second MDS row. 293 | cs.create_gate("s_mds_1", |cs| { 294 | let sbox_out_0 = cs.query_advice(sbox_out_col, Rotation::prev()); 295 | let sbox_out_1 = cs.query_advice(sbox_out_col, Rotation::cur()); 296 | let sbox_out_2 = cs.query_advice(sbox_out_col, Rotation::next()); 297 | let mds_out_1 = cs.query_advice(mds_out_col, Rotation::cur()); 298 | let s_mds_1 = cs.query_selector(s_mds[1], Rotation::cur()); 299 | 300 | // The second MDS column. 301 | let m_0 = cs.query_fixed(mds_cols[1], Rotation::prev()); 302 | let m_1 = cs.query_fixed(mds_cols[1], Rotation::cur()); 303 | let m_2 = cs.query_fixed(mds_cols[1], Rotation::next()); 304 | 305 | // Dot product of the sbox outputs with the second MDS column. 306 | s_mds_1 * (sbox_out_0 * m_0 + sbox_out_1 * m_1 + sbox_out_2 * m_2 - mds_out_1) 307 | }); 308 | 309 | // Calculates the dot product of the sbox outputs with column `2` of the MDS matrix. Note 310 | // that `s_mds_2` is enabled in the third MDS row. 311 | cs.create_gate("s_mds_2", |cs| { 312 | let sbox_out_0 = cs.query_advice(sbox_out_col, Rotation(-2)); 313 | let sbox_out_1 = cs.query_advice(sbox_out_col, Rotation::prev()); 314 | let sbox_out_2 = cs.query_advice(sbox_out_col, Rotation::cur()); 315 | let mds_out_2 = cs.query_advice(mds_out_col, Rotation::cur()); 316 | let s_mds_2 = cs.query_selector(s_mds[2], Rotation::cur()); 317 | 318 | // The third MDS column. 319 | let m_0 = cs.query_fixed(mds_cols[2], Rotation(-2)); 320 | let m_1 = cs.query_fixed(mds_cols[2], Rotation::prev()); 321 | let m_2 = cs.query_fixed(mds_cols[2], Rotation::cur()); 322 | 323 | // Dot product of the sbox outputs with the third MDS column. 324 | s_mds_2 * (sbox_out_0 * m_0 + sbox_out_1 * m_1 + sbox_out_2 * m_2 - mds_out_2) 325 | }); 326 | 327 | // Copies a round's MDS output into the next round's state. 328 | let perm_output_to_input = 329 | Permutation::new(cs, &[mds_out_col.into(), a_col.into()]); 330 | 331 | // Copies a round's MDS output into the next round's sbox output. 332 | let perm_output_to_sbox_output = 333 | Permutation::new(cs, &[mds_out_col.into(), sbox_out_col.into()]); 334 | 335 | PoseidonChipConfig { 336 | a_col, 337 | sbox_out_col, 338 | mds_out_col, 339 | pre_key_col, 340 | post_key_col, 341 | mds_cols, 342 | s_sbox_pre_post, 343 | s_sbox_post, 344 | s_sbox_no_add, 345 | s_mds, 346 | perm_output_to_input, 347 | perm_output_to_sbox_output, 348 | } 349 | } 350 | 351 | fn alloc_full_round( 352 | &self, 353 | layouter: &mut impl Layouter, 354 | // We need to pass in the `state` as `MaybeAlloc` to allow `alloc_full_round()` to be called 355 | // for the first round, i.e. the state values are unallocated prior to the first round. 356 | state: &[MaybeAlloc], 357 | round: usize, 358 | ) -> Result, Error> { 359 | dbg!(round); 360 | assert!(is_full_round(round)); 361 | assert_eq!(state.len(), WIDTH); 362 | 363 | let is_first_round = round == 0; 364 | let is_last_round = round == R - 1; 365 | let is_pre_sparse_round = round == R_f - 1; 366 | 367 | if is_first_round { 368 | for state_i in state { 369 | assert!(state_i.is_unallocated()); 370 | } 371 | } else { 372 | for state_i in state { 373 | assert!(state_i.is_allocated()); 374 | } 375 | } 376 | 377 | let post_keys = &*POST_KEYS[round]; 378 | assert_eq!(post_keys.len(), if is_last_round { 0 } else { WIDTH }); 379 | 380 | // Copy field elements out of `MaybeAlloc`'s for more concise arithmetic. 381 | let state_values: Vec = state 382 | .iter() 383 | .map(|maybe_alloc| maybe_alloc.value()) 384 | .collect(); 385 | 386 | // Calculate the S-box output for each state element. 387 | let sbox_outputs = if is_first_round { 388 | sbox_pre_post(&state_values) 389 | } else if is_last_round { 390 | sbox_no_add(&state_values) 391 | } else { 392 | sbox_post(&state_values, &post_keys) 393 | }; 394 | 395 | // Calculate the MDS mixing output for each state element. 396 | let m = if is_pre_sparse_round { &*PRE_SPARSE } else { &*MDS }; 397 | let mds_outputs = vec_matrix_mul(&sbox_outputs, m); 398 | 399 | // Store the allocated outputs of MDS mixing. 400 | let mut mds_outputs_alloc: Vec> = vec![None; WIDTH]; 401 | 402 | layouter.assign_region( 403 | || format!("alloc (full) round {}", round), 404 | |mut region| { 405 | // Allocate each state element's row in the constraint system. 406 | for row_offset in 0..WIDTH { 407 | dbg!(row_offset); 408 | // If this is the first round, we are allocating the state element for the first 409 | // time, otherwise we are reallocating an element output by the previous round. 410 | let a_cell = region.assign_advice( 411 | || format!("a_{} (round {})", row_offset, round), 412 | self.config.a_col, 413 | row_offset, 414 | || Ok(state_values[row_offset]), 415 | )?; 416 | 417 | if !is_first_round { 418 | region.constrain_equal( 419 | &self.config.perm_output_to_input, 420 | state[row_offset].cell(), 421 | a_cell, 422 | )?; 423 | } 424 | 425 | // If this is the first round allocate a pre-S-box key. 426 | if is_first_round { 427 | region.assign_fixed( 428 | || format!("pre_key_{} (round {})", row_offset, round), 429 | self.config.pre_key_col, 430 | row_offset, 431 | || Ok(PRE_KEYS[row_offset]), 432 | )?; 433 | } 434 | 435 | // If this is not the last round allocate a post-S-box key. 436 | if !is_last_round { 437 | region.assign_fixed( 438 | || format!("post_key_{} (round {})", row_offset, round), 439 | self.config.post_key_col, 440 | row_offset, 441 | || Ok(post_keys[row_offset]), 442 | )?; 443 | } 444 | 445 | // Allocate the S-box output. 446 | region.assign_advice( 447 | || format!("sbox_out_{} (round {})", row_offset, round), 448 | self.config.sbox_out_col, 449 | row_offset, 450 | || Ok(sbox_outputs[row_offset]), 451 | )?; 452 | 453 | // Allocate the MDS mixing output. 454 | let mds_output_cell = region.assign_advice( 455 | || format!("mds_out_{} (round {})", row_offset, round), 456 | self.config.mds_out_col, 457 | row_offset, 458 | || Ok(mds_outputs[row_offset]), 459 | )?; 460 | 461 | // Keep a reference to the allocated MDS output. 462 | mds_outputs_alloc[row_offset] = Some(MaybeAlloc::Allocated(Alloc { 463 | cell: mds_output_cell, 464 | value: mds_outputs[row_offset], 465 | })); 466 | 467 | // Enable the S-box and MDS mixing selectors. 468 | if is_first_round { 469 | self.config.s_sbox_pre_post.enable(&mut region, row_offset)?; 470 | } else if is_last_round { 471 | self.config.s_sbox_no_add.enable(&mut region, row_offset)?; 472 | } else { 473 | self.config.s_sbox_post.enable(&mut region, row_offset)?; 474 | }; 475 | self.config.s_mds[row_offset].enable(&mut region, row_offset)?; 476 | 477 | // Allocate this MDS matrix row. 478 | for col in 0..WIDTH { 479 | region.assign_fixed( 480 | || format!( 481 | "{} row={}, col={} (round {})", 482 | if is_pre_sparse_round { "P" } else { "MDS" }, 483 | row_offset, 484 | col, 485 | round, 486 | ), 487 | self.config.mds_cols[col], 488 | row_offset, 489 | || Ok(m[row_offset][col]), 490 | )?; 491 | } 492 | } 493 | Ok(()) 494 | }, 495 | )?; 496 | 497 | let mds_outputs_alloc: Vec = mds_outputs_alloc 498 | .into_iter() 499 | .map(|opt| opt.unwrap()) 500 | .collect(); 501 | 502 | Ok(mds_outputs_alloc) 503 | } 504 | 505 | fn alloc_partial_round( 506 | &self, 507 | layouter: &mut impl Layouter, 508 | state: &[MaybeAlloc], 509 | round: usize, 510 | ) -> Result, Error> { 511 | dbg!(round); 512 | assert!(is_partial_round(round)); 513 | assert_eq!(state.len(), WIDTH); 514 | 515 | for state_i in state { 516 | assert!(state_i.is_allocated()); 517 | } 518 | 519 | assert_eq!(POST_KEYS[round].len(), 1); 520 | let post_key = POST_KEYS[round][0]; 521 | 522 | // Copy field elements out of `MaybeAlloc`'s for more concise arithmetic. 523 | let state_values: Vec = state 524 | .iter() 525 | .map(|maybe_alloc| maybe_alloc.value()) 526 | .collect(); 527 | 528 | // Calculate the S-box output for the first state element. 529 | let mut sbox_outputs: Vec = vec![pow5(state_values[0]) + post_key]; 530 | sbox_outputs.extend_from_slice(&state_values[1..]); 531 | 532 | // Calculate the MDS mixing output for each state element. 533 | let sparse_index = round - R_f; 534 | let m = &*SPARSE[sparse_index]; 535 | let mds_outputs = vec_matrix_mul(&sbox_outputs, m); 536 | 537 | // Store the allocated outputs of MDS mixing. 538 | let mut mds_outputs_alloc: Vec> = vec![None; WIDTH]; 539 | 540 | layouter.assign_region( 541 | || format!("alloc (partial) round {}", round), 542 | |mut region| { 543 | // Allocate values that are exclusive to the first row. 544 | let row_offset = 0; 545 | 546 | // Reallocate the first state element which was output by the previous round. 547 | let a_cell = region.assign_advice( 548 | || format!("a_0 (round {})", round), 549 | self.config.a_col, 550 | row_offset, 551 | || Ok(state_values[0]), 552 | )?; 553 | 554 | region.constrain_equal( 555 | &self.config.perm_output_to_input, 556 | state[0].cell(), 557 | a_cell, 558 | )?; 559 | 560 | // Allocate the first state element's post-S-box key. 561 | region.assign_fixed( 562 | || format!("post_key_0 (round {})", round), 563 | self.config.post_key_col, 564 | row_offset, 565 | || Ok(post_key), 566 | )?; 567 | 568 | // Enable the first state element's S-box selector. 569 | self.config.s_sbox_post.enable(&mut region, row_offset)?; 570 | 571 | // Allocate the remaining round values. 572 | for row_offset in 0..WIDTH { 573 | // If this is the first row (`row_offset = 0`), allocate the first state 574 | // element's S-box output. If this is not the first row, reallocate the state 575 | // element output by the previous round. 576 | let sbox_out_cell = region.assign_advice( 577 | || format!("sbox_out_{} (round {})", row_offset, round), 578 | self.config.sbox_out_col, 579 | row_offset, 580 | || Ok(sbox_outputs[row_offset]), 581 | )?; 582 | 583 | if row_offset > 0 { 584 | region.constrain_equal( 585 | &self.config.perm_output_to_sbox_output, 586 | state[row_offset].cell(), 587 | sbox_out_cell, 588 | )?; 589 | } 590 | 591 | // Allocate the state element's MDS mixing output. 592 | let mds_out_cell = region.assign_advice( 593 | || format!("mds_out_{} (round {})", row_offset, round), 594 | self.config.mds_out_col, 595 | row_offset, 596 | || Ok(mds_outputs[row_offset]), 597 | )?; 598 | 599 | // Keep a reference to the allocated MDS output. 600 | mds_outputs_alloc[row_offset] = Some(MaybeAlloc::Allocated(Alloc { 601 | cell: mds_out_cell, 602 | value: mds_outputs[row_offset], 603 | })); 604 | 605 | // Enable the MDS mixing selector for this state element. 606 | self.config.s_mds[row_offset].enable(&mut region, row_offset)?; 607 | 608 | // Allocate this MDS matrix row. 609 | for col in 0..WIDTH { 610 | region.assign_fixed( 611 | || format!( 612 | "S{} row={}, col={} (round {})", 613 | sparse_index, 614 | row_offset, 615 | col, 616 | round, 617 | ), 618 | self.config.mds_cols[col], 619 | row_offset, 620 | || Ok(m[row_offset][col]), 621 | )?; 622 | } 623 | } 624 | Ok(()) 625 | }, 626 | )?; 627 | 628 | let mds_outputs_alloc: Vec = mds_outputs_alloc 629 | .into_iter() 630 | .map(|opt| opt.unwrap()) 631 | .collect(); 632 | 633 | Ok(mds_outputs_alloc) 634 | } 635 | } 636 | 637 | struct PoseidonCircuit { 638 | // Private inputs. 639 | initial_state: Vec, 640 | // Public inputs. 641 | digest: Fp, 642 | } 643 | 644 | #[derive(Clone)] 645 | struct PoseidonCircuitConfig { 646 | poseidon_config: PoseidonChipConfig, 647 | digest_col: Column, 648 | pub_col: Column, 649 | s_pub: Selector, 650 | perm_digest: Permutation, 651 | } 652 | 653 | impl Circuit for PoseidonCircuit { 654 | type Config = PoseidonCircuitConfig; 655 | 656 | fn configure(cs: &mut ConstraintSystem) -> Self::Config { 657 | let digest_col = cs.advice_column(); 658 | let poseidon_config = PoseidonChip::configure(cs, digest_col); 659 | let pub_col = cs.instance_column(); 660 | let s_pub = cs.selector(); 661 | let perm_digest = Permutation::new(cs, &[digest_col.into(), pub_col.into()]); 662 | 663 | // let (poseidon_config, io_cols) = PoseidonChip::configure(cs); 664 | // let preimg_cols = io_cols[..]; 665 | // let digest_col = io_cols[1]; 666 | // let pub_col = cs.instance_column(); 667 | // let s_pub = cs.selector(); 668 | // let perm_digest = Permutation::new(cs, &[digest_col.into(), pub_col.into()]); 669 | 670 | cs.create_gate("public input", |cs| { 671 | let digest = cs.query_advice(digest_col, Rotation::cur()); 672 | let pi = cs.query_instance(pub_col, Rotation::cur()); 673 | let s_pub = cs.query_selector(s_pub, Rotation::cur()); 674 | s_pub * (digest - pi) 675 | }); 676 | 677 | PoseidonCircuitConfig { 678 | poseidon_config, 679 | digest_col, 680 | pub_col, 681 | s_pub, 682 | perm_digest, 683 | } 684 | // PoseidonCircuitConfig { 685 | // poseidon_config, 686 | // preimg_col, 687 | // digest_col, 688 | // pub_col, 689 | // s_pub, 690 | // perm_digest, 691 | // } 692 | } 693 | 694 | fn synthesize(&self, cs: &mut impl Assignment, config: Self::Config) -> Result<(), Error> { 695 | let mut layouter = SingleChipLayouter::new(cs)?; 696 | let poseidon_chip = PoseidonChip::new(config.poseidon_config.clone()); 697 | 698 | let mut state_alloc: Vec = self.initial_state 699 | .iter() 700 | .map(|state_i| MaybeAlloc::Unallocated(state_i.clone())) 701 | .collect(); 702 | 703 | for round in 0..R_f { 704 | state_alloc = poseidon_chip.alloc_full_round(&mut layouter, &state_alloc, round)?; 705 | } 706 | 707 | for round in R_f..R_f + R_P { 708 | state_alloc = poseidon_chip.alloc_partial_round(&mut layouter, &state_alloc, round)?; 709 | } 710 | 711 | for round in R_f + R_P..R { 712 | state_alloc = poseidon_chip.alloc_full_round(&mut layouter, &state_alloc, round)?; 713 | } 714 | 715 | // The calculated digest is the second element of the output state vector. 716 | let digest_alloc = &state_alloc[1]; 717 | dbg!(digest_alloc.value()); 718 | 719 | layouter.assign_region( 720 | || "digest equality", 721 | |mut region| { 722 | let row_offset = 0; 723 | let digest_copy_cell = region.assign_advice( 724 | || "digest copy", 725 | config.digest_col, 726 | row_offset, 727 | || Ok(digest_alloc.value()), 728 | )?; 729 | region.constrain_equal(&config.perm_digest, digest_alloc.cell(), digest_copy_cell)?; 730 | config.s_pub.enable(&mut region, row_offset)?; 731 | Ok(()) 732 | }, 733 | ); 734 | 735 | Ok(()) 736 | } 737 | } 738 | 739 | fn poseidon(preimg: &[Fp]) -> Fp { 740 | let mut state = sbox_pre_post(&preimg); 741 | state = vec_matrix_mul(&state, &*MDS); 742 | 743 | for round in 1..R_f { 744 | state = sbox_post(&state, &*POST_KEYS[round]); 745 | let m = if round == R_f - 1 { &*PRE_SPARSE } else { &*MDS }; 746 | state = vec_matrix_mul(&state, m); 747 | } 748 | 749 | for round in R_f..R_f + R_P { 750 | state[0] = pow5(state[0].clone()) + POST_KEYS[round][0]; 751 | let sparse_index = round - R_f; 752 | state = vec_matrix_mul(&state, &*SPARSE[sparse_index]); 753 | } 754 | 755 | for round in R_f + R_P..R - 1 { 756 | state = sbox_post(&state, &*POST_KEYS[round]); 757 | state = vec_matrix_mul(&state, &*MDS); 758 | } 759 | 760 | state = sbox_no_add(&state); 761 | state = vec_matrix_mul(&state, &*MDS); 762 | state[1] 763 | } 764 | 765 | fn main() { 766 | // There are `WIDTH` number of rows per round; add one row for checking that the calculated 767 | // digest is equal to the public digest. 768 | const N_ROWS_USED: u32 = (R * WIDTH + 1) as u32; 769 | const PUB_INPUT_ROW_INDEX: usize = N_ROWS_USED as usize - 1; 770 | 771 | // The public digest. 772 | let pub_input = Fp::from_bytes(&[ 773 | 105u8, 223, 174, 214, 135, 774 | 10, 246, 134, 56, 44, 775 | 82, 200, 244, 29, 158, 776 | 165, 255, 6, 80, 24, 777 | 144, 74, 184, 235, 28, 778 | 196, 134, 44, 131, 236, 779 | 207, 13, 780 | ]).unwrap(); 781 | 782 | // Verifier's public inputs. 783 | let k = (N_ROWS_USED as f32).log2().ceil() as u32; 784 | let n_rows = 1 << k; 785 | let mut pub_inputs = vec![Fp::zero(); n_rows]; 786 | pub_inputs[PUB_INPUT_ROW_INDEX] = Fp::from(pub_input); 787 | 788 | let preimg = vec![Fp::from(55), Fp::from(101), Fp::from(237)]; 789 | dbg!(poseidon(&preimg)); 790 | 791 | // Prover's circuit contains public and private inputs. 792 | let circuit = PoseidonCircuit { 793 | initial_state: preimg, 794 | digest: pub_input, 795 | }; 796 | 797 | let prover = MockProver::run(k, &circuit, vec![pub_inputs.clone()]).unwrap(); 798 | dbg!(prover.verify()); 799 | // assert!(prover.verify().is_ok()); 800 | } 801 | -------------------------------------------------------------------------------- /src/poseidon_wrapped.rs: -------------------------------------------------------------------------------- 1 | const WIDTH: usize = 3; 2 | #[allow(non_upper_case_globals)] 3 | const R_f: usize = 3; 4 | const R_P: usize = 2; 5 | const R: usize = 2 * R_f + R_P; 6 | 7 | use halo2::{ 8 | arithmetic::{Field, FieldExt}, 9 | circuit::{layouter::SingleChipLayouter, Cell, Chip, Layouter}, 10 | dev::{MockProver, VerifyFailure}, 11 | pasta::Fp, 12 | plonk::{ 13 | Advice, Assignment, Circuit, Column, ConstraintSystem, Error, Fixed, Instance, Permutation, 14 | Selector, 15 | }, 16 | poly::Rotation, 17 | }; 18 | use lazy_static::lazy_static; 19 | use rand::{Rng, SeedableRng}; 20 | use rand_chacha::ChaCha8Rng; 21 | 22 | lazy_static! { 23 | static ref PRE_KEYS: Vec = rand_pre_keys([1; 32]); 24 | static ref POST_KEYS: Vec> = rand_post_keys([2; 32]); 25 | static ref MDS: Vec> = rand_matrix([3; 32]); 26 | static ref PRE_SPARSE: Vec> = rand_matrix([4; 32]); 27 | static ref SPARSE: Vec>> = (0..R_P) 28 | .map(|i| rand_matrix([5 + i as u8; 32])) 29 | .collect(); 30 | 31 | /* 32 | static ref PRE_KEYS: Vec = vec![Fp::from(1), Fp::from(2), Fp::from(3)]; 33 | static ref POST_KEYS: Vec> = vec![ 34 | vec![Fp::from(1), Fp::from(2), Fp::from(3)], 35 | vec![Fp::from(4), Fp::from(5), Fp::from(6)], 36 | vec![Fp::from(7), Fp::from(8), Fp::from(9)], 37 | vec![Fp::from(1)], 38 | vec![Fp::from(2)], 39 | vec![Fp::from(3), Fp::from(4), Fp::from(5)], 40 | vec![Fp::from(6), Fp::from(7), Fp::from(8)], 41 | vec![], 42 | ]; 43 | static ref MDS: Vec> = vec![ 44 | vec![Fp::from(1), Fp::from(2), Fp::from(3)], 45 | vec![Fp::from(4), Fp::from(5), Fp::from(6)], 46 | vec![Fp::from(7), Fp::from(8), Fp::from(9)], 47 | ]; 48 | static ref PRE_SPARSE: Vec> = vec![ 49 | vec![Fp::from(3), Fp::from(4), Fp::from(5)], 50 | vec![Fp::from(6), Fp::from(7), Fp::from(8)], 51 | vec![Fp::from(9), Fp::from(1), Fp::from(2)], 52 | ]; 53 | static ref SPARSE: Vec>> = vec![ 54 | vec![ 55 | vec![Fp::from(5), Fp::from(6), Fp::from(7)], 56 | vec![Fp::from(8), Fp::from(9), Fp::from(1)], 57 | vec![Fp::from(2), Fp::from(3), Fp::from(4)], 58 | ], 59 | vec![ 60 | vec![Fp::from(7), Fp::from(8), Fp::from(9)], 61 | vec![Fp::from(1), Fp::from(2), Fp::from(3)], 62 | vec![Fp::from(4), Fp::from(5), Fp::from(6)], 63 | ], 64 | ]; 65 | */ 66 | } 67 | 68 | fn rand_pre_keys(seed: [u8; 32]) -> Vec { 69 | let mut rng = ChaCha8Rng::from_seed(seed); 70 | (0..WIDTH) 71 | .map(|_| Fp::from(rng.gen::())) 72 | .collect() 73 | } 74 | 75 | fn rand_post_keys(seed: [u8; 32]) -> Vec> { 76 | let mut rng = ChaCha8Rng::from_seed(seed); 77 | (0..R) 78 | .map(|round| { 79 | if is_full_round(round) && round != R - 1 { 80 | (0..WIDTH).map(|_| Fp::random(&mut rng)).collect() 81 | } else if is_partial_round(round) { 82 | vec![Fp::random(&mut rng)] 83 | } else { 84 | vec![] 85 | } 86 | }) 87 | .collect() 88 | } 89 | 90 | fn rand_matrix(seed: [u8; 32]) -> Vec> { 91 | let mut rng = ChaCha8Rng::from_seed(seed); 92 | (0..WIDTH) 93 | // .map(|_| (0..WIDTH).map(|_| Fp::random(&mut rng)).collect()) 94 | .map(|_| (0..WIDTH).map(|_| Fp::from(rng.gen::())).collect()) 95 | .collect() 96 | } 97 | 98 | fn is_full_round(round: usize) -> bool { 99 | round < R_f || round >= R_f + R_P 100 | } 101 | 102 | fn is_partial_round(round: usize) -> bool { 103 | round >= R_f && round < R_f + R_P 104 | } 105 | 106 | fn pow5(x: Fp) -> Fp { 107 | x.square().square() * x 108 | } 109 | 110 | fn sbox_pre_post(state: &[Fp]) -> Vec { 111 | (0..WIDTH) 112 | .map(|i| pow5(state[i] + PRE_KEYS[i]) + POST_KEYS[0][i]) 113 | .collect() 114 | } 115 | 116 | fn sbox_post(state: &[Fp], post_keys: &[Fp]) -> Vec { 117 | (0..WIDTH) 118 | .map(|i| pow5(state[i]) + post_keys[i]) 119 | .collect() 120 | } 121 | 122 | fn sbox_no_add(state: &[Fp]) -> Vec { 123 | (0..WIDTH) 124 | .map(|i| pow5(state[i])) 125 | .collect() 126 | } 127 | 128 | fn vec_matrix_mul(v: &[Fp], m: &[Vec]) -> Vec { 129 | let n = v.len(); 130 | assert_eq!(m.len(), n); 131 | (0..n) 132 | .map(|col| { 133 | let mut lc = Fp::zero(); 134 | for i in 0..n { 135 | lc += v[i] * m[i][col]; 136 | } 137 | lc 138 | }) 139 | .collect() 140 | } 141 | 142 | fn poseidon(preimg: &[Fp]) -> Fp { 143 | let mut state = sbox_pre_post(&preimg); 144 | state = vec_matrix_mul(&state, &*MDS); 145 | 146 | for round in 1..R_f { 147 | state = sbox_post(&state, &*POST_KEYS[round]); 148 | let m = if round == R_f - 1 { &*PRE_SPARSE } else { &*MDS }; 149 | state = vec_matrix_mul(&state, m); 150 | } 151 | 152 | for round in R_f..R_f + R_P { 153 | state[0] = pow5(state[0].clone()) + POST_KEYS[round][0]; 154 | let sparse_index = round - R_f; 155 | state = vec_matrix_mul(&state, &*SPARSE[sparse_index]); 156 | } 157 | 158 | for round in R_f + R_P..R - 1 { 159 | state = sbox_post(&state, &*POST_KEYS[round]); 160 | state = vec_matrix_mul(&state, &*MDS); 161 | } 162 | 163 | state = sbox_no_add(&state); 164 | state = vec_matrix_mul(&state, &*MDS); 165 | state[1] 166 | } 167 | 168 | #[derive(Clone, Debug)] 169 | struct Alloc { 170 | cell: Cell, 171 | value: Fp, 172 | } 173 | 174 | struct PoseidonChip { 175 | config: PoseidonChipConfig, 176 | } 177 | 178 | #[derive(Clone, Debug)] 179 | struct PoseidonChipConfig { 180 | a_col: Column, 181 | sbox_out_col: Column, 182 | mds_out_col: Column, 183 | pre_key_col: Column, 184 | post_key_col: Column, 185 | mds_cols: Vec>, 186 | s_sbox_pre_post: Selector, 187 | s_sbox_post: Selector, 188 | s_sbox_no_add: Selector, 189 | s_mds: Vec, 190 | perm_output_to_input: Permutation, 191 | perm_output_to_sbox_output: Permutation, 192 | perm_preimg: Vec, 193 | } 194 | 195 | 196 | impl Chip for PoseidonChip { 197 | type Config = PoseidonChipConfig; 198 | type Loaded = (); 199 | 200 | fn config(&self) -> &Self::Config { 201 | &self.config 202 | } 203 | 204 | fn loaded(&self) -> &Self::Loaded { 205 | &() 206 | } 207 | } 208 | 209 | impl PoseidonChip { 210 | fn new(config: PoseidonChipConfig) -> Self { 211 | PoseidonChip { config } 212 | } 213 | 214 | fn configure(cs: &mut ConstraintSystem) -> 215 | (PoseidonChipConfig, Vec>, Column) 216 | { 217 | let a_col = cs.advice_column(); 218 | let sbox_out_col = cs.advice_column(); 219 | let mds_out_col = cs.advice_column(); 220 | let pre_key_col = cs.fixed_column(); 221 | let post_key_col = cs.fixed_column(); 222 | let mds_cols = vec![cs.fixed_column(), cs.fixed_column(), cs.fixed_column()]; 223 | 224 | let s_sbox_pre_post = cs.selector(); 225 | let s_sbox_post = cs.selector(); 226 | let s_sbox_no_add = cs.selector(); 227 | let s_mds = vec![cs.selector(), cs.selector(), cs.selector()]; 228 | 229 | cs.create_gate("s_sbox_pre_post", |cs| { 230 | let a = cs.query_advice(a_col, Rotation::cur()); 231 | let pre_key = cs.query_fixed(pre_key_col, Rotation::cur()); 232 | let post_key = cs.query_fixed(post_key_col, Rotation::cur()); 233 | let sbox_out = cs.query_advice(sbox_out_col, Rotation::cur()); 234 | let s_sbox_pre_post = cs.query_selector(s_sbox_pre_post, Rotation::cur()); 235 | // (a + pre_key)^5 + post_key = out 236 | let a_plus_pre = a + pre_key; 237 | s_sbox_pre_post * ( 238 | a_plus_pre.clone() * a_plus_pre.clone() * a_plus_pre.clone() * a_plus_pre.clone() * 239 | a_plus_pre + post_key - sbox_out 240 | ) 241 | }); 242 | 243 | cs.create_gate("s_sbox_post", |cs| { 244 | let a = cs.query_advice(a_col, Rotation::cur()); 245 | let post_key = cs.query_fixed(post_key_col, Rotation::cur()); 246 | let sbox_out = cs.query_advice(sbox_out_col, Rotation::cur()); 247 | let s_sbox_post = cs.query_selector(s_sbox_post, Rotation::cur()); 248 | // a^5 + post_key = b 249 | s_sbox_post * (a.clone() * a.clone() * a.clone() * a.clone() * a + post_key - sbox_out) 250 | }); 251 | 252 | cs.create_gate("s_sbox_no_add", |cs| { 253 | let a = cs.query_advice(a_col, Rotation::cur()); 254 | let sbox_out = cs.query_advice(sbox_out_col, Rotation::cur()); 255 | let s_sbox_no_add = cs.query_selector(s_sbox_no_add, Rotation::cur()); 256 | // a^5 = b 257 | s_sbox_no_add * (a.clone() * a.clone() * a.clone() * a.clone() * a - sbox_out) 258 | }); 259 | 260 | // Calculates the dot product of the sbox outputs with column `0` of the MDS matrix. Note 261 | // that `s_mds_0` is enabled in the first MDS row. 262 | cs.create_gate("s_mds_0", |cs| { 263 | let sbox_out_0 = cs.query_advice(sbox_out_col, Rotation::cur()); 264 | let sbox_out_1 = cs.query_advice(sbox_out_col, Rotation::next()); 265 | let sbox_out_2 = cs.query_advice(sbox_out_col, Rotation(2)); 266 | let mds_out_0 = cs.query_advice(mds_out_col, Rotation::cur()); 267 | let s_mds_0 = cs.query_selector(s_mds[0], Rotation::cur()); 268 | 269 | // The first MDS column. 270 | let m_0 = cs.query_fixed(mds_cols[0], Rotation::cur()); 271 | let m_1 = cs.query_fixed(mds_cols[0], Rotation::next()); 272 | let m_2 = cs.query_fixed(mds_cols[0], Rotation(2)); 273 | 274 | // Dot product of sbox outputs with the first MDS column. 275 | s_mds_0 * (sbox_out_0 * m_0 + sbox_out_1 * m_1 + sbox_out_2 * m_2 - mds_out_0) 276 | }); 277 | 278 | // Calculates the dot product of the sbox outputs with column `1` of the MDS matrix. Note 279 | // that `s_mds_1` is enabled in the second MDS row. 280 | cs.create_gate("s_mds_1", |cs| { 281 | let sbox_out_0 = cs.query_advice(sbox_out_col, Rotation::prev()); 282 | let sbox_out_1 = cs.query_advice(sbox_out_col, Rotation::cur()); 283 | let sbox_out_2 = cs.query_advice(sbox_out_col, Rotation::next()); 284 | let mds_out_1 = cs.query_advice(mds_out_col, Rotation::cur()); 285 | let s_mds_1 = cs.query_selector(s_mds[1], Rotation::cur()); 286 | 287 | // The second MDS column. 288 | let m_0 = cs.query_fixed(mds_cols[1], Rotation::prev()); 289 | let m_1 = cs.query_fixed(mds_cols[1], Rotation::cur()); 290 | let m_2 = cs.query_fixed(mds_cols[1], Rotation::next()); 291 | 292 | // Dot product of the sbox outputs with the second MDS column. 293 | s_mds_1 * (sbox_out_0 * m_0 + sbox_out_1 * m_1 + sbox_out_2 * m_2 - mds_out_1) 294 | }); 295 | 296 | // Calculates the dot product of the sbox outputs with column `2` of the MDS matrix. Note 297 | // that `s_mds_2` is enabled in the third MDS row. 298 | cs.create_gate("s_mds_2", |cs| { 299 | let sbox_out_0 = cs.query_advice(sbox_out_col, Rotation(-2)); 300 | let sbox_out_1 = cs.query_advice(sbox_out_col, Rotation::prev()); 301 | let sbox_out_2 = cs.query_advice(sbox_out_col, Rotation::cur()); 302 | let mds_out_2 = cs.query_advice(mds_out_col, Rotation::cur()); 303 | let s_mds_2 = cs.query_selector(s_mds[2], Rotation::cur()); 304 | 305 | // The third MDS column. 306 | let m_0 = cs.query_fixed(mds_cols[2], Rotation(-2)); 307 | let m_1 = cs.query_fixed(mds_cols[2], Rotation::prev()); 308 | let m_2 = cs.query_fixed(mds_cols[2], Rotation::cur()); 309 | 310 | // Dot product of the sbox outputs with the third MDS column. 311 | s_mds_2 * (sbox_out_0 * m_0 + sbox_out_1 * m_1 + sbox_out_2 * m_2 - mds_out_2) 312 | }); 313 | 314 | // Copies a round's MDS output into the next round's state. 315 | let perm_output_to_input = 316 | Permutation::new(cs, &[mds_out_col.into(), a_col.into()]); 317 | 318 | // Copies a round's MDS output into the next round's sbox output. 319 | let perm_output_to_sbox_output = 320 | Permutation::new(cs, &[mds_out_col.into(), sbox_out_col.into()]); 321 | 322 | // Copy the first, second, and third preimage elements into the first column. 323 | let perm_preimg = { 324 | let perm_a_to_a = Permutation::new(cs, &[a_col.into()]); 325 | let perm_sbox_out_to_a = Permutation::new(cs, &[sbox_out_col.into(), a_col.into()]); 326 | let perm_mds_out_to_a = perm_output_to_input.clone(); 327 | vec![perm_a_to_a, perm_sbox_out_to_a, perm_mds_out_to_a] 328 | }; 329 | 330 | let poseidon_config = PoseidonChipConfig { 331 | a_col, 332 | sbox_out_col, 333 | mds_out_col, 334 | pre_key_col, 335 | post_key_col, 336 | mds_cols, 337 | s_sbox_pre_post, 338 | s_sbox_post, 339 | s_sbox_no_add, 340 | s_mds, 341 | perm_output_to_input, 342 | perm_output_to_sbox_output, 343 | perm_preimg, 344 | }; 345 | let preimg_cols = vec![a_col, sbox_out_col, mds_out_col]; 346 | let digest_col = mds_out_col; 347 | (poseidon_config, preimg_cols, digest_col) 348 | } 349 | 350 | fn hash( 351 | &self, 352 | layouter: &mut impl Layouter, 353 | preimg_alloc: Vec, 354 | ) -> Result { 355 | let mut state_alloc = preimg_alloc; 356 | 357 | for round in 0..R_f { 358 | state_alloc = self.alloc_full_round(layouter, &state_alloc, round)?; 359 | } 360 | 361 | for round in R_f..R_f + R_P { 362 | state_alloc = self.alloc_partial_round(layouter, &state_alloc, round)?; 363 | } 364 | 365 | for round in R_f + R_P..R { 366 | state_alloc = self.alloc_full_round(layouter, &state_alloc, round)?; 367 | } 368 | 369 | Ok(state_alloc.remove(1)) 370 | } 371 | 372 | fn alloc_full_round( 373 | &self, 374 | layouter: &mut impl Layouter, 375 | state_alloc: &[Alloc], 376 | round: usize, 377 | ) -> Result, Error> { 378 | assert!(is_full_round(round)); 379 | assert_eq!(state_alloc.len(), WIDTH); 380 | 381 | let is_first_round = round == 0; 382 | let is_last_round = round == R - 1; 383 | let is_pre_sparse_round = round == R_f - 1; 384 | 385 | let post_keys = &*POST_KEYS[round]; 386 | assert_eq!(post_keys.len(), if is_last_round { 0 } else { WIDTH }); 387 | 388 | // Copy field elements out of the `Alloc`'s for easier arithmetic. 389 | let state_values: Vec = state_alloc 390 | .iter() 391 | .map(|alloc| alloc.value) 392 | .collect(); 393 | 394 | // Calculate the S-box output for each state element. 395 | let sbox_outputs = if is_first_round { 396 | sbox_pre_post(&state_values) 397 | } else if is_last_round { 398 | sbox_no_add(&state_values) 399 | } else { 400 | sbox_post(&state_values, &post_keys) 401 | }; 402 | 403 | // Calculate the MDS mixing output for each state element. 404 | let m = if is_pre_sparse_round { &*PRE_SPARSE } else { &*MDS }; 405 | let mds_outputs = vec_matrix_mul(&sbox_outputs, m); 406 | 407 | // Store the allocated outputs of MDS mixing. 408 | let mut mds_outputs_alloc: Vec> = vec![None; WIDTH]; 409 | 410 | layouter.assign_region( 411 | || format!("alloc (full) round {}", round), 412 | |mut region| { 413 | for row_offset in 0..WIDTH { 414 | let a_cell = region.assign_advice( 415 | || format!("a_{} (round {})", row_offset, round), 416 | self.config.a_col, 417 | row_offset, 418 | || Ok(state_values[row_offset]), 419 | )?; 420 | 421 | // If this is the first round, we reallocate values from the preimage row. 422 | // Otherwise, we reallocate values output by the previous round. 423 | let perm = if is_first_round { 424 | &self.config.perm_preimg[row_offset] 425 | } else { 426 | &self.config.perm_output_to_input 427 | }; 428 | region.constrain_equal(perm, state_alloc[row_offset].cell, a_cell)?; 429 | 430 | // If this is the first round allocate a pre-S-box key. 431 | if is_first_round { 432 | region.assign_fixed( 433 | || format!("pre_key_{} (round {})", row_offset, round), 434 | self.config.pre_key_col, 435 | row_offset, 436 | || Ok(PRE_KEYS[row_offset]), 437 | )?; 438 | } 439 | 440 | // If this is not the last round allocate a post-S-box key. 441 | if !is_last_round { 442 | region.assign_fixed( 443 | || format!("post_key_{} (round {})", row_offset, round), 444 | self.config.post_key_col, 445 | row_offset, 446 | || Ok(post_keys[row_offset]), 447 | )?; 448 | } 449 | 450 | // Allocate the S-box output. 451 | region.assign_advice( 452 | || format!("sbox_out_{} (round {})", row_offset, round), 453 | self.config.sbox_out_col, 454 | row_offset, 455 | || Ok(sbox_outputs[row_offset]), 456 | )?; 457 | 458 | // Allocate the MDS mixing output. 459 | let mds_output_cell = region.assign_advice( 460 | || format!("mds_out_{} (round {})", row_offset, round), 461 | self.config.mds_out_col, 462 | row_offset, 463 | || Ok(mds_outputs[row_offset]), 464 | )?; 465 | 466 | // Keep a reference to the allocated MDS output. 467 | mds_outputs_alloc[row_offset] = Some(Alloc { 468 | cell: mds_output_cell, 469 | value: mds_outputs[row_offset], 470 | }); 471 | 472 | // Enable the S-box and MDS mixing selectors. 473 | if is_first_round { 474 | self.config.s_sbox_pre_post.enable(&mut region, row_offset)?; 475 | } else if is_last_round { 476 | self.config.s_sbox_no_add.enable(&mut region, row_offset)?; 477 | } else { 478 | self.config.s_sbox_post.enable(&mut region, row_offset)?; 479 | }; 480 | self.config.s_mds[row_offset].enable(&mut region, row_offset)?; 481 | 482 | // Allocate this MDS matrix row. 483 | for col in 0..WIDTH { 484 | region.assign_fixed( 485 | || format!( 486 | "{} row={}, col={} (round {})", 487 | if is_pre_sparse_round { "P" } else { "MDS" }, 488 | row_offset, 489 | col, 490 | round, 491 | ), 492 | self.config.mds_cols[col], 493 | row_offset, 494 | || Ok(m[row_offset][col]), 495 | )?; 496 | } 497 | } 498 | Ok(()) 499 | }, 500 | )?; 501 | 502 | let mds_outputs_alloc: Vec = mds_outputs_alloc 503 | .into_iter() 504 | .map(|opt| opt.unwrap()) 505 | .collect(); 506 | 507 | Ok(mds_outputs_alloc) 508 | } 509 | 510 | fn alloc_partial_round( 511 | &self, 512 | layouter: &mut impl Layouter, 513 | state_alloc: &[Alloc], 514 | round: usize, 515 | ) -> Result, Error> { 516 | assert!(is_partial_round(round)); 517 | assert_eq!(state_alloc.len(), WIDTH); 518 | 519 | assert_eq!(POST_KEYS[round].len(), 1); 520 | let post_key = POST_KEYS[round][0]; 521 | 522 | // Copy field elements out of `Alloc`'s for easier arithmetic. 523 | let state_values: Vec = state_alloc 524 | .iter() 525 | .map(|alloc| alloc.value) 526 | .collect(); 527 | 528 | // Calculate the S-box output for the first state element. 529 | let mut sbox_outputs: Vec = vec![pow5(state_values[0]) + post_key]; 530 | sbox_outputs.extend_from_slice(&state_values[1..]); 531 | 532 | // Calculate the MDS mixing output for each state element. 533 | let sparse_index = round - R_f; 534 | let m = &*SPARSE[sparse_index]; 535 | let mds_outputs = vec_matrix_mul(&sbox_outputs, m); 536 | 537 | // Store the allocated outputs of MDS mixing. 538 | let mut mds_outputs_alloc: Vec> = vec![None; WIDTH]; 539 | 540 | layouter.assign_region( 541 | || format!("alloc (partial) round {}", round), 542 | |mut region| { 543 | // Allocate values that are exclusive to the first row. 544 | let row_offset = 0; 545 | 546 | // Reallocate the first state element which was output by the previous round. 547 | let a_cell = region.assign_advice( 548 | || format!("a_0 (round {})", round), 549 | self.config.a_col, 550 | row_offset, 551 | || Ok(state_values[0]), 552 | )?; 553 | 554 | region.constrain_equal( 555 | &self.config.perm_output_to_input, 556 | state_alloc[0].cell, 557 | a_cell, 558 | )?; 559 | 560 | // Allocate the first state element's post-S-box key. 561 | region.assign_fixed( 562 | || format!("post_key_0 (round {})", round), 563 | self.config.post_key_col, 564 | row_offset, 565 | || Ok(post_key), 566 | )?; 567 | 568 | // Enable the first state element's S-box selector. 569 | self.config.s_sbox_post.enable(&mut region, row_offset)?; 570 | 571 | // Allocate the remaining round values. 572 | for row_offset in 0..WIDTH { 573 | // If this is the first row (`row_offset = 0`), allocate the first state 574 | // element's S-box output. If this is not the first row, reallocate the state 575 | // element output by the previous round. 576 | let sbox_out_cell = region.assign_advice( 577 | || format!("sbox_out_{} (round {})", row_offset, round), 578 | self.config.sbox_out_col, 579 | row_offset, 580 | || Ok(sbox_outputs[row_offset]), 581 | )?; 582 | 583 | if row_offset > 0 { 584 | region.constrain_equal( 585 | &self.config.perm_output_to_sbox_output, 586 | state_alloc[row_offset].cell, 587 | sbox_out_cell, 588 | )?; 589 | } 590 | 591 | // Allocate the state element's MDS mixing output. 592 | let mds_out_cell = region.assign_advice( 593 | || format!("mds_out_{} (round {})", row_offset, round), 594 | self.config.mds_out_col, 595 | row_offset, 596 | || Ok(mds_outputs[row_offset]), 597 | )?; 598 | 599 | // Keep a reference to the allocated MDS output. 600 | mds_outputs_alloc[row_offset] = Some(Alloc { 601 | cell: mds_out_cell, 602 | value: mds_outputs[row_offset], 603 | }); 604 | 605 | // Enable the MDS mixing selector for this state element. 606 | self.config.s_mds[row_offset].enable(&mut region, row_offset)?; 607 | 608 | // Allocate this MDS matrix row. 609 | for col in 0..WIDTH { 610 | region.assign_fixed( 611 | || format!( 612 | "S{} row={}, col={} (round {})", 613 | sparse_index, 614 | row_offset, 615 | col, 616 | round, 617 | ), 618 | self.config.mds_cols[col], 619 | row_offset, 620 | || Ok(m[row_offset][col]), 621 | )?; 622 | } 623 | } 624 | Ok(()) 625 | }, 626 | )?; 627 | 628 | let mds_outputs_alloc: Vec = mds_outputs_alloc 629 | .into_iter() 630 | .map(|opt| opt.unwrap()) 631 | .collect(); 632 | 633 | Ok(mds_outputs_alloc) 634 | } 635 | } 636 | 637 | struct HasherCircuit { 638 | // Private inputs. 639 | preimg: Vec, 640 | } 641 | 642 | #[derive(Clone)] 643 | struct HasherCircuitConfig { 644 | preimg_cols: Vec>, 645 | digest_col: Column, 646 | pub_col: Column, 647 | s_pub: Selector, 648 | perm_digest: Permutation, 649 | poseidon_config: PoseidonChipConfig, 650 | } 651 | 652 | impl Circuit for HasherCircuit { 653 | type Config = HasherCircuitConfig; 654 | 655 | fn configure(cs: &mut ConstraintSystem) -> Self::Config { 656 | let (poseidon_config, preimg_cols, digest_col) = PoseidonChip::configure(cs); 657 | let perm_digest = Permutation::new(cs, &[digest_col.into()]); 658 | let pub_col = cs.instance_column(); 659 | let s_pub = cs.selector(); 660 | 661 | cs.create_gate("public input", |cs| { 662 | let digest = cs.query_advice(digest_col, Rotation::cur()); 663 | let pi = cs.query_instance(pub_col, Rotation::cur()); 664 | let s_pub = cs.query_selector(s_pub, Rotation::cur()); 665 | s_pub * (digest - pi) 666 | }); 667 | 668 | HasherCircuitConfig { 669 | preimg_cols, 670 | digest_col, 671 | pub_col, 672 | s_pub, 673 | perm_digest, 674 | poseidon_config, 675 | } 676 | } 677 | 678 | fn synthesize(&self, cs: &mut impl Assignment, config: Self::Config) -> Result<(), Error> { 679 | assert_eq!(self.preimg.len(), WIDTH); 680 | let mut layouter = SingleChipLayouter::new(cs)?; 681 | let poseidon_chip = PoseidonChip::new(config.poseidon_config.clone()); 682 | 683 | // Allocate the preimage in the first row of the constraint system. 684 | let mut preimg_alloc: Vec> = vec![None; WIDTH]; 685 | 686 | layouter.assign_region( 687 | || "alloc preimage", 688 | |mut region| { 689 | let row_offset = 0; 690 | for i in 0..WIDTH { 691 | let preimg_cell = region.assign_advice( 692 | || format!("preimg_{}", i), 693 | config.preimg_cols[i], 694 | row_offset, 695 | || Ok(self.preimg[i]), 696 | )?; 697 | preimg_alloc[i] = Some(Alloc { 698 | cell: preimg_cell, 699 | value: self.preimg[i].clone(), 700 | }); 701 | } 702 | Ok(()) 703 | }, 704 | )?; 705 | 706 | let preimg_alloc: Vec = preimg_alloc.into_iter().map(|opt| opt.unwrap()).collect(); 707 | let digest_alloc = poseidon_chip.hash(&mut layouter, preimg_alloc)?; 708 | dbg!(digest_alloc.value); 709 | 710 | // Copy the digest out of the chip and check that agrees with the public input. 711 | layouter.assign_region( 712 | || "digest equality", 713 | |mut region| { 714 | let row_offset = 0; 715 | let copied_digest_cell = region.assign_advice( 716 | || "copied digest", 717 | config.digest_col, 718 | row_offset, 719 | || Ok(digest_alloc.value), 720 | )?; 721 | region.constrain_equal(&config.perm_digest, digest_alloc.cell, copied_digest_cell)?; 722 | config.s_pub.enable(&mut region, row_offset)?; 723 | Ok(()) 724 | }, 725 | )?; 726 | 727 | Ok(()) 728 | } 729 | } 730 | 731 | fn main() { 732 | // There are `WIDTH` number of rows per round; add one row for allocating the preimage and one 733 | // row for checking that the calculated digest agrees with the public input digest. 734 | const N_ROWS_USED: u32 = (R * WIDTH + 2) as u32; 735 | const PUB_INPUT_ROW_INDEX: usize = N_ROWS_USED as usize - 1; 736 | 737 | // The public digest. 738 | let pub_input = Fp::from_bytes(&[ 739 | 92, 156, 142, 19, 149, 223, 255, 67, 5, 168, 740 | 243, 206, 123, 14, 94, 31, 226, 187, 207, 47, 741 | 97, 158, 70, 2, 132, 63, 106, 142, 219, 243, 742 | 144, 17, 743 | ]).unwrap(); 744 | 745 | // Verifier's public inputs. 746 | let k = (N_ROWS_USED as f32).log2().ceil() as u32; 747 | let n_rows = 1 << k; 748 | let mut pub_inputs = vec![Fp::zero(); n_rows]; 749 | pub_inputs[PUB_INPUT_ROW_INDEX] = Fp::from(pub_input); 750 | 751 | // Prover's private inputs. 752 | let preimg = vec![Fp::from(55), Fp::from(101), Fp::from(237)]; 753 | dbg!(poseidon(&preimg)); 754 | // println!("{:?}", poseidon(&preimg).to_bytes()); 755 | let circuit = HasherCircuit { preimg }; 756 | 757 | let prover = MockProver::run(k, &circuit, vec![pub_inputs.clone()]).unwrap(); 758 | // dbg!(prover.verify()); 759 | assert!(prover.verify().is_ok()); 760 | } 761 | --------------------------------------------------------------------------------