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 (AddressPair(_, to), di) in delegations.iter::<AddressPair, types::DelegationInfo>() {
135            let total = by_destination.entry(to).or_default();
136            *total = total.checked_add(di.shares).ok_or(Error::InvalidArgument)?;
137        }
138
139        Ok(by_destination)
140    })
141}
142
143/// Retrive delegation metadata for all delegations.
144pub fn get_all_delegations() -> Result<Vec<types::CompleteDelegationInfo>, Error> {
145    CurrentState::with_store(|store| {
146        let store = storage::PrefixStore::new(store, &MODULE_NAME);
147        let delegations = storage::TypedStore::new(storage::PrefixStore::new(store, &DELEGATIONS));
148
149        Ok(delegations
150            .iter::<AddressPair, types::DelegationInfo>()
151            .map(
152                |(AddressPair(from, to), di)| types::CompleteDelegationInfo {
153                    from,
154                    to,
155                    shares: di.shares,
156                },
157            )
158            .collect())
159    })
160}
161
162/// Record new undelegation and add to undelegation queue.
163///
164/// In case an undelegation for the given (from, to, epoch) tuple already exists, the undelegation
165/// entry is merged by adding shares. When a non-zero receipt identifier is passed, the identifier
166/// is set in case the existing entry has no such identifier yet.
167///
168/// It returns the receipt identifier of the undelegation done receipt.
169pub fn add_undelegation(
170    from: Address,
171    to: Address,
172    epoch: EpochTime,
173    shares: u128,
174    receipt: u64,
175) -> Result<u64, Error> {
176    CurrentState::with_store(|mut root_store| {
177        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
178        let undelegations = storage::PrefixStore::new(store, &UNDELEGATIONS);
179        let account = storage::PrefixStore::new(undelegations, &to);
180        let mut entry = storage::TypedStore::new(storage::PrefixStore::new(account, &from));
181        let mut di: types::DelegationInfo = entry.get(epoch.to_storage_key()).unwrap_or_default();
182
183        if receipt > 0 && di.receipt == 0 {
184            di.receipt = receipt;
185        }
186        let done_receipt = di.receipt;
187
188        di.shares = di
189            .shares
190            .checked_add(shares)
191            .ok_or(Error::InvalidArgument)?;
192
193        entry.insert(epoch.to_storage_key(), di);
194
195        // Add to undelegation queue (if existing item is there, this will have no effect).
196        let store = storage::PrefixStore::new(root_store, &MODULE_NAME);
197        let mut queue = storage::PrefixStore::new(store, &UNDELEGATION_QUEUE);
198        queue.insert(
199            &queue_entry_key(from, to, epoch),
200            &[0xF6], /* CBOR NULL */
201        );
202
203        Ok(done_receipt)
204    })
205}
206
207fn queue_entry_key(from: Address, to: Address, epoch: EpochTime) -> Vec<u8> {
208    [&epoch.to_storage_key(), to.as_ref(), from.as_ref()].concat()
209}
210
211/// Remove an existing undelegation and return it.
212///
213/// In case the undelegation doesn't exist, returns a default-constructed DelegationInfo.
214pub fn take_undelegation(ud: &Undelegation) -> Result<types::DelegationInfo, Error> {
215    CurrentState::with_store(|mut root_store| {
216        // Get and remove undelegation metadata.
217        let store = storage::PrefixStore::new(&mut root_store, &MODULE_NAME);
218        let undelegations = storage::PrefixStore::new(store, &UNDELEGATIONS);
219        let account = storage::PrefixStore::new(undelegations, &ud.to);
220        let mut entry = storage::TypedStore::new(storage::PrefixStore::new(account, &ud.from));
221        let di: types::DelegationInfo = entry.get(ud.epoch.to_storage_key()).unwrap_or_default();
222        entry.remove(ud.epoch.to_storage_key());
223
224        // Remove queue entry.
225        let store = storage::PrefixStore::new(root_store, &MODULE_NAME);
226        let mut queue = storage::PrefixStore::new(store, &UNDELEGATION_QUEUE);
227        queue.remove(&queue_entry_key(ud.from, ud.to, ud.epoch));
228
229        Ok(di)
230    })
231}
232
233struct AddressWithEpoch {
234    from: Address,
235    epoch: EpochTime,
236}
237
238impl TryFrom<&[u8]> for AddressWithEpoch {
239    type Error = anyhow::Error;
240
241    fn try_from(value: &[u8]) -> Result<Self, Self::Error> {
242        if value.len() != Address::SIZE + 8 {
243            anyhow::bail!("incorrect address with epoch key size");
244        }
245
246        Ok(Self {
247            from: Address::try_from(&value[..Address::SIZE])?,
248            epoch: EpochTime::from_be_bytes(value[Address::SIZE..].try_into()?),
249        })
250    }
251}
252
253/// Retrieve all undelegation metadata to a given address.
254pub fn get_undelegations(to: Address) -> Result<Vec<types::UndelegationInfo>, Error> {
255    CurrentState::with_store(|store| {
256        let store = storage::PrefixStore::new(store, &MODULE_NAME);
257        let undelegations = storage::PrefixStore::new(store, &UNDELEGATIONS);
258        let account = storage::TypedStore::new(storage::PrefixStore::new(undelegations, &to));
259
260        Ok(account
261            .iter()
262            .map(
263                |(ae, di): (AddressWithEpoch, types::DelegationInfo)| -> types::UndelegationInfo {
264                    types::UndelegationInfo {
265                        from: ae.from,
266                        epoch: ae.epoch,
267                        shares: di.shares,
268                    }
269                },
270            )
271            .collect())
272    })
273}
274
275/// This is needed to properly iterate over the UNDELEGATIONS map.
276#[derive(Clone, PartialEq, PartialOrd, Eq, Ord)]
277struct AddressPairWithEpoch(AddressPair, EpochTime);
278
279impl TryFrom<&[u8]> for AddressPairWithEpoch {
280    type Error = anyhow::Error;
281
282    fn try_from(bytes: &[u8]) -> Result<Self, Self::Error> {
283        if bytes.len() != (2 * Address::SIZE) + 8 {
284            anyhow::bail!("incorrect address pair with epoch key size");
285        }
286        let a = Address::try_from(&bytes[..Address::SIZE])
287            .map_err(|_| anyhow::anyhow!("malformed address"))?;
288        let b = Address::try_from(&bytes[Address::SIZE..bytes.len() - 8])
289            .map_err(|_| anyhow::anyhow!("malformed address"))?;
290        let epoch = EpochTime::from_be_bytes(bytes[bytes.len() - 8..].try_into()?);
291        Ok(AddressPairWithEpoch(AddressPair(a, b), epoch))
292    }
293}
294
295/// Retrieve all undelegation metadata.
296pub fn get_all_undelegations() -> Result<Vec<types::CompleteUndelegationInfo>, Error> {
297    CurrentState::with_store(|store| {
298        let store = storage::PrefixStore::new(store, &MODULE_NAME);
299        let undelegations =
300            storage::TypedStore::new(storage::PrefixStore::new(store, &UNDELEGATIONS));
301
302        Ok(undelegations
303            .iter::<AddressPairWithEpoch, types::DelegationInfo>()
304            .map(|(AddressPairWithEpoch(AddressPair(to, from), epoch), di)| {
305                types::CompleteUndelegationInfo {
306                    from,
307                    to,
308                    epoch,
309                    shares: di.shares,
310                }
311            })
312            .collect())
313    })
314}
315
316/// Undelegation metadata.
317pub struct Undelegation {
318    pub from: Address,
319    pub to: Address,
320    pub epoch: EpochTime,
321}
322
323impl<'a> TryFrom<&'a [u8]> for Undelegation {
324    type Error = anyhow::Error;
325
326    fn try_from(value: &'a [u8]) -> Result<Self, Self::Error> {
327        // Decode a storage key of the format (epoch, to, from).
328        if value.len() != 2 * Address::SIZE + 8 {
329            anyhow::bail!("incorrect undelegation key size");
330        }
331
332        Ok(Self {
333            epoch: EpochTime::from_be_bytes(value[..8].try_into()?),
334            to: Address::from_bytes(&value[8..8 + Address::SIZE])?,
335            from: Address::from_bytes(&value[8 + Address::SIZE..])?,
336        })
337    }
338}
339
340/// Retrieve all queued undelegations for epochs earlier than or equal to the passed epoch.
341pub fn get_queued_undelegations(epoch: EpochTime) -> Result<Vec<Undelegation>, Error> {
342    CurrentState::with_store(|store| {
343        let store = storage::PrefixStore::new(store, &MODULE_NAME);
344        let queue = storage::TypedStore::new(storage::PrefixStore::new(store, &UNDELEGATION_QUEUE));
345
346        Ok(queue
347            .iter()
348            .map(|(k, _): (Undelegation, ())| k)
349            .take_while(|ud| ud.epoch <= epoch)
350            .collect())
351    })
352}
353
354/// Store the given receipt.
355pub fn set_receipt(owner: Address, kind: types::ReceiptKind, id: u64, receipt: types::Receipt) {
356    CurrentState::with_store(|store| {
357        let store = storage::PrefixStore::new(store, &MODULE_NAME);
358        let receipts = storage::PrefixStore::new(store, &RECEIPTS);
359        let of_owner = storage::PrefixStore::new(receipts, &owner);
360        let kind = [kind as u8];
361        let mut of_kind = storage::TypedStore::new(storage::PrefixStore::new(of_owner, &kind));
362
363        of_kind.insert(id.to_be_bytes(), receipt);
364    });
365}
366
367/// Remove the given receipt from storage if it exists and return it, otherwise return `None`.
368pub fn take_receipt(owner: Address, kind: types::ReceiptKind, id: u64) -> Option<types::Receipt> {
369    CurrentState::with_store(|store| {
370        let store = storage::PrefixStore::new(store, &MODULE_NAME);
371        let receipts = storage::PrefixStore::new(store, &RECEIPTS);
372        let of_owner = storage::PrefixStore::new(receipts, &owner);
373        let kind = [kind as u8];
374        let mut of_kind = storage::TypedStore::new(storage::PrefixStore::new(of_owner, &kind));
375
376        let receipt = of_kind.get(id.to_be_bytes());
377        of_kind.remove(id.to_be_bytes());
378
379        receipt
380    })
381}
382
383/// A trait that exists solely to convert `beacon::EpochTime` to bytes for use as a storage key.
384trait ToStorageKey {
385    fn to_storage_key(&self) -> [u8; 8];
386}
387
388impl ToStorageKey for EpochTime {
389    fn to_storage_key(&self) -> [u8; 8] {
390        self.to_be_bytes()
391    }
392}
393
394#[cfg(test)]
395mod test {
396    use super::*;
397    use crate::testing::{keys, mock};
398
399    #[test]
400    fn test_delegation() {
401        let _mock = mock::Mock::default();
402
403        add_delegation(keys::alice::address(), keys::bob::address(), 500).unwrap();
404        add_delegation(keys::alice::address(), keys::bob::address(), 500).unwrap();
405
406        let di = get_delegation(keys::bob::address(), keys::alice::address()).unwrap();
407        assert_eq!(di.shares, 0);
408        let di = get_delegation(keys::alice::address(), keys::bob::address()).unwrap();
409        assert_eq!(di.shares, 1000);
410
411        let dis = get_delegations(keys::bob::address()).unwrap();
412        assert!(dis.is_empty());
413        let dis = get_delegations(keys::alice::address()).unwrap();
414        assert_eq!(dis.len(), 1);
415        assert_eq!(dis[0].shares, 1000);
416
417        let totals = get_delegations_by_destination().unwrap();
418        assert_eq!(totals.len(), 1);
419        assert_eq!(totals[&keys::bob::address()], 1000);
420
421        sub_delegation(keys::alice::address(), keys::bob::address(), 100).unwrap();
422
423        let di = get_delegation(keys::alice::address(), keys::bob::address()).unwrap();
424        assert_eq!(di.shares, 900);
425
426        let totals = get_delegations_by_destination().unwrap();
427        assert_eq!(totals.len(), 1);
428        assert_eq!(totals[&keys::bob::address()], 900);
429
430        add_delegation(keys::bob::address(), keys::bob::address(), 200).unwrap();
431
432        let totals = get_delegations_by_destination().unwrap();
433        assert_eq!(totals.len(), 1);
434        assert_eq!(totals[&keys::bob::address()], 1100);
435
436        add_delegation(keys::bob::address(), keys::alice::address(), 100).unwrap();
437
438        let totals = get_delegations_by_destination().unwrap();
439        assert_eq!(totals.len(), 2);
440        assert_eq!(totals[&keys::alice::address()], 100);
441        assert_eq!(totals[&keys::bob::address()], 1100);
442    }
443
444    #[test]
445    fn test_undelegation() {
446        let _mock = mock::Mock::default();
447
448        add_undelegation(keys::alice::address(), keys::bob::address(), 42, 500, 12).unwrap();
449        add_undelegation(keys::alice::address(), keys::bob::address(), 42, 500, 24).unwrap();
450        add_undelegation(keys::alice::address(), keys::bob::address(), 84, 200, 36).unwrap();
451
452        let qd = get_queued_undelegations(10).unwrap();
453        assert!(qd.is_empty());
454        let qd = get_queued_undelegations(42).unwrap();
455        assert_eq!(qd.len(), 1);
456        assert_eq!(qd[0].from, keys::alice::address());
457        assert_eq!(qd[0].to, keys::bob::address());
458        assert_eq!(qd[0].epoch, 42);
459        let qd = get_queued_undelegations(43).unwrap();
460        assert_eq!(qd.len(), 1);
461        assert_eq!(qd[0].from, keys::alice::address());
462        assert_eq!(qd[0].to, keys::bob::address());
463        assert_eq!(qd[0].epoch, 42);
464
465        let udis = get_undelegations(keys::alice::address()).unwrap();
466        assert!(udis.is_empty());
467        let udis = get_undelegations(keys::bob::address()).unwrap();
468        assert_eq!(udis.len(), 2);
469        assert_eq!(udis[0].from, keys::alice::address());
470        assert_eq!(udis[0].shares, 1000);
471        assert_eq!(udis[0].epoch, 42);
472        assert_eq!(udis[1].from, keys::alice::address());
473        assert_eq!(udis[1].shares, 200);
474        assert_eq!(udis[1].epoch, 84);
475
476        let di = take_undelegation(&qd[0]).unwrap();
477        assert_eq!(di.shares, 1000);
478        assert_eq!(di.receipt, 12, "receipt id should not be overwritten");
479
480        let qd = get_queued_undelegations(42).unwrap();
481        assert!(qd.is_empty());
482
483        let udis = get_undelegations(keys::bob::address()).unwrap();
484        assert_eq!(udis.len(), 1);
485    }
486
487    #[test]
488    fn test_receipts() {
489        let _mock = mock::Mock::default();
490
491        let receipt = types::Receipt {
492            shares: 123,
493            ..Default::default()
494        };
495        set_receipt(
496            keys::alice::address(),
497            types::ReceiptKind::Delegate,
498            42,
499            receipt.clone(),
500        );
501
502        let dec_receipt = take_receipt(keys::alice::address(), types::ReceiptKind::Delegate, 10);
503        assert!(dec_receipt.is_none(), "missing receipt should return None");
504
505        let dec_receipt = take_receipt(
506            keys::alice::address(),
507            types::ReceiptKind::UndelegateStart,
508            42,
509        );
510        assert!(dec_receipt.is_none(), "missing receipt should return None");
511
512        let dec_receipt = take_receipt(
513            keys::alice::address(),
514            types::ReceiptKind::UndelegateDone,
515            42,
516        );
517        assert!(dec_receipt.is_none(), "missing receipt should return None");
518
519        let dec_receipt = take_receipt(keys::alice::address(), types::ReceiptKind::Delegate, 42);
520        assert_eq!(dec_receipt, Some(receipt), "receipt should be correct");
521
522        let dec_receipt = take_receipt(keys::alice::address(), types::ReceiptKind::Delegate, 42);
523        assert!(dec_receipt.is_none(), "receipt should have been removed");
524    }
525}
oasis_runtime_sdk/modules/consensus_accounts/state.rs

oasis_runtime_sdk/modules/consensus_accounts/
state.rs
