Skip to main content

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