Skip to content

Commit 111ab88

Browse files
committed
fixup! pyth: introduce pyth accumulator library
1 parent 1e3d1b8 commit 111ab88

File tree

3 files changed

+88
-32
lines changed

3 files changed

+88
-32
lines changed

Cargo.lock

Lines changed: 13 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyth/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ rand = "0.7.0"
2828
serde_json = "1.0.96"
2929
solana-client = { path = "../client" }
3030
solana-sdk = { path = "../sdk" }
31+
proptest = "1.1.0"
3132

3233
[package.metadata.docs.rs]
3334
targets = ["x86_64-unknown-linux-gnu"]

pyth/src/accumulators/merkle.rs

Lines changed: 74 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ impl<'a, H: Hasher + 'a> Accumulator<'a> for MerkleAccumulator<H> {
100100
type Proof = MerklePath<H>;
101101

102102
fn from_set(items: impl Iterator<Item = &'a [u8]>) -> Option<Self> {
103-
let items: Vec<H::Hash> = items.map(|i| hash_leaf::<H>(i)).collect();
103+
let items: Vec<&[u8]> = items.collect();
104104
Self::new(&items)
105105
}
106106

@@ -121,12 +121,12 @@ impl<'a, H: Hasher + 'a> Accumulator<'a> for MerkleAccumulator<H> {
121121

122122
// This code is adapted from the solana-merkle-tree crate to use a generic hasher.
123123
impl<H: Hasher> MerkleAccumulator<H> {
124-
pub fn new(items: &[H::Hash]) -> Option<Self> {
124+
pub fn new(items: &[&[u8]]) -> Option<Self> {
125125
if items.is_empty() {
126126
return None;
127127
}
128128

129-
let depth = (items.len() as f64).log2().ceil() as u32;
129+
let depth = items.len().next_power_of_two().trailing_zeros();
130130
let mut tree: Vec<H::Hash> = vec![Default::default(); 1 << (depth + 1)];
131131

132132
// Filling the leaf hashes
@@ -154,13 +154,11 @@ impl<H: Hasher> MerkleAccumulator<H> {
154154
})
155155
}
156156

157-
fn find_path(&self, index: usize) -> MerklePath<H> {
157+
fn find_path(&self, mut index: usize) -> MerklePath<H> {
158158
let mut path = Vec::new();
159-
let depth = (self.nodes.len() as f64).log2().ceil() as u32;
160-
let mut idx = (1 << depth) + index;
161-
while idx > 1 {
162-
path.push(self.nodes[idx ^ 1].clone());
163-
idx /= 2;
159+
while index > 1 {
160+
path.push(self.nodes[index ^ 1].clone());
161+
index /= 2;
164162
}
165163
MerklePath::new(path)
166164
}
@@ -170,6 +168,8 @@ impl<H: Hasher> MerkleAccumulator<H> {
170168
mod test {
171169
use {
172170
super::*,
171+
proptest::prelude::*,
172+
rand::RngCore,
173173
std::{
174174
collections::HashSet,
175175
mem::size_of,
@@ -203,6 +203,34 @@ mod test {
203203
}
204204
}
205205

206+
#[derive(Debug)]
207+
struct MerkleAccumulatorDataWrapper {
208+
pub accumulator: MerkleAccumulator,
209+
pub data: HashSet<Vec<u8>>,
210+
}
211+
212+
impl Arbitrary for MerkleAccumulatorDataWrapper {
213+
type Parameters = usize;
214+
215+
fn arbitrary_with(size: Self::Parameters) -> Self::Strategy {
216+
let size = size.saturating_add(1);
217+
prop::collection::vec(
218+
prop::collection::vec(any::<u8>(), 1..=10),
219+
size..=size.saturating_add(100),
220+
)
221+
.prop_map(|v| {
222+
let data: HashSet<Vec<u8>> = v.into_iter().collect();
223+
let accumulator =
224+
MerkleAccumulator::<Keccak256>::from_set(data.iter().map(|i| i.as_ref()))
225+
.unwrap();
226+
MerkleAccumulatorDataWrapper { accumulator, data }
227+
})
228+
.boxed()
229+
}
230+
231+
type Strategy = BoxedStrategy<Self>;
232+
}
233+
206234
#[test]
207235
fn test_merkle() {
208236
let mut set: HashSet<&[u8]> = HashSet::new();
@@ -231,27 +259,45 @@ mod test {
231259

232260
let accumulator = MerkleAccumulator::<Keccak256>::from_set(set.into_iter()).unwrap();
233261
let proof = accumulator.prove(&item_a).unwrap();
262+
234263
assert!(accumulator.check(proof, &item_a));
235264
let proof = accumulator.prove(&item_a).unwrap();
236-
println!(
237-
"proof: {:#?}",
238-
proof.0.iter().map(|x| format!("{x:?}")).collect::<Vec<_>>()
239-
);
240-
println!("accumulator root: {:?}", accumulator.get_root().unwrap());
241-
println!(
242-
r"
243-
Sizes:
244-
MerkleAccumulator::Proof {:?}
245-
Keccak256Hasher::Hash {:?}
246-
MerkleNode {:?}
247-
MerklePath {:?}
248-
249-
",
250-
size_of::<<MerkleAccumulator as Accumulator>::Proof>(),
251-
size_of::<<Keccak256 as Hasher>::Hash>(),
252-
size_of::<MerkleNode<Keccak256>>(),
253-
size_of::<MerklePath<Keccak256>>()
254-
);
265+
assert_eq!(size_of::<<Keccak256 as Hasher>::Hash>(), 32);
266+
255267
assert!(!accumulator.check(proof, &item_d));
256268
}
269+
270+
// We're using proptest to generate arbitrary Merkle trees as part of our fuzzing strategy.
271+
// This will help us identify any edge cases or unexpected behavior in the implementation.
272+
proptest! {
273+
#[test]
274+
fn test_merkle_tree(v in any::<MerkleAccumulatorDataWrapper>()) {
275+
let accumulator = v.accumulator;
276+
let data = v.data;
277+
for d in data.iter() {
278+
let proof = accumulator.prove(d).unwrap();
279+
let mut invalid_proof = proof.clone();
280+
invalid_proof.0.push(hash_leaf::<Keccak256>(&[0]));
281+
assert!(accumulator.check(proof, d));
282+
assert!(!accumulator.check(invalid_proof, d));
283+
284+
}
285+
286+
let false_data = data.clone()
287+
.into_iter()
288+
.map(|d| {
289+
let mut d = d.to_vec();
290+
d.extend_from_slice(&[0]);
291+
d
292+
})
293+
.filter(|d| !data.contains(d))
294+
.collect::<HashSet<Vec<u8>>>();
295+
296+
297+
for f_d in false_data.iter() {
298+
let proof = accumulator.prove(f_d);
299+
assert!(proof.is_none());
300+
}
301+
}
302+
}
257303
}

0 commit comments

Comments
 (0)