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
76
77
78
79
80
81
82
83
84
85

use std::collections::{HashMap, HashSet};

use solana_sdk::signature::Signature;

use crate::common::typedefs::account::Account;

use crate::common::typedefs::hash::Hash;

use crate::ingester::parser::indexer_events::PathNode;

#[derive(Debug, Clone)]
pub struct EnrichedPathNode {
    pub node: PathNode,
    pub slot: u64,
    pub tree: [u8; 32],
    pub seq: u64,
    pub level: usize,
    pub tree_depth: usize,
    pub leaf_index: Option<u32>,
}

pub struct PathUpdate {
    pub tree: [u8; 32],
    pub path: Vec<PathNode>,
    pub seq: u64,
}

#[derive(Hash, Eq, Clone, PartialEq, Debug)]
pub struct Transaction {
    pub signature: Signature,
    pub slot: u64,
}

#[derive(Hash, PartialEq, Eq)]
pub struct AccountTransaction {
    pub hash: Hash,
    pub signature: Signature,
    pub slot: u64,
}

#[derive(Default)]
/// Representation of state update of the compression system that is optimal for simple persistance.
pub struct StateUpdate {
    pub in_accounts: Vec<Account>,
    pub out_accounts: Vec<Account>,
    pub path_nodes: HashMap<([u8; 32], u32), EnrichedPathNode>,
    pub account_transactions: HashSet<AccountTransaction>,
}

impl StateUpdate {
    pub fn new() -> Self {
        StateUpdate::default()
    }

    pub fn prune_redundant_updates(&mut self) {
        let in_account_set: HashSet<Hash> =
            self.in_accounts.iter().map(|a| a.hash.clone()).collect();
        // NOTE: For snapshot verification, we might need to persist accounts data until accounts are
        //       removed from the tree through the nullifier crank.
        self.out_accounts
            .retain(|a| !in_account_set.contains(&a.hash));
    }

    pub fn merge_updates(updates: Vec<StateUpdate>) -> StateUpdate {
        let mut merged = StateUpdate::default();
        for update in updates {
            merged.in_accounts.extend(update.in_accounts);
            merged.out_accounts.extend(update.out_accounts);
            merged
                .account_transactions
                .extend(update.account_transactions);

            for (key, node) in update.path_nodes {
                if let Some(existing) = merged.path_nodes.get_mut(&key) {
                    if (*existing).seq < node.seq {
                        *existing = node;
                    }
                } else {
                    merged.path_nodes.insert(key, node);
                }
            }
        }
        merged
    }
}