use alloc::collections::BTreeSet;
#[cfg(feature = "with-serde")]
use serde::Deserialize;
#[cfg(feature = "with-serde")]
use serde::Serialize;
use super::node_hash::AccumulatorHash;
use super::node_hash::BitcoinNodeHash;
use super::proof::NodesAndRootsOldNew;
use super::proof::Proof;
use super::util;
use crate::prelude::*;
#[derive(Debug, Clone, Default)]
pub struct UpdateData<Hash: AccumulatorHash> {
pub(crate) to_destroy: Vec<u64>,
pub(crate) prev_num_leaves: u64,
pub(crate) new_add: Vec<(u64, Hash)>,
pub(crate) new_del: Vec<(u64, Hash)>,
}
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum StumpError {
EmptyProof,
Io(io::ErrorKind),
InvalidProof(String),
}
impl From<io::Error> for StumpError {
fn from(err: io::Error) -> Self {
Self::Io(err.kind())
}
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "with-serde", derive(Serialize, Deserialize))]
pub struct Stump<Hash: AccumulatorHash = BitcoinNodeHash> {
pub leaves: u64,
pub roots: Vec<Hash>,
}
impl Default for Stump {
fn default() -> Self {
Self::new()
}
}
impl Stump {
pub fn new() -> Self {
Self {
leaves: 0,
roots: Vec::new(),
}
}
pub fn serialize<Dst: Write>(&self, mut writer: &mut Dst) -> Result<u64, StumpError> {
let mut len = 8;
writer.write_all(&self.leaves.to_le_bytes())?;
writer.write_all(&(self.roots.len() as u64).to_le_bytes())?;
for root in self.roots.iter() {
len += 32;
root.write(&mut writer)?;
}
Ok(len)
}
pub fn undo(&mut self, old_state: Self) {
self.leaves = old_state.leaves;
self.roots = old_state.roots;
}
}
impl<Hash: AccumulatorHash> Stump<Hash> {
pub fn verify(&self, proof: &Proof<Hash>, del_hashes: &[Hash]) -> Result<bool, StumpError> {
proof
.verify(del_hashes, &self.roots, self.leaves)
.map_err(StumpError::InvalidProof)
}
pub fn new_with_hash() -> Self {
Self {
leaves: 0,
roots: Vec::new(),
}
}
pub fn modify(
&self,
utxos: &[Hash],
del_hashes: &[Hash],
proof: &Proof<Hash>,
) -> Result<(Self, UpdateData<Hash>), StumpError> {
let (intermediate, mut computed_roots) = self.remove(del_hashes, proof)?;
let mut new_roots = vec![];
for root in self.roots.iter() {
if let Some(pos) = computed_roots.iter().position(|(old, _new)| old == root) {
let (old_root, new_root) = computed_roots.remove(pos);
if old_root == *root {
new_roots.push(new_root);
continue;
}
}
new_roots.push(*root);
}
if !computed_roots.is_empty() {
return Err(StumpError::EmptyProof);
}
let (roots, updated, destroyed) = Self::add(new_roots, utxos, self.leaves);
let new_stump = Self {
leaves: self.leaves + utxos.len() as u64,
roots,
};
let update_data = UpdateData {
new_add: updated,
prev_num_leaves: self.leaves,
to_destroy: destroyed,
new_del: intermediate,
};
Ok((new_stump, update_data))
}
pub fn deserialize<Source: Read>(mut data: Source) -> Result<Self, StumpError> {
let leaves = util::read_u64(&mut data)?;
let roots_len = util::read_u64(&mut data)?;
let mut roots = vec![];
for _ in 0..roots_len {
let root = Hash::read(&mut data)?;
roots.push(root);
}
Ok(Self { leaves, roots })
}
fn remove(
&self,
del_hashes: &[Hash],
proof: &Proof<Hash>,
) -> Result<NodesAndRootsOldNew<Hash>, StumpError> {
if del_hashes.is_empty() {
return Ok((
vec![],
self.roots.iter().map(|root| (*root, *root)).collect(),
));
}
let del_hashes = del_hashes
.iter()
.map(|hash| (*hash, Hash::empty()))
.collect::<Vec<_>>();
proof
.calculate_hashes_delete(&del_hashes, self.leaves)
.map_err(StumpError::InvalidProof)
}
fn add(
mut roots: Vec<Hash>,
utxos: &[Hash],
mut leaves: u64,
) -> (Vec<Hash>, Vec<(u64, Hash)>, Vec<u64>) {
let after_rows = util::tree_rows(leaves + (utxos.len() as u64));
let mut updated_subtree: BTreeSet<(u64, Hash)> = BTreeSet::new();
let all_deleted = util::roots_to_destroy(utxos.len() as u64, leaves, &roots);
for (i, add) in utxos.iter().enumerate() {
let mut pos = leaves;
let deleted = util::roots_to_destroy((utxos.len() - i) as u64, leaves, &roots);
for del in deleted {
if util::is_ancestor(util::parent(del, after_rows), pos, after_rows).unwrap() {
pos = util::calc_next_pos(pos, del, after_rows).unwrap();
}
}
let mut h = 0;
let mut to_add = *add;
while (leaves >> h) & 1 == 1 {
let root = roots.pop();
if let Some(root) = root {
if !root.is_empty() {
updated_subtree.insert((util::left_sibling(pos), root));
updated_subtree.insert((pos, to_add));
pos = util::parent(pos, after_rows);
to_add = AccumulatorHash::parent_hash(&root, &to_add);
}
}
h += 1;
}
updated_subtree.insert((pos, to_add));
roots.push(to_add);
leaves += 1;
}
(roots, updated_subtree.into_iter().collect(), all_deleted)
}
}
#[cfg(test)]
mod test {
use core::fmt;
use core::fmt::Display;
use core::str::FromStr;
use io::Cursor;
use serde::Deserialize;
use super::Stump;
use super::*;
use crate::node_hash::AccumulatorHash;
use crate::node_hash::BitcoinNodeHash;
use crate::proof::Proof;
use crate::util::hash_from_u8;
#[derive(Debug, Deserialize)]
struct TestCase {
leaf_preimages: Vec<u8>,
target_values: Option<Vec<u64>>,
expected_roots: Vec<String>,
proofhashes: Option<Vec<String>>,
}
#[test]
fn test_stump() {
let s = Stump::new();
assert!(s.leaves == 0);
assert!(s.roots.is_empty());
}
#[test]
fn test_custom_hash_type() {
#[derive(Debug, Default, Clone, Copy, Hash, PartialEq, Eq, PartialOrd, Ord)]
struct CustomHash([u8; 32]);
impl Display for CustomHash {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{:?}", self.0)
}
}
impl AccumulatorHash for CustomHash {
fn empty() -> Self {
Self([0; 32])
}
fn is_empty(&self) -> bool {
self.0.iter().all(|&x| x == 0)
}
fn parent_hash(left: &Self, right: &Self) -> Self {
let mut hash = [0; 32];
#[allow(clippy::needless_range_loop)]
for i in 0..32 {
hash[i] = left.0[i] ^ right.0[i];
}
Self(hash)
}
fn read<R: Read>(reader: &mut R) -> io::Result<Self> {
let mut hash = [0; 32];
reader
.read_exact(&mut hash)
.map_err(|e| e.to_string())
.unwrap();
Ok(Self(hash))
}
fn write<W: Write>(&self, writer: &mut W) -> io::Result<()> {
writer
.write_all(&self.0)
.map_err(|e| e.to_string())
.unwrap();
Ok(())
}
fn is_placeholder(&self) -> bool {
false
}
fn placeholder() -> Self {
Self([0; 32])
}
}
let s = Stump::<CustomHash>::new_with_hash();
assert!(s.leaves == 0);
assert!(s.roots.is_empty());
let hashes = [0, 1, 2, 3, 4, 5, 6, 7]
.iter()
.map(|&el| CustomHash([el; 32]))
.collect::<Vec<_>>();
let (stump, _) = s
.modify(
&hashes,
&[],
&Proof::<CustomHash>::new_with_hash(Vec::new(), Vec::new()),
)
.unwrap();
assert_eq!(stump.leaves, 8);
assert_eq!(stump.roots.len(), 1);
}
#[test]
fn test_updated_data() {
#[derive(Debug, Deserialize)]
struct TestData {
roots: Vec<String>,
leaves: u64,
additional_preimages: Vec<u64>,
del_hashes: Vec<String>,
proof_hashes: Vec<String>,
proof_targets: Vec<u64>,
new_add_pos: Vec<u64>,
new_add_hash: Vec<String>,
new_del_pos: Vec<u64>,
new_del_hashes: Vec<String>,
to_destroy: Vec<u64>,
}
let contents = include_str!("../../test_values/update_data_tests.json");
let tests =
serde_json::from_str::<Vec<TestData>>(contents).expect("JSON deserialization error");
for data in tests {
let roots = data
.roots
.iter()
.map(|hash| BitcoinNodeHash::from_str(hash).unwrap())
.collect();
let stump = Stump {
leaves: data.leaves,
roots,
};
let utxos = data
.additional_preimages
.iter()
.map(|preimage| hash_from_u8(*preimage as u8))
.collect::<Vec<_>>();
let del_hashes = data
.del_hashes
.iter()
.map(|hash| BitcoinNodeHash::from_str(hash).unwrap())
.collect::<Vec<_>>();
let proof_hashes = data
.proof_hashes
.iter()
.map(|hash| BitcoinNodeHash::from_str(hash).unwrap())
.collect::<Vec<_>>();
let proof = Proof::new(data.proof_targets, proof_hashes);
let (_, updated) = stump.modify(&utxos, &del_hashes, &proof).unwrap();
let new_add_hash: Vec<_> = data
.new_add_hash
.iter()
.map(|hash| BitcoinNodeHash::from_str(hash).unwrap())
.collect();
let new_add: Vec<_> = data.new_add_pos.into_iter().zip(new_add_hash).collect();
let new_del_hash: Vec<_> = data
.new_del_hashes
.iter()
.map(|hash| BitcoinNodeHash::from_str(hash).unwrap())
.collect();
let new_del: Vec<_> = data.new_del_pos.into_iter().zip(new_del_hash).collect();
assert_eq!(updated.prev_num_leaves, data.leaves);
assert_eq!(updated.to_destroy, data.to_destroy);
assert_eq!(updated.new_add, new_add);
for del in new_del.iter() {
assert!(updated.new_del.contains(del));
}
}
}
#[test]
fn test_update_data_add() {
let preimages = vec![0, 1, 2, 3];
let hashes = preimages.into_iter().map(hash_from_u8).collect::<Vec<_>>();
let (_, updated) = Stump::new()
.modify(&hashes, &[], &Proof::default())
.unwrap();
let positions = vec![0, 1, 2, 3, 4, 5, 6];
let hashes: Vec<_> = [
"6e340b9cffb37a989ca544e6bb780a2c78901d3fb33738768511a30617afa01d",
"4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459a",
"dbc1b4c900ffe48d575b5da5c638040125f65db0fe3e24494b76ea986457d986",
"084fed08b978af4d7d196a7446a86b58009e636b611db16211b65a9aadff29c5",
"02242b37d8e851f1e86f46790298c7097df06893d6226b7c1453c213e91717de",
"9576f4ade6e9bc3a6458b506ce3e4e890df29cb14cb5d3d887672aef55647a2b",
"df46b17be5f66f0750a4b3efa26d4679db170a72d41eb56c3e4ff75a58c65386",
]
.iter()
.map(|hash| BitcoinNodeHash::from_str(hash).unwrap())
.collect();
let positions: Vec<_> = positions.into_iter().zip(hashes).collect();
assert_eq!(positions, updated.new_add);
}
#[test]
#[cfg(feature = "with-serde")]
fn test_serde_rtt() {
let stump = Stump::new()
.modify(&[hash_from_u8(0), hash_from_u8(1)], &[], &Proof::default())
.unwrap()
.0;
let serialized = serde_json::to_string(&stump).expect("Serialization failed");
let deserialized: Stump =
serde_json::from_str(&serialized).expect("Deserialization failed");
assert_eq!(stump, deserialized);
}
fn run_case_with_deletion(case: TestCase) {
let leaf_hashes = case
.leaf_preimages
.into_iter()
.map(hash_from_u8)
.collect::<Vec<_>>();
let target_hashes = case
.target_values
.as_ref()
.unwrap()
.iter()
.map(|target| hash_from_u8(*target as u8))
.collect::<Vec<_>>();
let proof_hashes = case
.proofhashes
.unwrap_or_default()
.into_iter()
.map(|hash| {
BitcoinNodeHash::from_str(hash.as_str()).expect("Test case hashes are valid")
})
.collect::<Vec<_>>();
let proof = Proof::new(case.target_values.unwrap(), proof_hashes);
let roots = case
.expected_roots
.into_iter()
.map(|hash| {
BitcoinNodeHash::from_str(hash.as_str()).expect("Test case hashes are valid")
})
.collect::<Vec<BitcoinNodeHash>>();
let (stump, _) = Stump::new()
.modify(&leaf_hashes, &[], &Proof::default())
.expect("This stump is valid");
let (stump, _) = stump.modify(&[], &target_hashes, &proof).unwrap();
assert_eq!(stump.roots, roots);
}
fn run_single_addition_case(case: TestCase) {
let s = Stump::new();
let test_values = case.leaf_preimages;
let roots = case.expected_roots;
let hashes = test_values
.iter()
.map(|value| hash_from_u8(*value))
.collect::<Vec<_>>();
let (s, _) = s
.modify(&hashes, &[], &Proof::default())
.expect("Stump from test cases are valid");
assert_eq!(s.leaves, hashes.len() as u64);
for (i, root) in roots.into_iter().enumerate() {
assert_eq!(root, s.roots[i].to_string());
}
}
#[test]
fn test_undo() {
let mut hashes = vec![];
for i in 0..100_u8 {
hashes.push(hash_from_u8(i));
}
let s_old = Stump::new();
let s_old = s_old
.modify(&hashes, &[], &Proof::default())
.expect("Stump from test cases are valid")
.0;
let mut s_new = s_old.clone();
for i in 0..100_u8 {
hashes.push(hash_from_u8(i));
}
s_new
.modify(&hashes, &[], &Proof::default())
.expect("Stump from test cases are valid");
let s_old_copy = s_old.clone();
s_new.undo(s_old);
assert!(s_new == s_old_copy);
}
#[test]
fn test_serialize() {
let hashes = [0, 1, 2, 3, 4, 5, 6, 7]
.iter()
.map(|&el| BitcoinNodeHash::from([el; 32]))
.collect::<Vec<_>>();
let (stump, _) = Stump::new()
.modify(&hashes, &[], &Proof::default())
.unwrap();
let mut writer = Vec::new();
stump.serialize(&mut writer).unwrap();
let mut reader = Cursor::new(writer);
let stump2 = Stump::deserialize(&mut reader).unwrap();
assert_eq!(stump, stump2);
}
#[test]
fn run_test_cases() {
#[derive(Deserialize)]
struct TestsJSON {
insertion_tests: Vec<TestCase>,
deletion_tests: Vec<TestCase>,
}
let contents = include_str!("../../test_values/test_cases.json");
let tests =
serde_json::from_str::<TestsJSON>(contents).expect("JSON deserialization error");
for i in tests.insertion_tests {
run_single_addition_case(i);
}
for i in tests.deletion_tests {
run_case_with_deletion(i);
}
}
}
#[cfg(bench)]
mod bench {
extern crate test;
use std::convert::TryFrom;
use test::Bencher;
use super::Stump;
use crate::accumulator::node_hash::AccumulatorHash;
use crate::accumulator::proof::Proof;
use crate::accumulator::util::hash_from_u8;
#[bench]
fn bench_addition(bencher: &mut Bencher) {
let hash = [
"4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459a",
"4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459b",
"4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459c",
"4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459d",
"4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459e",
"4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459f",
]
.iter()
.map(|&hash| AccumulatorHash::try_from(hash).unwrap())
.collect::<Vec<_>>();
let proof = &Proof::default();
bencher.iter(move || {
let _ = Stump::new().modify(&hash, &[], &proof);
});
}
#[bench]
fn bench_add_del(bencher: &mut Bencher) {
let leaf_preimages = [0_u8, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
let leaves = leaf_preimages
.iter()
.map(|preimage| hash_from_u8(*preimage))
.collect::<Vec<_>>();
let target_values = [0, 4, 5, 6, 7, 8];
let target_hashes = target_values
.iter()
.map(|val| hash_from_u8(*val as u8))
.collect::<Vec<_>>();
let proofhashes = [
"4bf5122f344554c53bde2ebb8cd2b7e3d1600ad631c385a5d7cce23c7785459a",
"2b4c342f5433ebe591a1da77e013d1b72475562d48578dca8b84bac6651c3cb9",
"9576f4ade6e9bc3a6458b506ce3e4e890df29cb14cb5d3d887672aef55647a2b",
"c413035120e8c9b0ca3e40c93d06fe60a0d056866138300bb1f1dd172b4923c3",
]
.iter()
.map(|&value| AccumulatorHash::try_from(value).unwrap())
.collect::<Vec<_>>();
let acc = Stump::new()
.modify(&leaves, &vec![], &Proof::default())
.unwrap()
.0;
let proof = Proof::new(target_values.to_vec(), proofhashes);
bencher.iter(move || acc.modify(&leaves, &target_hashes, &proof));
}
}