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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use blake2::{
  digest::{consts::U32, Digest},
  Blake2b,
};

use scale::Encode;
use borsh::{BorshSerialize, BorshDeserialize};
use serai_client::{
  in_instructions::primitives::{Batch, SignedBatch},
  primitives::ExternalNetworkId,
  validator_sets::primitives::{ExternalValidatorSet, Session},
};

pub use serai_db::*;

use ::tributary::ReadWrite;
use crate::tributary::{TributarySpec, Transaction, scanner::RecognizedIdType};

create_db!(
  MainDb {
    HandledMessageDb: (network: ExternalNetworkId) -> u64,
    ActiveTributaryDb: () -> Vec<u8>,
    RetiredTributaryDb: (set: ExternalValidatorSet) -> (),
    FirstPreprocessDb: (
      network: ExternalNetworkId,
      id_type: RecognizedIdType,
      id: &[u8]
    ) -> Vec<Vec<u8>>,
    LastReceivedBatchDb: (network: ExternalNetworkId) -> u32,
    ExpectedBatchDb: (network: ExternalNetworkId, id: u32) -> [u8; 32],
    BatchDb: (network: ExternalNetworkId, id: u32)  -> SignedBatch,
    LastVerifiedBatchDb: (network: ExternalNetworkId) -> u32,
    HandoverBatchDb: (set: ExternalValidatorSet) -> u32,
    LookupHandoverBatchDb: (network: ExternalNetworkId, batch: u32) -> Session,
    QueuedBatchesDb: (set: ExternalValidatorSet) -> Vec<u8>
  }
);

impl ActiveTributaryDb {
  pub fn active_tributaries<G: Get>(getter: &G) -> (Vec<u8>, Vec<TributarySpec>) {
    let bytes = Self::get(getter).unwrap_or_default();
    let mut bytes_ref: &[u8] = bytes.as_ref();

    let mut tributaries = vec![];
    while !bytes_ref.is_empty() {
      tributaries.push(TributarySpec::deserialize_reader(&mut bytes_ref).unwrap());
    }

    (bytes, tributaries)
  }

  pub fn add_participating_in_tributary(txn: &mut impl DbTxn, spec: &TributarySpec) {
    let (mut existing_bytes, existing) = ActiveTributaryDb::active_tributaries(txn);
    for tributary in &existing {
      if tributary == spec {
        return;
      }
    }

    spec.serialize(&mut existing_bytes).unwrap();
    ActiveTributaryDb::set(txn, &existing_bytes);
  }

  pub fn retire_tributary(txn: &mut impl DbTxn, set: ExternalValidatorSet) {
    let mut active = Self::active_tributaries(txn).1;
    for i in 0 .. active.len() {
      if active[i].set() == set {
        active.remove(i);
        break;
      }
    }

    let mut bytes = vec![];
    for active in active {
      active.serialize(&mut bytes).unwrap();
    }
    Self::set(txn, &bytes);
    RetiredTributaryDb::set(txn, set, &());
  }
}

impl FirstPreprocessDb {
  pub fn save_first_preprocess(
    txn: &mut impl DbTxn,
    network: ExternalNetworkId,
    id_type: RecognizedIdType,
    id: &[u8],
    preprocess: &Vec<Vec<u8>>,
  ) {
    if let Some(existing) = FirstPreprocessDb::get(txn, network, id_type, id) {
      assert_eq!(&existing, preprocess, "saved a distinct first preprocess");
      return;
    }
    FirstPreprocessDb::set(txn, network, id_type, id, preprocess);
  }
}

impl ExpectedBatchDb {
  pub fn save_expected_batch(txn: &mut impl DbTxn, batch: &Batch) {
    LastReceivedBatchDb::set(txn, batch.network, &batch.id);
    Self::set(
      txn,
      batch.network,
      batch.id,
      &Blake2b::<U32>::digest(batch.instructions.encode()).into(),
    );
  }
}

impl HandoverBatchDb {
  pub fn set_handover_batch(txn: &mut impl DbTxn, set: ExternalValidatorSet, batch: u32) {
    Self::set(txn, set, &batch);
    LookupHandoverBatchDb::set(txn, set.network, batch, &set.session);
  }
}
impl QueuedBatchesDb {
  pub fn queue(txn: &mut impl DbTxn, set: ExternalValidatorSet, batch: &Transaction) {
    let mut batches = Self::get(txn, set).unwrap_or_default();
    batch.write(&mut batches).unwrap();
    Self::set(txn, set, &batches);
  }

  pub fn take(txn: &mut impl DbTxn, set: ExternalValidatorSet) -> Vec<Transaction> {
    let batches_vec = Self::get(txn, set).unwrap_or_default();
    txn.del(Self::key(set));

    let mut batches: &[u8] = &batches_vec;
    let mut res = vec![];
    while !batches.is_empty() {
      res.push(Transaction::read(&mut batches).unwrap());
    }
    res
  }
}