use crate::trie::{
bytes_to_nibbles,
proof_decode::{self, DecodedTrieProof, StorageValue},
proof_encode::ProofBuilder,
};
use alloc::vec::Vec;
use hashbrown::HashSet;
pub use proof_decode::ParseError;
#[derive(Debug, derive_more::Display, derive_more::Error)]
pub enum MinimizeProofError {
KeyNotFound,
IncompleteProof,
}
pub fn minimize_proof<T: AsRef<[u8]>>(
decoded_proof: &DecodedTrieProof<T>,
trie_root_merkle_value: &[u8; 32],
key: &[u8],
) -> Result<Vec<u8>, MinimizeProofError> {
let mut builder = ProofBuilder::new();
let mut key_nibbles_if_first_iter =
Some(bytes_to_nibbles(key.iter().copied()).collect::<Vec<_>>());
loop {
let Some(missing) = builder
.missing_node_values()
.next()
.map(|v| Vec::from(v))
.or_else(|| key_nibbles_if_first_iter.take())
else {
break;
};
if let Some(ancestor_key) = decoded_proof
.closest_ancestor_in_proof(trie_root_merkle_value, missing.iter().copied())
.map_err(|_| MinimizeProofError::KeyNotFound)?
{
let ancestor = decoded_proof
.proof_entry(trie_root_merkle_value, ancestor_key)
.unwrap();
builder.set_node_value(
&missing,
ancestor.node_value,
match ancestor.trie_node_info.storage_value {
StorageValue::Known { value, .. } => Some(value),
_ => None,
},
);
} else {
debug_assert!(itertools::equal(
missing.iter().copied(),
bytes_to_nibbles(key.iter().copied())
));
let root = decoded_proof
.trie_root_proof_entry(trie_root_merkle_value)
.ok_or(MinimizeProofError::KeyNotFound)?;
builder.set_node_value(&missing, root.node_value, None);
};
}
Ok(builder.build_to_vec())
}
pub fn merge_proofs<'a>(mut proofs: impl Iterator<Item = &'a [u8]>) -> Result<Vec<u8>, ParseError> {
let entries = proofs.try_fold(
HashSet::with_hasher(fnv::FnvBuildHasher::default()),
|mut acc, proof| {
let proof_entries = proof_decode::decode_proof(&proof)?;
acc.extend(proof_entries);
Ok(acc)
},
)?;
let mut ret = Vec::with_capacity(8 + entries.iter().map(|e| e.len() + 8).sum::<usize>());
ret.extend_from_slice(crate::util::encode_scale_compact_usize(entries.len()).as_ref());
for entry in entries {
ret.extend_from_slice(crate::util::encode_scale_compact_usize(entry.len()).as_ref());
ret.extend_from_slice(entry);
}
Ok(ret)
}