1use std::collections::HashMap;
2
3use hotmint_types::validator::ValidatorId;
4
5use crate::types::{StakeEntry, UnbondingEntry, ValidatorState};
6
7pub trait StakingStore {
13 fn get_validator(&self, id: ValidatorId) -> Option<ValidatorState>;
14 fn set_validator(&mut self, id: ValidatorId, state: ValidatorState);
15 fn remove_validator(&mut self, id: ValidatorId);
16 fn all_validator_ids(&self) -> Vec<ValidatorId>;
17
18 fn get_stake(&self, staker: &[u8], validator: ValidatorId) -> Option<StakeEntry>;
19 fn set_stake(&mut self, staker: &[u8], validator: ValidatorId, entry: StakeEntry);
20 fn remove_stake(&mut self, staker: &[u8], validator: ValidatorId);
21 fn stakers_of(&self, validator: ValidatorId) -> Vec<(Vec<u8>, StakeEntry)>;
23
24 fn push_unbonding(&mut self, entry: UnbondingEntry);
28
29 fn drain_mature_unbondings(&mut self, current_height: u64) -> Vec<UnbondingEntry>;
31
32 fn all_unbondings(&self) -> Vec<UnbondingEntry>;
34
35 fn replace_unbondings(&mut self, entries: Vec<UnbondingEntry>);
37}
38
39#[derive(Default)]
41pub struct InMemoryStakingStore {
42 validators: HashMap<ValidatorId, ValidatorState>,
43 stakes: HashMap<(Vec<u8>, ValidatorId), StakeEntry>,
45 unbondings: Vec<UnbondingEntry>,
46}
47
48impl InMemoryStakingStore {
49 pub fn new() -> Self {
50 Self::default()
51 }
52}
53
54impl StakingStore for InMemoryStakingStore {
55 fn get_validator(&self, id: ValidatorId) -> Option<ValidatorState> {
56 self.validators.get(&id).cloned()
57 }
58
59 fn set_validator(&mut self, id: ValidatorId, state: ValidatorState) {
60 self.validators.insert(id, state);
61 }
62
63 fn remove_validator(&mut self, id: ValidatorId) {
64 self.validators.remove(&id);
65 }
66
67 fn all_validator_ids(&self) -> Vec<ValidatorId> {
68 self.validators.keys().copied().collect()
69 }
70
71 fn get_stake(&self, staker: &[u8], validator: ValidatorId) -> Option<StakeEntry> {
72 self.stakes.get(&(staker.to_vec(), validator)).cloned()
73 }
74
75 fn set_stake(&mut self, staker: &[u8], validator: ValidatorId, entry: StakeEntry) {
76 self.stakes.insert((staker.to_vec(), validator), entry);
77 }
78
79 fn remove_stake(&mut self, staker: &[u8], validator: ValidatorId) {
80 self.stakes.remove(&(staker.to_vec(), validator));
81 }
82
83 fn stakers_of(&self, validator: ValidatorId) -> Vec<(Vec<u8>, StakeEntry)> {
84 self.stakes
85 .iter()
86 .filter(|((_, vid), _)| *vid == validator)
87 .map(|((addr, _), entry)| (addr.clone(), entry.clone()))
88 .collect()
89 }
90
91 fn push_unbonding(&mut self, entry: UnbondingEntry) {
92 self.unbondings.push(entry);
93 }
94
95 fn drain_mature_unbondings(&mut self, current_height: u64) -> Vec<UnbondingEntry> {
96 let (mature, pending): (Vec<_>, Vec<_>) = self
97 .unbondings
98 .drain(..)
99 .partition(|e| e.completion_height <= current_height);
100 self.unbondings = pending;
101 mature
102 }
103
104 fn all_unbondings(&self) -> Vec<UnbondingEntry> {
105 self.unbondings.clone()
106 }
107
108 fn replace_unbondings(&mut self, entries: Vec<UnbondingEntry>) {
109 self.unbondings = entries;
110 }
111}