1#![forbid(unsafe_code)]
17#![allow(clippy::too_many_arguments)]
18#![warn(clippy::cast_possible_truncation)]
19
20mod partial_solution;
21pub use partial_solution::*;
22
23mod solution;
24pub use solution::*;
25
26mod solution_id;
27pub use solution_id::*;
28
29mod solutions;
30pub use solutions::*;
31
32use console::{
33 account::Address,
34 algorithms::Sha3_256,
35 collections::kary_merkle_tree::KaryMerkleTree,
36 prelude::{
37 FromBits,
38 Network,
39 Result,
40 anyhow,
41 bail,
42 cfg_into_iter,
43 cfg_iter,
44 cfg_keys,
45 cfg_values,
46 ensure,
47 has_duplicates,
48 },
49 types::U64,
50};
51
52use aleo_std::prelude::*;
53use core::num::NonZeroUsize;
54use indexmap::IndexMap;
55#[cfg(feature = "locktick")]
56use locktick::parking_lot::RwLock;
57use lru::LruCache;
58#[cfg(not(feature = "locktick"))]
59use parking_lot::RwLock;
60use rand::SeedableRng;
61use rand_chacha::ChaChaRng;
62use std::sync::Arc;
63
64#[cfg(not(feature = "serial"))]
65use rayon::prelude::*;
66
67const ARITY: u8 = 8;
69const CACHE_SIZE: usize = 1 << 10;
71
72type MerkleTree = KaryMerkleTree<Sha3_256, Sha3_256, 9, { ARITY }>;
74
75pub trait PuzzleTrait<N: Network>: Send + Sync {
77 fn new() -> Self
79 where
80 Self: Sized;
81
82 fn to_leaves(&self, epoch_hash: N::BlockHash, rng: &mut ChaChaRng) -> Result<Vec<Vec<bool>>>;
84
85 fn to_all_leaves(&self, epoch_hash: N::BlockHash, rngs: Vec<ChaChaRng>) -> Result<Vec<Vec<Vec<bool>>>>;
87}
88
89#[derive(Clone)]
90pub struct Puzzle<N: Network> {
91 inner: Arc<dyn PuzzleTrait<N>>,
93 proof_target_cache: Arc<RwLock<LruCache<SolutionID<N>, u64>>>,
95}
96
97impl<N: Network> Puzzle<N> {
98 pub fn new<P: PuzzleTrait<N> + 'static>() -> Self {
100 Self {
101 inner: Arc::new(P::new()),
102 proof_target_cache: Arc::new(RwLock::new(LruCache::new(NonZeroUsize::new(CACHE_SIZE).unwrap()))),
103 }
104 }
105
106 pub fn get_leaves(&self, solution: &PartialSolution<N>) -> Result<Vec<Vec<bool>>> {
108 let mut rng = ChaChaRng::seed_from_u64(*solution.id());
110 self.inner.to_leaves(solution.epoch_hash(), &mut rng)
112 }
113
114 pub fn get_all_leaves(&self, solutions: &PuzzleSolutions<N>) -> Result<Vec<Vec<Vec<bool>>>> {
116 ensure!(
118 cfg_values!(solutions).all(|solution| solution.epoch_hash() == solutions[0].epoch_hash()),
119 "The solutions are for different epochs"
120 );
121 let rngs = cfg_keys!(solutions).map(|solution_id| ChaChaRng::seed_from_u64(**solution_id)).collect::<Vec<_>>();
123 self.inner.to_all_leaves(solutions[0].epoch_hash(), rngs)
125 }
126
127 pub fn get_proof_target(&self, solution: &Solution<N>) -> Result<u64> {
129 let proof_target = self.get_proof_target_unchecked(solution)?;
131 ensure!(solution.target() == proof_target, "The proof target does not match the expected proof target");
133 Ok(proof_target)
135 }
136
137 pub fn get_proof_target_unchecked(&self, solution: &Solution<N>) -> Result<u64> {
141 self.get_proof_target_from_partial_solution(solution.partial_solution())
143 }
144
145 pub fn get_proof_target_from_partial_solution(&self, partial_solution: &PartialSolution<N>) -> Result<u64> {
147 if let Some(proof_target) = self.proof_target_cache.write().get(&partial_solution.id()) {
149 return Ok(*proof_target);
150 }
151
152 let leaves = self.get_leaves(partial_solution)?;
154 let proof_target = Self::leaves_to_proof_target(&leaves)?;
156
157 self.proof_target_cache.write().put(partial_solution.id(), proof_target);
159 Ok(proof_target)
161 }
162
163 pub fn get_proof_targets(&self, solutions: &PuzzleSolutions<N>) -> Result<Vec<u64>> {
165 let mut targets = vec![0u64; solutions.len()];
167
168 let mut to_compute = Vec::new();
170 for (i, (id, solution)) in solutions.iter().enumerate() {
172 match self.proof_target_cache.write().get(id) {
174 Some(proof_target) => {
176 ensure!(
178 solution.target() == *proof_target,
179 "The proof target does not match the cached proof target"
180 );
181 targets[i] = *proof_target
182 }
183 None => to_compute.push((i, id, *solution)),
185 }
186 }
187
188 if !to_compute.is_empty() {
189 let solutions_subset = PuzzleSolutions::new(to_compute.iter().map(|(_, _, solution)| *solution).collect())?;
191 let leaves = self.get_all_leaves(&solutions_subset)?;
193 let targets_subset = cfg_iter!(leaves)
195 .zip(cfg_iter!(solutions_subset))
196 .map(|(leaves, (solution_id, solution))| {
197 let proof_target = Self::leaves_to_proof_target(leaves)?;
199 ensure!(
201 solution.target() == proof_target,
202 "The proof target does not match the computed proof target"
203 );
204 self.proof_target_cache.write().put(*solution_id, proof_target);
206 Ok((solution_id, proof_target))
208 })
209 .collect::<Result<IndexMap<_, _>>>()?;
210
211 for (i, id, _) in &to_compute {
213 targets[*i] = targets_subset[id];
214 }
215 }
216
217 Ok(targets)
219 }
220
221 pub fn get_combined_proof_target(&self, solutions: &PuzzleSolutions<N>) -> Result<u128> {
223 self.get_proof_targets(solutions)?.into_iter().try_fold(0u128, |combined, proof_target| {
224 combined.checked_add(proof_target as u128).ok_or_else(|| anyhow!("Combined proof target overflowed"))
225 })
226 }
227
228 pub fn prove(
230 &self,
231 epoch_hash: N::BlockHash,
232 address: Address<N>,
233 counter: u64,
234 minimum_proof_target: Option<u64>,
235 ) -> Result<Solution<N>> {
236 let partial_solution = PartialSolution::new(epoch_hash, address, counter)?;
238 let proof_target = self.get_proof_target_from_partial_solution(&partial_solution)?;
240 if let Some(minimum_proof_target) = minimum_proof_target {
242 if proof_target < minimum_proof_target {
243 bail!("Solution was below the minimum proof target ({proof_target} < {minimum_proof_target})")
244 }
245 }
246
247 Ok(Solution::new(partial_solution, proof_target))
249 }
250
251 pub fn check_solution(
253 &self,
254 solution: &Solution<N>,
255 expected_epoch_hash: N::BlockHash,
256 expected_proof_target: u64,
257 ) -> Result<()> {
258 if solution.epoch_hash() != expected_epoch_hash {
260 bail!(
261 "Solution does not match the expected epoch hash (found '{}', expected '{expected_epoch_hash}')",
262 solution.epoch_hash()
263 )
264 }
265 let proof_target = self.get_proof_target(solution)?;
267 if proof_target < expected_proof_target {
268 bail!("Solution does not meet the proof target requirement ({proof_target} < {expected_proof_target})")
269 }
270 Ok(())
271 }
272
273 pub fn check_solution_mut(
276 &self,
277 solution: &mut Solution<N>,
278 expected_epoch_hash: N::BlockHash,
279 expected_proof_target: u64,
280 ) -> Result<()> {
281 if solution.epoch_hash() != expected_epoch_hash {
283 bail!(
284 "Solution does not match the expected epoch hash (found '{}', expected '{expected_epoch_hash}')",
285 solution.epoch_hash()
286 )
287 }
288 let proof_target = self.get_proof_target_unchecked(solution)?;
290
291 solution.target = proof_target;
293
294 if proof_target < expected_proof_target {
296 bail!("Solution does not meet the proof target requirement ({proof_target} < {expected_proof_target})")
297 }
298 Ok(())
299 }
300
301 pub fn check_solutions(
303 &self,
304 solutions: &PuzzleSolutions<N>,
305 expected_epoch_hash: N::BlockHash,
306 expected_proof_target: u64,
307 ) -> Result<()> {
308 let timer = timer!("Puzzle::verify");
309
310 ensure!(!solutions.is_empty(), "The solutions are empty");
312 if solutions.len() > N::MAX_SOLUTIONS {
314 bail!("Exceed the maximum number of solutions ({} > {})", solutions.len(), N::MAX_SOLUTIONS)
315 }
316 if has_duplicates(solutions.solution_ids()) {
318 bail!("The solutions contain duplicate solution IDs");
319 }
320 lap!(timer, "Perform initial checks");
321
322 cfg_iter!(solutions).try_for_each(|(solution_id, solution)| {
324 if solution.epoch_hash() != expected_epoch_hash {
325 bail!("Solution '{solution_id}' did not match the expected epoch hash (found '{}', expected '{expected_epoch_hash}')", solution.epoch_hash())
326 }
327 Ok(())
328 })?;
329 lap!(timer, "Verify each epoch hash matches");
330
331 cfg_into_iter!(self.get_proof_targets(solutions)?).enumerate().try_for_each(|(i, proof_target)| {
333 if proof_target < expected_proof_target {
334 bail!(
335 "Solution '{:?}' did not meet the proof target requirement ({proof_target} < {expected_proof_target})",
336 solutions.get_index(i).map(|(id, _)| id)
337 )
338 }
339 Ok(())
340 })?;
341 finish!(timer, "Verify each solution");
342 Ok(())
343 }
344
345 fn leaves_to_proof_target(leaves: &[Vec<bool>]) -> Result<u64> {
347 let merkle_tree = MerkleTree::new(&Sha3_256::default(), &Sha3_256::default(), leaves)?;
349 let root = merkle_tree.root();
351 match *U64::<N>::from_bits_be(&root[0..64])? {
353 0 => Ok(u64::MAX),
354 value => Ok(u64::MAX / value),
355 }
356 }
357}
358
359#[cfg(test)]
360mod tests {
361 use super::*;
362 use console::{
363 account::{Address, PrivateKey},
364 network::Network,
365 prelude::{FromBytes, TestRng, ToBits as TBits, ToBytes, Uniform},
366 types::Field,
367 };
368
369 use anyhow::Result;
370 use core::marker::PhantomData;
371 use rand::{CryptoRng, Rng, RngCore, SeedableRng};
372 use rand_chacha::ChaChaRng;
373
374 type CurrentNetwork = console::network::MainnetV0;
375
376 const ITERATIONS: u64 = 100;
377
378 pub struct SimplePuzzle<N: Network>(PhantomData<N>);
379
380 impl<N: Network> PuzzleTrait<N> for SimplePuzzle<N> {
381 fn new() -> Self {
383 Self(PhantomData)
384 }
385
386 fn to_leaves(&self, epoch_hash: N::BlockHash, rng: &mut ChaChaRng) -> Result<Vec<Vec<bool>>> {
388 let num_leaves = self.num_leaves(epoch_hash)?;
390 let leaves = (0..num_leaves).map(|_| Field::<N>::rand(rng).to_bits_le()).collect::<Vec<_>>();
392 Ok(leaves)
394 }
395
396 fn to_all_leaves(&self, epoch_hash: N::BlockHash, rngs: Vec<ChaChaRng>) -> Result<Vec<Vec<Vec<bool>>>> {
398 let num_leaves = self.num_leaves(epoch_hash)?;
400 let mut leaves = Vec::with_capacity(rngs.len());
402 for mut rng in rngs {
404 leaves.push((0..num_leaves).map(|_| Field::<N>::rand(&mut rng).to_bits_le()).collect::<Vec<_>>());
406 }
407 Ok(leaves)
409 }
410 }
411
412 impl<N: Network> SimplePuzzle<N> {
413 pub fn num_leaves(&self, epoch_hash: N::BlockHash) -> Result<usize> {
415 const MIN_NUMBER_OF_LEAVES: usize = 100;
416 const MAX_NUMBER_OF_LEAVES: usize = 200;
417
418 let seed = u64::from_bytes_le(&epoch_hash.to_bytes_le()?[0..8])?;
420 let mut epoch_rng = ChaChaRng::seed_from_u64(seed);
422 Ok(epoch_rng.gen_range(MIN_NUMBER_OF_LEAVES..MAX_NUMBER_OF_LEAVES))
424 }
425 }
426
427 fn sample_puzzle() -> Puzzle<CurrentNetwork> {
429 Puzzle::<CurrentNetwork>::new::<SimplePuzzle<CurrentNetwork>>()
430 }
431
432 #[test]
433 fn test_puzzle() {
434 let mut rng = TestRng::default();
435
436 let puzzle = sample_puzzle();
438
439 let epoch_hash = rng.gen();
441
442 for batch_size in 1..=CurrentNetwork::MAX_SOLUTIONS {
443 let solutions = (0..batch_size)
445 .map(|_| puzzle.prove(epoch_hash, rng.gen(), rng.gen(), None).unwrap())
446 .collect::<Vec<_>>();
447 let solutions = PuzzleSolutions::new(solutions).unwrap();
448
449 assert!(puzzle.check_solutions(&solutions, epoch_hash, 0u64).is_ok());
451
452 let bad_epoch_hash = rng.gen();
454 assert!(puzzle.check_solutions(&solutions, bad_epoch_hash, 0u64).is_err());
455 }
456 }
457
458 #[test]
459 fn test_prove_with_minimum_proof_target() {
460 let mut rng = TestRng::default();
461
462 let puzzle = sample_puzzle();
464
465 let epoch_hash = rng.gen();
467
468 for _ in 0..ITERATIONS {
469 let private_key = PrivateKey::<CurrentNetwork>::new(&mut rng).unwrap();
470 let address = Address::try_from(private_key).unwrap();
471 let counter = u64::rand(&mut rng);
472
473 let solution = puzzle.prove(epoch_hash, address, counter, None).unwrap();
474 let proof_target = puzzle.get_proof_target(&solution).unwrap();
475
476 assert!(puzzle.prove(epoch_hash, address, counter, Some(proof_target)).is_ok());
478
479 assert!(puzzle.prove(epoch_hash, address, counter, Some(proof_target.saturating_add(1))).is_err());
481
482 let solutions = PuzzleSolutions::new(vec![solution]).unwrap();
483
484 assert!(puzzle.check_solutions(&solutions, epoch_hash, proof_target).is_ok());
486
487 assert!(puzzle.check_solutions(&solutions, epoch_hash, proof_target.saturating_add(1)).is_err());
489 }
490 }
491
492 #[test]
493 fn test_prove_with_no_minimum_proof_target() {
494 let mut rng = rand::thread_rng();
495
496 let puzzle = sample_puzzle();
498
499 let epoch_hash = rng.gen();
501
502 let private_key = PrivateKey::<CurrentNetwork>::new(&mut rng).unwrap();
504 let address = Address::try_from(private_key).unwrap();
505
506 let solution = puzzle.prove(epoch_hash, address, rng.gen(), None).unwrap();
508 assert!(puzzle.check_solution(&solution, epoch_hash, 0u64).is_ok());
509
510 let solutions = PuzzleSolutions::new(vec![solution]).unwrap();
511 assert!(puzzle.check_solutions(&solutions, epoch_hash, 0u64).is_ok());
512 }
513
514 #[test]
515 fn test_check_solution_with_incorrect_target_fails() {
516 let mut rng = rand::thread_rng();
517
518 let puzzle = sample_puzzle();
520
521 let epoch_hash = rng.gen();
523
524 let private_key = PrivateKey::<CurrentNetwork>::new(&mut rng).unwrap();
526 let address = Address::try_from(private_key).unwrap();
527
528 let solution = puzzle.prove(epoch_hash, address, rng.gen(), None).unwrap();
530
531 let incorrect_solution = Solution::new(*solution.partial_solution(), solution.target().saturating_add(1));
533
534 assert!(puzzle.check_solution(&incorrect_solution, epoch_hash, 0u64).is_err());
536
537 let new_puzzle = sample_puzzle();
539 assert!(new_puzzle.check_solution(&incorrect_solution, epoch_hash, 0u64).is_err());
540
541 let incorrect_solutions = PuzzleSolutions::new(vec![incorrect_solution]).unwrap();
543 assert!(puzzle.check_solutions(&incorrect_solutions, epoch_hash, 0u64).is_err());
544
545 let new_puzzle = sample_puzzle();
547 assert!(new_puzzle.check_solutions(&incorrect_solutions, epoch_hash, 0u64).is_err());
548 }
549
550 #[test]
551 fn test_check_solutions_with_incorrect_target_fails() {
552 let mut rng = TestRng::default();
553
554 let puzzle = sample_puzzle();
556
557 let epoch_hash = rng.gen();
559
560 for batch_size in 1..=CurrentNetwork::MAX_SOLUTIONS {
561 let incorrect_solutions = (0..batch_size)
563 .map(|_| {
564 let solution = puzzle.prove(epoch_hash, rng.gen(), rng.gen(), None).unwrap();
565 Solution::new(*solution.partial_solution(), solution.target().saturating_add(1))
566 })
567 .collect::<Vec<_>>();
568 let incorrect_solutions = PuzzleSolutions::new(incorrect_solutions).unwrap();
569
570 assert!(puzzle.check_solutions(&incorrect_solutions, epoch_hash, 0u64).is_err());
572
573 let new_puzzle = sample_puzzle();
575 assert!(new_puzzle.check_solutions(&incorrect_solutions, epoch_hash, 0u64).is_err());
576 }
577 }
578
579 #[test]
580 fn test_check_solutions_with_duplicate_nonces() {
581 let mut rng = TestRng::default();
582
583 let puzzle = sample_puzzle();
585
586 let epoch_hash = rng.gen();
588 let address = rng.gen();
590 let counter = rng.gen();
592
593 for batch_size in 1..=CurrentNetwork::MAX_SOLUTIONS {
594 let solutions =
596 (0..batch_size).map(|_| puzzle.prove(epoch_hash, address, counter, None).unwrap()).collect::<Vec<_>>();
597 let solutions = match batch_size {
599 1 => PuzzleSolutions::new(solutions).unwrap(),
600 _ => {
601 assert!(PuzzleSolutions::new(solutions).is_err());
602 continue;
603 }
604 };
605 match batch_size {
606 1 => assert!(puzzle.check_solutions(&solutions, epoch_hash, 0u64).is_ok()),
607 _ => unreachable!("There are duplicates that should not reach this point in the test"),
608 }
609 }
610 }
611
612 #[test]
613 fn test_get_proof_targets_without_cache() {
614 let mut rng = TestRng::default();
615
616 let epoch_hash = rng.gen();
618
619 for batch_size in 1..=CurrentNetwork::MAX_SOLUTIONS {
620 let puzzle = sample_puzzle();
622 let solutions = (0..batch_size)
624 .map(|_| puzzle.prove(epoch_hash, rng.gen(), rng.gen(), None).unwrap())
625 .collect::<Vec<_>>();
626 let solutions = PuzzleSolutions::new(solutions).unwrap();
627
628 let puzzle = sample_puzzle();
630
631 let proof_targets = puzzle.get_proof_targets(&solutions).unwrap();
633
634 for ((_, solution), proof_target) in solutions.iter().zip(proof_targets) {
636 assert_eq!(puzzle.get_proof_target(solution).unwrap(), proof_target);
637 }
638 }
639 }
640
641 #[test]
642 fn test_get_proof_targets_with_partial_cache() {
643 let mut rng = TestRng::default();
644
645 let epoch_hash = rng.gen();
647
648 for batch_size in 1..=CurrentNetwork::MAX_SOLUTIONS {
649 let puzzle = sample_puzzle();
651 let solutions = (0..batch_size)
653 .map(|_| puzzle.prove(epoch_hash, rng.gen(), rng.gen(), None).unwrap())
654 .collect::<Vec<_>>();
655 let solutions = PuzzleSolutions::new(solutions).unwrap();
656
657 let puzzle = sample_puzzle();
659
660 for solution in solutions.values() {
662 if rng.gen::<bool>() {
664 puzzle.get_proof_target(solution).unwrap();
666 }
667 }
668
669 let proof_targets = puzzle.get_proof_targets(&solutions).unwrap();
671
672 for ((_, solution), proof_target) in solutions.iter().zip(proof_targets) {
674 assert_eq!(puzzle.get_proof_target(solution).unwrap(), proof_target);
675 }
676 }
677 }
678
679 #[ignore]
681 #[test]
682 fn test_profiler() -> Result<()> {
683 fn sample_address_and_counter(rng: &mut (impl CryptoRng + RngCore)) -> (Address<CurrentNetwork>, u64) {
684 let private_key = PrivateKey::new(rng).unwrap();
685 let address = Address::try_from(private_key).unwrap();
686 let counter = rng.next_u64();
687 (address, counter)
688 }
689
690 let mut rng = rand::thread_rng();
691
692 let puzzle = sample_puzzle();
694
695 let epoch_hash = rng.gen();
697
698 for batch_size in [1, 2, <CurrentNetwork as Network>::MAX_SOLUTIONS] {
699 let solutions = (0..batch_size)
701 .map(|_| {
702 let (address, counter) = sample_address_and_counter(&mut rng);
703 puzzle.prove(epoch_hash, address, counter, None).unwrap()
704 })
705 .collect::<Vec<_>>();
706 let solutions = PuzzleSolutions::new(solutions).unwrap();
708 puzzle.check_solutions(&solutions, epoch_hash, 0u64).unwrap();
710 }
711
712 bail!("\n\nRemember to #[ignore] this test!\n\n")
713 }
714}