1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

//! Arbitrary value generation and test harnesses for high-level typed database access.

#![allow(unused_imports)]

use std::ops::Deref;

use zebra_chain::{amount::NonNegative, block::Block, sprout, value_balance::ValueBalance};

use crate::service::finalized_state::{
    disk_db::{DiskDb, DiskWriteBatch, WriteDisk},
    ZebraDb,
};

// Enable older test code to automatically access the inner database via Deref coercion.
impl Deref for ZebraDb {
    type Target = DiskDb;

    fn deref(&self) -> &Self::Target {
        self.db()
    }
}

impl ZebraDb {
    /// Returns the inner database.
    ///
    /// This is a test-only and shielded-scan-only method, because it allows write access
    /// and raw read access to the RocksDB instance.
    pub fn db(&self) -> &DiskDb {
        &self.db
    }

    /// Allow to set up a fake value pool in the database for testing purposes.
    #[cfg(any(test, feature = "proptest-impl"))]
    pub fn set_finalized_value_pool(&self, fake_value_pool: ValueBalance<NonNegative>) {
        let mut batch = DiskWriteBatch::new();
        let value_pool_cf = self.db().cf_handle("tip_chain_value_pool").unwrap();

        batch.zs_insert(&value_pool_cf, (), fake_value_pool);
        self.db().write(batch).unwrap();
    }

    /// Artificially prime the note commitment tree anchor sets with anchors
    /// referenced in a block, for testing purposes _only_.
    #[cfg(any(test, feature = "proptest-impl"))]
    pub fn populate_with_anchors(&self, block: &Block) {
        let mut batch = DiskWriteBatch::new();

        let sprout_anchors = self.db().cf_handle("sprout_anchors").unwrap();
        let sapling_anchors = self.db().cf_handle("sapling_anchors").unwrap();
        let orchard_anchors = self.db().cf_handle("orchard_anchors").unwrap();

        for transaction in block.transactions.iter() {
            // Sprout
            for joinsplit in transaction.sprout_groth16_joinsplits() {
                batch.zs_insert(
                    &sprout_anchors,
                    joinsplit.anchor,
                    sprout::tree::NoteCommitmentTree::default(),
                );
            }

            // Sapling
            for anchor in transaction.sapling_anchors() {
                batch.zs_insert(&sapling_anchors, anchor, ());
            }

            // Orchard
            if let Some(orchard_shielded_data) = transaction.orchard_shielded_data() {
                batch.zs_insert(&orchard_anchors, orchard_shielded_data.shared_anchor, ());
            }
        }

        self.db().write(batch).unwrap();
    }
}