Merge pull request #1248 from zcash/zcs-test-state-orchard-tree

zcash_client_sqlite: Track Orchard commitment tree sizes in `TestState`
This commit is contained in:
str4d 2024-03-09 19:40:39 +00:00 committed by GitHub
commit 58f46e4636
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 174 additions and 59 deletions

View File

@ -414,6 +414,7 @@ mod tests {
AddressType::DefaultExternal,
NonNegativeAmount::const_from_u64(8),
2,
2,
);
st.generate_next_block(
&dfvk,

View File

@ -166,10 +166,54 @@ impl<Cache> TestBuilder<Cache> {
}
}
pub(crate) struct CachedBlock {
height: BlockHeight,
hash: BlockHash,
sapling_end_size: u32,
orchard_end_size: u32,
}
impl CachedBlock {
fn none(sapling_activation_height: BlockHeight) -> Self {
Self {
height: sapling_activation_height,
hash: BlockHash([0; 32]),
sapling_end_size: 0,
orchard_end_size: 0,
}
}
fn at(
height: BlockHeight,
hash: BlockHash,
sapling_tree_size: u32,
orchard_tree_size: u32,
) -> Self {
Self {
height,
hash,
sapling_end_size: sapling_tree_size,
orchard_end_size: orchard_tree_size,
}
}
fn roll_forward(self, cb: &CompactBlock) -> Self {
assert_eq!(self.height + 1, cb.height());
Self {
height: cb.height(),
hash: cb.hash(),
sapling_end_size: self.sapling_end_size
+ cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
orchard_end_size: self.orchard_end_size
+ cb.vtx.iter().map(|tx| tx.actions.len() as u32).sum::<u32>(),
}
}
}
/// The state for a `zcash_client_sqlite` test.
pub(crate) struct TestState<Cache> {
cache: Cache,
latest_cached_block: Option<(BlockHeight, BlockHash, u32)>,
latest_cached_block: Option<CachedBlock>,
_data_file: NamedTempFile,
db_data: WalletDb<Connection, Network>,
test_account: Option<(
@ -190,8 +234,8 @@ where
self.cache.block_source()
}
pub(crate) fn latest_cached_block(&self) -> &Option<(BlockHeight, BlockHash, u32)> {
&self.latest_cached_block
pub(crate) fn latest_cached_block(&self) -> Option<&CachedBlock> {
self.latest_cached_block.as_ref()
}
/// Creates a fake block at the expected next height containing a single output of the
@ -202,19 +246,22 @@ where
req: AddressType,
value: NonNegativeAmount,
) -> (BlockHeight, Cache::InsertResult, Fvk::Nullifier) {
let (height, prev_hash, initial_sapling_tree_size) = self
let cached_block = self
.latest_cached_block
.map(|(prev_height, prev_hash, end_size)| (prev_height + 1, prev_hash, end_size))
.unwrap_or_else(|| (self.sapling_activation_height(), BlockHash([0; 32]), 0));
.take()
.unwrap_or_else(|| CachedBlock::none(self.sapling_activation_height() - 1));
let height = cached_block.height + 1;
let (res, nf) = self.generate_block_at(
height,
prev_hash,
cached_block.hash,
fvk,
req,
value,
initial_sapling_tree_size,
cached_block.sapling_end_size,
cached_block.orchard_end_size,
);
assert!(self.latest_cached_block.is_some());
(height, res, nf)
}
@ -224,6 +271,7 @@ where
///
/// This generated block will be treated as the latest block, and subsequent calls to
/// [`Self::generate_next_block`] will build on it.
#[allow(clippy::too_many_arguments)]
pub(crate) fn generate_block_at<Fvk: TestFvk>(
&mut self,
height: BlockHeight,
@ -232,6 +280,7 @@ where
req: AddressType,
value: NonNegativeAmount,
initial_sapling_tree_size: u32,
initial_orchard_tree_size: u32,
) -> (Cache::InsertResult, Fvk::Nullifier) {
let (cb, nf) = fake_compact_block(
&self.network(),
@ -241,15 +290,19 @@ where
req,
value,
initial_sapling_tree_size,
initial_orchard_tree_size,
);
let res = self.cache.insert(&cb);
self.latest_cached_block = Some((
height,
cb.hash(),
initial_sapling_tree_size
+ cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
));
self.latest_cached_block = Some(
CachedBlock::at(
height - 1,
cb.hash(),
initial_sapling_tree_size,
initial_orchard_tree_size,
)
.roll_forward(&cb),
);
(res, nf)
}
@ -263,29 +316,26 @@ where
to: impl Into<Address>,
value: NonNegativeAmount,
) -> (BlockHeight, Cache::InsertResult) {
let (height, prev_hash, initial_sapling_tree_size) = self
let cached_block = self
.latest_cached_block
.map(|(prev_height, prev_hash, end_size)| (prev_height + 1, prev_hash, end_size))
.unwrap_or_else(|| (self.sapling_activation_height(), BlockHash([0; 32]), 0));
.take()
.unwrap_or_else(|| CachedBlock::none(self.sapling_activation_height() - 1));
let height = cached_block.height + 1;
let cb = fake_compact_block_spending(
&self.network(),
height,
prev_hash,
cached_block.hash,
note,
fvk,
to.into(),
value,
initial_sapling_tree_size,
cached_block.sapling_end_size,
cached_block.orchard_end_size,
);
let res = self.cache.insert(&cb);
self.latest_cached_block = Some((
height,
cb.hash(),
initial_sapling_tree_size
+ cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
));
self.latest_cached_block = Some(cached_block.roll_forward(&cb));
(height, res)
}
@ -320,27 +370,23 @@ where
tx_index: usize,
tx: &Transaction,
) -> (BlockHeight, Cache::InsertResult) {
let (height, prev_hash, initial_sapling_tree_size) = self
let cached_block = self
.latest_cached_block
.map(|(prev_height, prev_hash, end_size)| (prev_height + 1, prev_hash, end_size))
.unwrap_or_else(|| (self.sapling_activation_height(), BlockHash([0; 32]), 0));
.take()
.unwrap_or_else(|| CachedBlock::none(self.sapling_activation_height() - 1));
let height = cached_block.height + 1;
let cb = fake_compact_block_from_tx(
height,
prev_hash,
cached_block.hash,
tx_index,
tx,
initial_sapling_tree_size,
0,
cached_block.sapling_end_size,
cached_block.orchard_end_size,
);
let res = self.cache.insert(&cb);
self.latest_cached_block = Some((
height,
cb.hash(),
initial_sapling_tree_size
+ cb.vtx.iter().map(|tx| tx.outputs.len() as u32).sum::<u32>(),
));
self.latest_cached_block = Some(cached_block.roll_forward(&cb));
(height, res)
}
@ -399,10 +445,12 @@ where
self.cache
.block_source()
.with_blocks::<_, Infallible>(None, None, |block: CompactBlock| {
self.latest_cached_block = Some((
let chain_metadata = block.chain_metadata.unwrap();
self.latest_cached_block = Some(CachedBlock::at(
BlockHeight::from_u32(block.height.try_into().unwrap()),
BlockHash::from_slice(block.hash.as_slice()),
block.chain_metadata.unwrap().sapling_commitment_tree_size,
chain_metadata.sapling_commitment_tree_size,
chain_metadata.orchard_commitment_tree_size,
));
Ok(())
})
@ -1070,6 +1118,7 @@ fn fake_compact_tx<R: RngCore + CryptoRng>(rng: &mut R) -> CompactTx {
/// Create a fake CompactBlock at the given height, containing a single output paying
/// an address. Returns the CompactBlock and the nullifier for the new note.
#[allow(clippy::too_many_arguments)]
fn fake_compact_block<P: consensus::Parameters, Fvk: TestFvk>(
params: &P,
height: BlockHeight,
@ -1078,6 +1127,7 @@ fn fake_compact_block<P: consensus::Parameters, Fvk: TestFvk>(
req: AddressType,
value: NonNegativeAmount,
initial_sapling_tree_size: u32,
initial_orchard_tree_size: u32,
) -> (CompactBlock, Fvk::Nullifier) {
// Create a fake Note for the account
let mut rng = OsRng;
@ -1094,8 +1144,13 @@ fn fake_compact_block<P: consensus::Parameters, Fvk: TestFvk>(
&mut rng,
);
let cb =
fake_compact_block_from_compact_tx(ctx, height, prev_hash, initial_sapling_tree_size, 0);
let cb = fake_compact_block_from_compact_tx(
ctx,
height,
prev_hash,
initial_sapling_tree_size,
initial_orchard_tree_size,
);
(cb, nf)
}
@ -1152,6 +1207,7 @@ fn fake_compact_block_spending<P: consensus::Parameters, Fvk: TestFvk>(
to: Address,
value: NonNegativeAmount,
initial_sapling_tree_size: u32,
initial_orchard_tree_size: u32,
) -> CompactBlock {
let mut rng = OsRng;
let mut ctx = fake_compact_tx(&mut rng);
@ -1229,7 +1285,13 @@ fn fake_compact_block_spending<P: consensus::Parameters, Fvk: TestFvk>(
}
}
fake_compact_block_from_compact_tx(ctx, height, prev_hash, initial_sapling_tree_size, 0)
fake_compact_block_from_compact_tx(
ctx,
height,
prev_hash,
initial_sapling_tree_size,
initial_orchard_tree_size,
)
}
fn fake_compact_block_from_compact_tx(

View File

@ -1248,10 +1248,23 @@ pub(crate) fn birthday_in_anchor_shard<T: ShieldedPoolTester>() {
let received_tx_height = birthday.height() + 10;
let initial_sapling_tree_size =
u64::from(birthday.sapling_frontier().value().unwrap().position() + 1)
.try_into()
.unwrap();
let initial_sapling_tree_size = birthday
.sapling_frontier()
.value()
.map(|f| u64::from(f.position() + 1))
.unwrap_or(0)
.try_into()
.unwrap();
#[cfg(feature = "orchard")]
let initial_orchard_tree_size = birthday
.orchard_frontier()
.value()
.map(|f| u64::from(f.position() + 1))
.unwrap_or(0)
.try_into()
.unwrap();
#[cfg(not(feature = "orchard"))]
let initial_orchard_tree_size = 0;
// Generate 9 blocks that have no value for us, starting at the birthday height.
let not_our_key = T::sk_to_fvk(&T::sk(&[]));
@ -1263,6 +1276,7 @@ pub(crate) fn birthday_in_anchor_shard<T: ShieldedPoolTester>() {
AddressType::DefaultExternal,
not_our_value,
initial_sapling_tree_size,
initial_orchard_tree_size,
);
for _ in 1..9 {
st.generate_next_block(&not_our_key, AddressType::DefaultExternal, not_our_value);
@ -1340,7 +1354,8 @@ pub(crate) fn checkpoint_gaps<T: ShieldedPoolTester>() {
&not_our_key,
AddressType::DefaultExternal,
not_our_value,
st.latest_cached_block().unwrap().2,
st.latest_cached_block().unwrap().sapling_end_size,
st.latest_cached_block().unwrap().orchard_end_size,
);
// Scan the block

View File

@ -2645,6 +2645,7 @@ mod tests {
AddressType::DefaultExternal,
not_our_value,
17,
17,
);
st.scan_cached_blocks(end_height, 1);
@ -2661,6 +2662,7 @@ mod tests {
AddressType::DefaultExternal,
not_our_value,
0,
0,
);
st.scan_cached_blocks(start_height, 1);

View File

@ -608,7 +608,9 @@ pub(crate) mod tests {
// We'll start inserting leaf notes 5 notes after the end of the third subtree, with a gap
// of 10 blocks. After `scan_cached_blocks`, the scan queue should have a requested scan
// range of 300..310 with `FoundNote` priority, 310..320 with `Scanned` priority.
// We set both Sapling and Orchard to the same initial tree size for simplicity.
let initial_sapling_tree_size = (0x1 << 16) * 3 + 5;
let initial_orchard_tree_size = (0x1 << 16) * 3 + 5;
let initial_height = sapling_activation_height + 310;
let value = NonNegativeAmount::const_from_u64(50000);
@ -619,6 +621,7 @@ pub(crate) mod tests {
AddressType::DefaultExternal,
value,
initial_sapling_tree_size,
initial_orchard_tree_size,
);
for _ in 1..=10 {
@ -702,18 +705,18 @@ pub(crate) mod tests {
.with_block_cache()
.with_test_account(|network| {
// We use Canopy activation as an arbitrary birthday height that's greater than Sapling
// activation. We set the Canopy frontier to be 1234 notes into the second shard.
// activation. We set the Canopy Sapling frontier to be 1234 notes into the second shard.
let birthday_height = network.activation_height(NetworkUpgrade::Canopy).unwrap();
let frontier_position = Position::from((0x1 << 16) + 1234);
let frontier = Frontier::from_parts(
frontier_position,
let sapling_frontier_position = Position::from((0x1 << 16) + 1234);
let sapling_frontier = Frontier::from_parts(
sapling_frontier_position,
Node::empty_leaf(),
vec![Node::empty_leaf(); frontier_position.past_ommer_count().into()],
vec![Node::empty_leaf(); sapling_frontier_position.past_ommer_count().into()],
)
.unwrap();
AccountBirthday::from_parts(
birthday_height,
frontier,
sapling_frontier,
#[cfg(feature = "orchard")]
Frontier::empty(),
None,
@ -921,6 +924,23 @@ pub(crate) mod tests {
assert_eq!(actual, expected);
// Now, scan the max scanned block.
let initial_sapling_tree_size = birthday
.sapling_frontier()
.value()
.map(|f| u64::from(f.position() + 1))
.unwrap_or(0)
.try_into()
.unwrap();
#[cfg(feature = "orchard")]
let initial_orchard_tree_size = birthday
.orchard_frontier()
.value()
.map(|f| u64::from(f.position() + 1))
.unwrap_or(0)
.try_into()
.unwrap();
#[cfg(not(feature = "orchard"))]
let initial_orchard_tree_size = 0;
st.generate_block_at(
max_scanned,
BlockHash([0u8; 32]),
@ -928,9 +948,8 @@ pub(crate) mod tests {
AddressType::DefaultExternal,
// 1235 notes into into the second shard
NonNegativeAmount::const_from_u64(10000),
u64::from(birthday.sapling_frontier().value().unwrap().position() + 1)
.try_into()
.unwrap(),
initial_sapling_tree_size,
initial_orchard_tree_size,
);
st.scan_cached_blocks(max_scanned, 1);
@ -1041,15 +1060,31 @@ pub(crate) mod tests {
assert_eq!(actual, expected);
// Now, scan the max scanned block.
let initial_sapling_tree_size = birthday
.sapling_frontier()
.value()
.map(|f| u64::from(f.position() + 1))
.unwrap_or(0)
.try_into()
.unwrap();
#[cfg(feature = "orchard")]
let initial_orchard_tree_size = birthday
.orchard_frontier()
.value()
.map(|f| u64::from(f.position() + 1))
.unwrap_or(0)
.try_into()
.unwrap();
#[cfg(not(feature = "orchard"))]
let initial_orchard_tree_size = 0;
st.generate_block_at(
max_scanned,
BlockHash([0u8; 32]),
&dfvk,
AddressType::DefaultExternal,
NonNegativeAmount::const_from_u64(10000),
u64::from(birthday.sapling_frontier().value().unwrap().position() + 1)
.try_into()
.unwrap(),
initial_sapling_tree_size,
initial_orchard_tree_size,
);
st.scan_cached_blocks(max_scanned, 1);