darkfi_sdk/crypto/smt/
util.rs

1/* This file is part of DarkFi (https://dark.fi)
2 *
3 * Copyright (C) 2020-2025 Dyne.org foundation
4 *
5 * This program is free software: you can redistribute it and/or modify
6 * it under the terms of the GNU Affero General Public License as
7 * published by the Free Software Foundation, either version 3 of the
8 * License, or (at your option) any later version.
9 *
10 * This program is distributed in the hope that it will be useful,
11 * but WITHOUT ANY WARRANTY; without even the implied warranty of
12 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13 * GNU Affero General Public License for more details.
14 *
15 * You should have received a copy of the GNU Affero General Public License
16 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
17 */
18
19use std::marker::PhantomData;
20
21use halo2_gadgets::poseidon::{
22    primitives as poseidon,
23    primitives::{ConstantLength, P128Pow5T3, Spec},
24};
25use num_bigint::BigUint;
26use pasta_curves::group::ff::{PrimeField, WithSmallOrderMulGroup};
27
28pub trait FieldElement: WithSmallOrderMulGroup<3> + Ord + PrimeField {
29    fn as_biguint(&self) -> BigUint;
30}
31
32impl FieldElement for pasta_curves::Fp {
33    fn as_biguint(&self) -> BigUint {
34        let repr = self.to_repr();
35        BigUint::from_bytes_le(&repr)
36    }
37}
38
39impl FieldElement for pasta_curves::Fq {
40    fn as_biguint(&self) -> BigUint {
41        let repr = self.to_repr();
42        BigUint::from_bytes_le(&repr)
43    }
44}
45
46pub trait FieldHasher<F: WithSmallOrderMulGroup<3> + Ord, const L: usize>: Clone {
47    fn hash(&self, inputs: [F; L]) -> F;
48    fn hasher() -> Self;
49}
50
51#[derive(Debug, Clone)]
52pub struct Poseidon<F: WithSmallOrderMulGroup<3> + Ord, const L: usize>(PhantomData<F>);
53
54impl<F: WithSmallOrderMulGroup<3> + Ord, const L: usize> Poseidon<F, L> {
55    pub fn new() -> Self {
56        Poseidon(PhantomData)
57    }
58}
59
60impl<F: WithSmallOrderMulGroup<3> + Ord, const L: usize> Default for Poseidon<F, L> {
61    fn default() -> Self {
62        Self::new()
63    }
64}
65
66impl<F: WithSmallOrderMulGroup<3> + Ord, const L: usize> FieldHasher<F, L> for Poseidon<F, L>
67where
68    P128Pow5T3: Spec<F, 3, 2>,
69{
70    fn hash(&self, inputs: [F; L]) -> F {
71        poseidon::Hash::<_, P128Pow5T3, ConstantLength<L>, 3, 2>::init().hash(inputs)
72    }
73
74    fn hasher() -> Self {
75        Poseidon(PhantomData)
76    }
77}
78
79#[inline]
80/// Converts a leaf position to the internal BigUint index for storage.
81pub(super) fn leaf_pos_to_index<const N: usize, F: FieldElement>(pos: &F) -> BigUint {
82    // Starting index for the last level
83    // 2^N - 1
84    let final_level_index = (BigUint::from(1u32) << (N as u64)) - 1u32;
85
86    final_level_index + pos.as_biguint()
87}
88
89/// Returns the log2 value of the given number. Used for converting the index to the level.
90#[inline]
91pub(super) fn log2(x: &BigUint) -> u64 {
92    (x + 1u32).bits() - 1
93}
94
95/// Returns the index of the left child, given an index.
96#[inline]
97pub(super) fn left_child(index: &BigUint) -> BigUint {
98    2u32 * index + 1u32
99}
100
101/// Returns the index of the right child, given an index.
102#[inline]
103pub(super) fn right_child(index: &BigUint) -> BigUint {
104    2u32 * index + 2u32
105}
106
107/// Returns true iff the given index represents a left child.
108#[inline]
109pub(super) fn is_left_child(index: &BigUint) -> bool {
110    // Any simple way to convert the (index % 2) into a u32 rather
111    // than converting 1 into a BigUint?
112    index % 2u32 == 1u32.into()
113}
114
115/// Returns the index of the parent, given an index.
116#[inline]
117pub(super) fn parent(index: &BigUint) -> Option<BigUint> {
118    if *index > 0u32.into() {
119        Some((index - 1u32) >> 1)
120    } else {
121        None
122    }
123}
124
125/// Returns the index of the sibling, given an index.
126#[inline]
127pub(super) fn sibling(index: &BigUint) -> Option<BigUint> {
128    if *index == 0u32.into() {
129        None
130    } else if is_left_child(index) {
131        Some(index + 1u32)
132    } else {
133        Some(index - 1u32)
134    }
135}