Skip to main content

hotmint_consensus/
store.rs

1use std::collections::{BTreeMap, HashMap};
2
3use hotmint_types::{Block, BlockHash, Height};
4
5pub trait BlockStore: Send + Sync {
6    fn put_block(&mut self, block: Block);
7    fn get_block(&self, hash: &BlockHash) -> Option<Block>;
8    fn get_block_by_height(&self, h: Height) -> Option<Block>;
9
10    /// Get blocks in [from, to] inclusive. Default iterates one-by-one.
11    fn get_blocks_in_range(&self, from: Height, to: Height) -> Vec<Block> {
12        let mut blocks = Vec::new();
13        let mut h = from.as_u64();
14        while h <= to.as_u64() {
15            if let Some(block) = self.get_block_by_height(Height(h)) {
16                blocks.push(block);
17            }
18            h += 1;
19        }
20        blocks
21    }
22
23    /// Return the highest stored block height.
24    fn tip_height(&self) -> Height {
25        Height::GENESIS
26    }
27}
28
29/// In-memory block store stub
30pub struct MemoryBlockStore {
31    by_hash: HashMap<BlockHash, Block>,
32    by_height: BTreeMap<u64, BlockHash>,
33}
34
35impl Default for MemoryBlockStore {
36    fn default() -> Self {
37        Self::new()
38    }
39}
40
41impl MemoryBlockStore {
42    pub fn new() -> Self {
43        let mut store = Self {
44            by_hash: HashMap::new(),
45            by_height: BTreeMap::new(),
46        };
47        let genesis = Block::genesis();
48        store.put_block(genesis);
49        store
50    }
51}
52
53impl BlockStore for MemoryBlockStore {
54    fn put_block(&mut self, block: Block) {
55        let hash = block.hash;
56        self.by_height.insert(block.height.as_u64(), hash);
57        self.by_hash.insert(hash, block);
58    }
59
60    fn get_block(&self, hash: &BlockHash) -> Option<Block> {
61        self.by_hash.get(hash).cloned()
62    }
63
64    fn get_block_by_height(&self, h: Height) -> Option<Block> {
65        self.by_height
66            .get(&h.as_u64())
67            .and_then(|hash| self.by_hash.get(hash))
68            .cloned()
69    }
70
71    fn get_blocks_in_range(&self, from: Height, to: Height) -> Vec<Block> {
72        self.by_height
73            .range(from.as_u64()..=to.as_u64())
74            .filter_map(|(_, hash)| self.by_hash.get(hash).cloned())
75            .collect()
76    }
77
78    fn tip_height(&self) -> Height {
79        self.by_height
80            .keys()
81            .next_back()
82            .map(|h| Height(*h))
83            .unwrap_or(Height::GENESIS)
84    }
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90    use hotmint_types::{ValidatorId, ViewNumber};
91
92    fn make_block(height: u64, parent: BlockHash) -> Block {
93        let hash = BlockHash([height as u8; 32]);
94        Block {
95            height: Height(height),
96            parent_hash: parent,
97            view: ViewNumber(height),
98            proposer: ValidatorId(0),
99            payload: vec![],
100            hash,
101        }
102    }
103
104    #[test]
105    fn test_genesis_present() {
106        let store = MemoryBlockStore::new();
107        let genesis = store.get_block(&BlockHash::GENESIS);
108        assert!(genesis.is_some());
109        assert_eq!(genesis.unwrap().height, Height::GENESIS);
110    }
111
112    #[test]
113    fn test_put_and_get_by_hash() {
114        let mut store = MemoryBlockStore::new();
115        let block = make_block(1, BlockHash::GENESIS);
116        let hash = block.hash;
117        store.put_block(block);
118        let retrieved = store.get_block(&hash).unwrap();
119        assert_eq!(retrieved.height, Height(1));
120    }
121
122    #[test]
123    fn test_get_by_height() {
124        let mut store = MemoryBlockStore::new();
125        let b1 = make_block(1, BlockHash::GENESIS);
126        let b2 = make_block(2, b1.hash);
127        store.put_block(b1);
128        store.put_block(b2);
129
130        assert!(store.get_block_by_height(Height(1)).is_some());
131        assert!(store.get_block_by_height(Height(2)).is_some());
132        assert!(store.get_block_by_height(Height(99)).is_none());
133    }
134
135    #[test]
136    fn test_get_nonexistent() {
137        let store = MemoryBlockStore::new();
138        assert!(store.get_block(&BlockHash([99u8; 32])).is_none());
139        assert!(store.get_block_by_height(Height(999)).is_none());
140    }
141
142    #[test]
143    fn test_get_blocks_in_range() {
144        let mut store = MemoryBlockStore::new();
145        let b1 = make_block(1, BlockHash::GENESIS);
146        let b2 = make_block(2, b1.hash);
147        let b3 = make_block(3, b2.hash);
148        store.put_block(b1);
149        store.put_block(b2);
150        store.put_block(b3);
151
152        let blocks = store.get_blocks_in_range(Height(1), Height(3));
153        assert_eq!(blocks.len(), 3);
154        assert_eq!(blocks[0].height, Height(1));
155        assert_eq!(blocks[2].height, Height(3));
156
157        // Partial range
158        let blocks = store.get_blocks_in_range(Height(2), Height(3));
159        assert_eq!(blocks.len(), 2);
160
161        // Out of range
162        let blocks = store.get_blocks_in_range(Height(10), Height(20));
163        assert!(blocks.is_empty());
164    }
165
166    #[test]
167    fn test_tip_height() {
168        let store = MemoryBlockStore::new();
169        assert_eq!(store.tip_height(), Height::GENESIS);
170
171        let mut store = MemoryBlockStore::new();
172        let b1 = make_block(1, BlockHash::GENESIS);
173        let b2 = make_block(2, b1.hash);
174        store.put_block(b1);
175        store.put_block(b2);
176        assert_eq!(store.tip_height(), Height(2));
177    }
178
179    #[test]
180    fn test_overwrite_same_height() {
181        let mut store = MemoryBlockStore::new();
182        let b1 = make_block(1, BlockHash::GENESIS);
183        store.put_block(b1);
184        // Different block at same height (different hash)
185        let mut b2 = make_block(1, BlockHash::GENESIS);
186        b2.hash = BlockHash([42u8; 32]);
187        b2.payload = vec![1, 2, 3];
188        store.put_block(b2);
189        // Height now points to new block
190        let retrieved = store.get_block_by_height(Height(1)).unwrap();
191        assert_eq!(retrieved.hash, BlockHash([42u8; 32]));
192    }
193}