snap_coin/core/
block_store.rs

1use std::{
2    collections::HashMap,
3    fs::{self, File},
4    io::Write,
5    sync::RwLock,
6};
7
8use bincode::{Decode, Encode};
9use serde::{Deserialize, Serialize};
10use thiserror::Error;
11
12use crate::{
13    core::{
14        block::{Block, BlockError}, transaction::{Transaction, TransactionId}, utxo::UTXODiff
15    },
16    crypto::Hash,
17    economics::GENESIS_PREVIOUS_BLOCK_HASH,
18};
19
20#[derive(Error, Serialize, Deserialize, Debug, Clone, Encode, Decode)]
21pub enum BlockStoreError {
22    #[error("No blocks left to pop")]
23    NoneToPop,
24
25    #[error("Encoding failed")]
26    Encode,
27
28    #[error("block error: {0}")]
29    BlockError(#[from] BlockError),
30
31    #[error("IO error: {0}")]
32    IO(String),
33}
34
35impl From<std::io::Error> for BlockStoreError {
36    fn from(value: std::io::Error) -> Self {
37        BlockStoreError::IO(value.to_string())
38    }
39}
40
41#[derive(Debug, Encode, Decode, Clone)]
42pub struct BlockIndex {
43    by_hash: HashMap<Hash, usize>,
44    by_height: HashMap<usize, Hash>,
45}
46
47impl BlockIndex {
48    pub fn new_empty() -> Self {
49        Self {
50            by_hash: HashMap::new(),
51            by_height: HashMap::new(),
52        }
53    }
54}
55
56#[derive(Debug, Encode, Decode)]
57pub struct BlockStore {
58    pub store_path: String,
59    block_index: RwLock<BlockIndex>, // RwLock's are justified, because they only get written to on block add or pop
60    height: RwLock<usize>,
61    last_block: RwLock<Hash>,
62}
63
64impl BlockStore {
65    pub fn new_empty(path: &str) -> Self {
66        Self {
67            store_path: path.to_owned(),
68            block_index: RwLock::new(BlockIndex::new_empty()),
69            height: RwLock::new(0usize),
70            last_block: RwLock::new(GENESIS_PREVIOUS_BLOCK_HASH),
71        }
72    }
73
74    /// Adds a block, writing it to disk, updating block lookup and current height
75    /// WARNING: block MUST be valid beforehand!
76    pub fn add_block(&self, block: Block, diffs: UTXODiff) -> Result<(), BlockStoreError> {
77        block.check_completeness()?;
78
79        let block_tmp = format!("{}.tmp", self.block_path_by_height(self.get_height()));
80        let diffs_tmp = format!("{}.tmp", self.utxo_diffs_path_by_height(self.get_height()));
81
82        // Serialize
83        let block_buffer = bincode::encode_to_vec(block.clone(), bincode::config::standard())
84            .map_err(|_| BlockStoreError::Encode)?;
85        let diffs_buffer = bincode::encode_to_vec(diffs, bincode::config::standard())
86            .map_err(|_| BlockStoreError::Encode)?;
87
88        // Write temp block
89        {
90            let mut f = File::create(&block_tmp)?;
91            f.write_all(&block_buffer)?;
92            f.sync_all()?;
93        }
94
95        // Write temp diffs
96        {
97            let mut f = File::create(&diffs_tmp)?;
98            f.write_all(&diffs_buffer)?;
99            f.sync_all()?;
100        }
101
102        // Atomic rename
103        fs::rename(&block_tmp, self.block_path_by_height(self.get_height()))?;
104        fs::rename(
105            &diffs_tmp,
106            self.utxo_diffs_path_by_height(self.get_height()),
107        )?;
108
109        // Update block index
110        {
111            let mut block_index = self.block_index.write().unwrap();
112            block_index
113                .by_hash
114                .insert(block.meta.hash.unwrap(), self.get_height()); // Unwraps are okay, we checked, block is complete
115            block_index
116                .by_height
117                .insert(self.get_height(), block.meta.hash.unwrap());
118        }
119
120        // Update block height and last block
121        *self.height.write().unwrap() = self.get_height() + 1;
122        *self.last_block.write().unwrap() = block.meta.hash.unwrap();
123        Ok(())
124    }
125
126    pub fn pop_block(&self) -> Result<(), BlockStoreError> {
127        if self.get_height() == 0 {
128            return Err(BlockStoreError::NoneToPop);
129        }
130        let height = self.get_height() - 1;
131
132        // Paths
133        let block_path = self.block_path_by_height(height);
134        let diffs_path = self.utxo_diffs_path_by_height(height);
135
136        let block_del = format!("{}.delete", block_path);
137        let diffs_del = format!("{}.delete", diffs_path);
138
139        // Atomically rename files out of the active set
140        fs::rename(&block_path, &block_del)?;
141        fs::rename(&diffs_path, &diffs_del)?;
142
143        // Remove block from index
144        {
145            let mut block_index = self.block_index.write().unwrap();
146            block_index.by_hash.remove(&self.get_last_block_hash());
147            block_index.by_height.remove(&height);
148        }
149
150        // Update height first
151        *self.height.write().unwrap() = height;
152
153        // Update last_block correctly
154        if height == 0 {
155            *self.last_block.write().unwrap() = GENESIS_PREVIOUS_BLOCK_HASH;
156        } else {
157            let last_block_after_pop = self.get_last_block().ok_or(BlockError::IncompleteBlock)?;
158            *self.last_block.write().unwrap() = last_block_after_pop.meta.hash.unwrap();
159        }
160
161        // Final delete (non-atomic, but safe now)
162        fs::remove_file(block_del)?;
163        fs::remove_file(diffs_del)?;
164
165        Ok(())
166    }
167
168    /// Gets current block height (count of all blocks)
169    pub fn get_height(&self) -> usize {
170        *self.height.read().unwrap()
171    }
172
173    /// Gets last added block hash (GENESIS_PREVIOUS_BLOCK_HASH if height = 0)
174    pub fn get_last_block_hash(&self) -> Hash {
175        *self.last_block.read().unwrap()
176    }
177
178    /// Gets last added block
179    pub fn get_last_block(&self) -> Option<Block> {
180        self.get_block_by_height(self.get_height().saturating_sub(1))
181    }
182
183    /// Gets block referenced by it's height
184    pub fn get_block_by_height(&self, height: usize) -> Option<Block> {
185        if height > self.get_height() {
186            return None;
187        }
188
189        let path = self.block_path_by_height(height);
190        Self::load_block_from_path(&path).ok()
191    }
192
193    /// Gets block referenced by it's hash
194    pub fn get_block_by_hash(&self, hash: Hash) -> Option<Block> {
195        self.get_block_by_height(*self.block_index.read().unwrap().by_hash.get(&hash)?)
196    }
197
198    /// Gets block hash referenced by it's height
199    pub fn get_block_hash_by_height(&self, height: usize) -> Option<Hash> {
200        self.block_index
201            .read()
202            .unwrap()
203            .by_height
204            .get(&height)
205            .copied()
206    }
207
208    /// Gets block height referenced by it's hash
209    pub fn get_block_height_by_hash(&self, hash: Hash) -> Option<usize> {
210        self.block_index.read().unwrap().by_hash.get(&hash).copied()
211    }
212
213    pub fn get_last_utxo_diffs(&self) -> Option<UTXODiff> {
214        let path = self.utxo_diffs_path_by_height(self.get_height() - 1);
215        Self::load_utxo_diffs_from_path(&path).ok()
216    }
217
218    /// Gets UTXO diffs by block height
219    pub fn get_utxo_diffs_by_height(&self, height: usize) -> Option<UTXODiff> {
220        if height == 0 || height > self.get_height() {
221            return None;
222        }
223
224        let path = self.utxo_diffs_path_by_height(height);
225        Self::load_utxo_diffs_from_path(&path).ok()
226    }
227
228    /// Gets UTXO diffs by block hash
229    pub fn get_utxo_diffs_by_hash(&self, hash: Hash) -> Option<UTXODiff> {
230        let height = self.get_block_height_by_hash(hash)?;
231        self.get_utxo_diffs_by_height(height)
232    }
233
234    fn load_utxo_diffs_from_path(path: &str) -> Result<UTXODiff, BlockStoreError> {
235        let data = fs::read(path).map_err(|e| BlockStoreError::IO(e.to_string()))?;
236        let (diffs, _) = bincode::decode_from_slice(&data, bincode::config::standard())
237            .map_err(|_| BlockStoreError::Encode)?;
238        Ok(diffs)
239    }
240
241    fn load_block_from_path(path: &str) -> Result<Block, BlockStoreError> {
242        let data = fs::read(path).map_err(|_| BlockStoreError::Encode)?;
243        let (block, _) = bincode::decode_from_slice(&data, bincode::config::standard())
244            .map_err(|_| BlockStoreError::Encode)?;
245        Ok(block)
246    }
247
248    fn block_path_by_height(&self, height: usize) -> String {
249        format!("{}{}.dat", self.store_path, height)
250    }
251    fn utxo_diffs_path_by_height(&self, height: usize) -> String {
252        format!("{}utxo-diffs-{}.dat", self.store_path, height)
253    }
254
255    pub fn get_transaction(&self, tx_id: TransactionId) -> Option<Transaction> {
256        // Iterate from newest to oldest for faster access if likely recent
257        let height = self.get_height();
258        for h in (0..height).rev() {
259            if let Some(block) = self.get_block_by_height(h) {
260                for tx in block.transactions {
261                    if tx.transaction_id == Some(tx_id) {
262                        return Some(tx);
263                    }
264                }
265            }
266        }
267        None
268    }
269
270    pub fn iter_blocks(&self) -> impl Iterator<Item = Result<Block, BlockStoreError>> + '_ {
271        let height = self.get_height();
272        (0..height).map(move |h| {
273            let path = self.block_path_by_height(h);
274            Self::load_block_from_path(&path)
275        })
276    }
277}
278
279impl Clone for BlockStore {
280    /// WARNING: SLOW
281    fn clone(&self) -> Self {
282        Self {
283            store_path: self.store_path.clone(),
284            block_index: RwLock::new(self.block_index.read().unwrap().clone()),
285            height: RwLock::new(*self.height.read().unwrap()),
286            last_block: RwLock::new(*self.last_block.read().unwrap()),
287        }
288    }
289}