Skip to main content

hotmint_staking/
store.rs

1use std::collections::HashMap;
2
3use hotmint_types::validator::ValidatorId;
4
5use crate::types::{StakeEntry, UnbondingEntry, ValidatorState};
6
7/// Abstract storage backend for staking state.
8///
9/// Implement this trait to plug in any persistence layer (in-memory, vsdb,
10/// RocksDB, etc.). The staking manager operates entirely through this
11/// interface and never assumes a specific storage backend.
12pub 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    /// Return all (staker_address, entry) pairs for a given validator.
22    fn stakers_of(&self, validator: ValidatorId) -> Vec<(Vec<u8>, StakeEntry)>;
23
24    // ── Unbonding queue ─────────────────────────────────────────────
25
26    /// Append an entry to the unbonding queue.
27    fn push_unbonding(&mut self, entry: UnbondingEntry);
28
29    /// Remove and return all entries whose `completion_height <= current_height`.
30    fn drain_mature_unbondings(&mut self, current_height: u64) -> Vec<UnbondingEntry>;
31
32    /// Return all pending unbonding entries (for slashing).
33    fn all_unbondings(&self) -> Vec<UnbondingEntry>;
34
35    /// Replace the entire unbonding queue (used after slashing adjustments).
36    fn replace_unbondings(&mut self, entries: Vec<UnbondingEntry>);
37}
38
39/// In-memory staking store for testing and demos.
40#[derive(Default)]
41pub struct InMemoryStakingStore {
42    validators: HashMap<ValidatorId, ValidatorState>,
43    /// Key: (staker_address, validator_id)
44    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}