oasis_runtime_sdk/modules/consensus_accounts/
state.rs

1//! State schema.
2use std::{
3    collections::BTreeMap,
4    convert::{TryFrom, TryInto},
5};
6
7use oasis_core_runtime::consensus::beacon::EpochTime;
8
9use crate::{
10    state::CurrentState,
11    storage::{self, Store},
12    types::address::Address,
13};
14
15use super::{types, Error, MODULE_NAME};
16
17/// Map of active delegations.
18pub const DELEGATIONS: &[u8] = &[0x01];
19/// Map of undelegations.
20pub const UNDELEGATIONS: &[u8] = &[0x02];
21/// An undelegation queue.
22pub const UNDELEGATION_QUEUE: &[u8] = &[0x03];
23/// Receipts.
24pub const RECEIPTS: &[u8] = &[0x04];
25
26/// Add delegation for a given (from, to) pair.
27///
28/// The given shares are added to any existing delegation that may exist for the same (from, to)
29/// address pair. If no delegation exists a new one is created.
30pub fn add_delegation(from: Address, to: Address, shares: u128) -> Result<(), Error> {
31    CurrentState::with_store(|store| {
32        let store = storage::PrefixStore::new(store, &MODULE_NAME);
33        let delegations = storage::PrefixStore::new(store, &DELEGATIONS);
34        let mut account = storage::TypedStore::new(storage::PrefixStore::new(delegations, &from));
35        let mut di: types::DelegationInfo = account.get(to).unwrap_or_default();
36
37        di.shares = di
38            .shares
39            .checked_add(shares)
40            .ok_or(Error::InvalidArgument)?;
41
42        account.insert(to, di);
43
44        Ok(())
45    })
46}
47
48/// Subtract delegation from a given (from, to) pair.
49pub fn sub_delegation(from: Address, to: Address, shares: u128) -> Result<(), Error> {
50    CurrentState::with_store(|store| {
51        let store = storage::PrefixStore::new(store, &MODULE_NAME);
52        let delegations = storage::PrefixStore::new(store, &DELEGATIONS);
53        let mut account = storage::TypedStore::new(storage::PrefixStore::new(delegations, &from));
54        let mut di: types::DelegationInfo = account.get(to).unwrap_or_default();
55
56        di.shares = di
57            .shares
58            .checked_sub(shares)
59            .ok_or(Error::InsufficientBalance)?;
60
61        if di.shares > 0 {
62            account.insert(to, di);
63        } else {
64            account.remove(to);
65        }
66
67        Ok(())
68    })
69}
70
71/// Retrieve delegation metadata for a given (from, to) pair.
72///
73/// In case no delegation exists for the given (from, to) address pair, an all-zero delegation
74/// metadata are returned.
75pub fn get_delegation(from: Address, to: Address) -> Result<types::DelegationInfo, Error> {
76    CurrentState::with_store(|store| {
77        let store = storage::PrefixStore::new(store, &MODULE_NAME);
78        let delegations = storage::PrefixStore::new(store, &DELEGATIONS);
79        let account = storage::TypedStore::new(storage::PrefixStore::new(delegations, &from));
80        Ok(account.get(to).unwrap_or_default())
81    })
82}
83
84/// Retrieve all delegation metadata originating from a given address.
85pub fn get_delegations(from: Address) -> Result<Vec<types::ExtendedDelegationInfo>, Error> {
86    CurrentState::with_store(|store| {
87        let store = storage::PrefixStore::new(store, &MODULE_NAME);
88        let delegations = storage::PrefixStore::new(store, &DELEGATIONS);
89        let account = storage::TypedStore::new(storage::PrefixStore::new(delegations, &from));
90
91        Ok(account
92            .iter()
93            .map(
94                |(to, di): (Address, types::DelegationInfo)| -> types::ExtendedDelegationInfo {
95                    types::ExtendedDelegationInfo {
96                        to,
97                        shares: di.shares,
98                    }
99                },
100            )
101            .collect())
102    })
103}
104
105/// This is needed to properly iterate over the DELEGATIONS map.
106#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
107struct AddressPair(Address, Address);
108
109#[derive(Error, Debug)]
110enum APError {
111    #[error("malformed address")]
112    MalformedAddress,
113}
114
115impl TryFrom<&[u8]> for AddressPair {
116    type Error = APError;
117
118    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
119        let a =
120            Address::try_from(&bytes[..Address::SIZE]).map_err(|_| APError::MalformedAddress)?;
121        let b =
122            Address::try_from(&bytes[Address::SIZE..]).map_err(|_| APError::MalformedAddress)?;
123        Ok(AddressPair(a, b))
124    }
125}
126
127/// Return the number of delegated shares for each destination escrow account.
128pub fn get_delegations_by_destination() -> Result<BTreeMap<Address, u128>, Error> {
129    CurrentState::with_store(|store| {
130        let store = storage::PrefixStore::new(store, &MODULE_NAME);
131        let delegations = storage::TypedStore::new(storage::PrefixStore::new(store, &DELEGATIONS));
132
133        let mut by_destination: BTreeMap<Address, u128> = BTreeMap::new();
134        for (ap, di) in delegations.iter::<AddressPair, types::DelegationInfo>() {
135            let total = by_destination.entry(ap.1).or_default();
136            *total = total.checked_add(di.shares).ok_or(Error::InvalidArgument)?;
137        }
138
139        Ok(by_destination)
140    })
141}
142
143/// Record new undelegation and add to undelegation queue.
144///
145/// In case an undelegation for the given (from, to, epoch) tuple already exists, the undelegation
146/// entry is merged by adding shares. When a non-zero receipt identifier is passed, the identifier
147/// is set in case the existing entry has no such identifier yet.
148///
149/// It returns the receipt identifier of the undelegation done receipt.
150pub fn add_undelegation(
151    from: Address,
152    to: Address,
153    epoch: EpochTime,
154    shares: u128,
155    receipt: u64,
156) -> Result<u64, Error> {
157    CurrentState::with_store(|mut root_store| {
158        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
159        let undelegations = storage::PrefixStore::new(store, &UNDELEGATIONS);
160        let account = storage::PrefixStore::new(undelegations, &to);
161        let mut entry = storage::TypedStore::new(storage::PrefixStore::new(account, &from));
162        let mut di: types::DelegationInfo = entry.get(epoch.to_storage_key()).unwrap_or_default();
163
164        if receipt > 0 && di.receipt == 0 {
165            di.receipt = receipt;
166        }
167        let done_receipt = di.receipt;
168
169        di.shares = di
170            .shares
171            .checked_add(shares)
172            .ok_or(Error::InvalidArgument)?;
173
174        entry.insert(epoch.to_storage_key(), di);
175
176        // Add to undelegation queue (if existing item is there, this will have no effect).
177        let store = storage::PrefixStore::new(root_store, &MODULE_NAME);
178        let mut queue = storage::PrefixStore::new(store, &UNDELEGATION_QUEUE);
179        queue.insert(
180            &queue_entry_key(from, to, epoch),
181            &[0xF6], /* CBOR NULL */
182        );
183
184        Ok(done_receipt)
185    })
186}
187
188fn queue_entry_key(from: Address, to: Address, epoch: EpochTime) -> Vec<u8> {
189    [&epoch.to_storage_key(), to.as_ref(), from.as_ref()].concat()
190}
191
192/// Remove an existing undelegation and return it.
193///
194/// In case the undelegation doesn't exist, returns a default-constructed DelegationInfo.
195pub fn take_undelegation(ud: &Undelegation) -> Result<types::DelegationInfo, Error> {
196    CurrentState::with_store(|mut root_store| {
197        // Get and remove undelegation metadata.
198        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
199        let undelegations = storage::PrefixStore::new(store, &UNDELEGATIONS);
200        let account = storage::PrefixStore::new(undelegations, &ud.to);
201        let mut entry = storage::TypedStore::new(storage::PrefixStore::new(account, &ud.from));
202        let di: types::DelegationInfo = entry.get(ud.epoch.to_storage_key()).unwrap_or_default();
203        entry.remove(ud.epoch.to_storage_key());
204
205        // Remove queue entry.
206        let store = storage::PrefixStore::new(root_store, &MODULE_NAME);
207        let mut queue = storage::PrefixStore::new(store, &UNDELEGATION_QUEUE);
208        queue.remove(&queue_entry_key(ud.from, ud.to, ud.epoch));
209
210        Ok(di)
211    })
212}
213
214struct AddressWithEpoch {
215    from: Address,
216    epoch: EpochTime,
217}
218
219impl TryFrom<&[u8]> for AddressWithEpoch {
220    type Error = anyhow::Error;
221
222    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
223        if value.len() != Address::SIZE + 8 {
224            anyhow::bail!("incorrect address with epoch key size");
225        }
226
227        Ok(Self {
228            from: Address::try_from(&value[..Address::SIZE])?,
229            epoch: EpochTime::from_be_bytes(value[Address::SIZE..].try_into()?),
230        })
231    }
232}
233
234/// Retrieve all undelegation metadata to a given address.
235pub fn get_undelegations(to: Address) -> Result<Vec<types::UndelegationInfo>, Error> {
236    CurrentState::with_store(|store| {
237        let store = storage::PrefixStore::new(store, &MODULE_NAME);
238        let undelegations = storage::PrefixStore::new(store, &UNDELEGATIONS);
239        let account = storage::TypedStore::new(storage::PrefixStore::new(undelegations, &to));
240
241        Ok(account
242            .iter()
243            .map(
244                |(ae, di): (AddressWithEpoch, types::DelegationInfo)| -> types::UndelegationInfo {
245                    types::UndelegationInfo {
246                        from: ae.from,
247                        epoch: ae.epoch,
248                        shares: di.shares,
249                    }
250                },
251            )
252            .collect())
253    })
254}
255
256/// Undelegation metadata.
257pub struct Undelegation {
258    pub from: Address,
259    pub to: Address,
260    pub epoch: EpochTime,
261}
262
263impl<'a> TryFrom<&'a [u8]> for Undelegation {
264    type Error = anyhow::Error;
265
266    fn try_from(value: &'a [u8]) -> Result<Self, Self::Error> {
267        // Decode a storage key of the format (epoch, to, from).
268        if value.len() != 2 * Address::SIZE + 8 {
269            anyhow::bail!("incorrect undelegation key size");
270        }
271
272        Ok(Self {
273            epoch: EpochTime::from_be_bytes(value[..8].try_into()?),
274            to: Address::from_bytes(&value[8..8 + Address::SIZE])?,
275            from: Address::from_bytes(&value[8 + Address::SIZE..])?,
276        })
277    }
278}
279
280/// Retrieve all queued undelegations for epochs earlier than or equal to the passed epoch.
281pub fn get_queued_undelegations(epoch: EpochTime) -> Result<Vec<Undelegation>, Error> {
282    CurrentState::with_store(|store| {
283        let store = storage::PrefixStore::new(store, &MODULE_NAME);
284        let queue = storage::TypedStore::new(storage::PrefixStore::new(store, &UNDELEGATION_QUEUE));
285
286        Ok(queue
287            .iter()
288            .map(|(k, _): (Undelegation, ())| k)
289            .take_while(|ud| ud.epoch <= epoch)
290            .collect())
291    })
292}
293
294/// Store the given receipt.
295pub fn set_receipt(owner: Address, kind: types::ReceiptKind, id: u64, receipt: types::Receipt) {
296    CurrentState::with_store(|store| {
297        let store = storage::PrefixStore::new(store, &MODULE_NAME);
298        let receipts = storage::PrefixStore::new(store, &RECEIPTS);
299        let of_owner = storage::PrefixStore::new(receipts, &owner);
300        let kind = [kind as u8];
301        let mut of_kind = storage::TypedStore::new(storage::PrefixStore::new(of_owner, &kind));
302
303        of_kind.insert(id.to_be_bytes(), receipt);
304    });
305}
306
307/// Remove the given receipt from storage if it exists and return it, otherwise return `None`.
308pub fn take_receipt(owner: Address, kind: types::ReceiptKind, id: u64) -> Option<types::Receipt> {
309    CurrentState::with_store(|store| {
310        let store = storage::PrefixStore::new(store, &MODULE_NAME);
311        let receipts = storage::PrefixStore::new(store, &RECEIPTS);
312        let of_owner = storage::PrefixStore::new(receipts, &owner);
313        let kind = [kind as u8];
314        let mut of_kind = storage::TypedStore::new(storage::PrefixStore::new(of_owner, &kind));
315
316        let receipt = of_kind.get(id.to_be_bytes());
317        of_kind.remove(id.to_be_bytes());
318
319        receipt
320    })
321}
322
323/// A trait that exists solely to convert `beacon::EpochTime` to bytes for use as a storage key.
324trait ToStorageKey {
325    fn to_storage_key(&self) -> [u8; 8];
326}
327
328impl ToStorageKey for EpochTime {
329    fn to_storage_key(&self) -> [u8; 8] {
330        self.to_be_bytes()
331    }
332}
333
334#[cfg(test)]
335mod test {
336    use super::*;
337    use crate::testing::{keys, mock};
338
339    #[test]
340    fn test_delegation() {
341        let _mock = mock::Mock::default();
342
343        add_delegation(keys::alice::address(), keys::bob::address(), 500).unwrap();
344        add_delegation(keys::alice::address(), keys::bob::address(), 500).unwrap();
345
346        let di = get_delegation(keys::bob::address(), keys::alice::address()).unwrap();
347        assert_eq!(di.shares, 0);
348        let di = get_delegation(keys::alice::address(), keys::bob::address()).unwrap();
349        assert_eq!(di.shares, 1000);
350
351        let dis = get_delegations(keys::bob::address()).unwrap();
352        assert!(dis.is_empty());
353        let dis = get_delegations(keys::alice::address()).unwrap();
354        assert_eq!(dis.len(), 1);
355        assert_eq!(dis[0].shares, 1000);
356
357        let totals = get_delegations_by_destination().unwrap();
358        assert_eq!(totals.len(), 1);
359        assert_eq!(totals[&keys::bob::address()], 1000);
360
361        sub_delegation(keys::alice::address(), keys::bob::address(), 100).unwrap();
362
363        let di = get_delegation(keys::alice::address(), keys::bob::address()).unwrap();
364        assert_eq!(di.shares, 900);
365
366        let totals = get_delegations_by_destination().unwrap();
367        assert_eq!(totals.len(), 1);
368        assert_eq!(totals[&keys::bob::address()], 900);
369
370        add_delegation(keys::bob::address(), keys::bob::address(), 200).unwrap();
371
372        let totals = get_delegations_by_destination().unwrap();
373        assert_eq!(totals.len(), 1);
374        assert_eq!(totals[&keys::bob::address()], 1100);
375
376        add_delegation(keys::bob::address(), keys::alice::address(), 100).unwrap();
377
378        let totals = get_delegations_by_destination().unwrap();
379        assert_eq!(totals.len(), 2);
380        assert_eq!(totals[&keys::alice::address()], 100);
381        assert_eq!(totals[&keys::bob::address()], 1100);
382    }
383
384    #[test]
385    fn test_undelegation() {
386        let _mock = mock::Mock::default();
387
388        add_undelegation(keys::alice::address(), keys::bob::address(), 42, 500, 12).unwrap();
389        add_undelegation(keys::alice::address(), keys::bob::address(), 42, 500, 24).unwrap();
390        add_undelegation(keys::alice::address(), keys::bob::address(), 84, 200, 36).unwrap();
391
392        let qd = get_queued_undelegations(10).unwrap();
393        assert!(qd.is_empty());
394        let qd = get_queued_undelegations(42).unwrap();
395        assert_eq!(qd.len(), 1);
396        assert_eq!(qd[0].from, keys::alice::address());
397        assert_eq!(qd[0].to, keys::bob::address());
398        assert_eq!(qd[0].epoch, 42);
399        let qd = get_queued_undelegations(43).unwrap();
400        assert_eq!(qd.len(), 1);
401        assert_eq!(qd[0].from, keys::alice::address());
402        assert_eq!(qd[0].to, keys::bob::address());
403        assert_eq!(qd[0].epoch, 42);
404
405        let udis = get_undelegations(keys::alice::address()).unwrap();
406        assert!(udis.is_empty());
407        let udis = get_undelegations(keys::bob::address()).unwrap();
408        assert_eq!(udis.len(), 2);
409        assert_eq!(udis[0].from, keys::alice::address());
410        assert_eq!(udis[0].shares, 1000);
411        assert_eq!(udis[0].epoch, 42);
412        assert_eq!(udis[1].from, keys::alice::address());
413        assert_eq!(udis[1].shares, 200);
414        assert_eq!(udis[1].epoch, 84);
415
416        let di = take_undelegation(&qd[0]).unwrap();
417        assert_eq!(di.shares, 1000);
418        assert_eq!(di.receipt, 12, "receipt id should not be overwritten");
419
420        let qd = get_queued_undelegations(42).unwrap();
421        assert!(qd.is_empty());
422
423        let udis = get_undelegations(keys::bob::address()).unwrap();
424        assert_eq!(udis.len(), 1);
425    }
426
427    #[test]
428    fn test_receipts() {
429        let _mock = mock::Mock::default();
430
431        let receipt = types::Receipt {
432            shares: 123,
433            ..Default::default()
434        };
435        set_receipt(
436            keys::alice::address(),
437            types::ReceiptKind::Delegate,
438            42,
439            receipt.clone(),
440        );
441
442        let dec_receipt = take_receipt(keys::alice::address(), types::ReceiptKind::Delegate, 10);
443        assert!(dec_receipt.is_none(), "missing receipt should return None");
444
445        let dec_receipt = take_receipt(
446            keys::alice::address(),
447            types::ReceiptKind::UndelegateStart,
448            42,
449        );
450        assert!(dec_receipt.is_none(), "missing receipt should return None");
451
452        let dec_receipt = take_receipt(
453            keys::alice::address(),
454            types::ReceiptKind::UndelegateDone,
455            42,
456        );
457        assert!(dec_receipt.is_none(), "missing receipt should return None");
458
459        let dec_receipt = take_receipt(keys::alice::address(), types::ReceiptKind::Delegate, 42);
460        assert_eq!(dec_receipt, Some(receipt), "receipt should be correct");
461
462        let dec_receipt = take_receipt(keys::alice::address(), types::ReceiptKind::Delegate, 42);
463        assert!(dec_receipt.is_none(), "receipt should have been removed");
464    }
465}
oasis_runtime_sdk/modules/consensus_accounts/state.rs

oasis_runtime_sdk/modules/consensus_accounts/
state.rs
