feat: accept hash or number in state query

Author: SupernaviX

common/src/queries/blocks.rs

@@ -9,15 +9,15 @@ pub enum BlocksStateQuery {
     GetLatestBlockTransactions,
     GetLatestBlockTransactionsCBOR,
     GetBlockInfo {
-        block_key: Vec<u8>,
+        block_key: BlockKey,
     },
     GetNextBlocks {
-        block_key: Vec<u8>,
+        block_key: BlockKey,
         limit: u64,
         skip: u64,
     },
     GetPreviousBlocks {
-        block_key: Vec<u8>,
+        block_key: BlockKey,
         limit: u64,
         skip: u64,
     },
@@ -29,16 +29,22 @@ pub enum BlocksStateQuery {
         slot: u64,
     },
     GetBlockTransactions {
-        block_key: Vec<u8>,
+        block_key: BlockKey,
     },
     GetBlockTransactionsCBOR {
-        block_key: Vec<u8>,
+        block_key: BlockKey,
     },
     GetBlockInvolvedAddresses {
-        block_key: Vec<u8>,
+        block_key: BlockKey,
     },
 }
 
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub enum BlockKey {
+    Hash(Vec<u8>),
+    Number(u64),
+}
+
 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
 pub enum BlocksStateQueryResponse {
     LatestBlock(BlockInfo),

modules/chain_store/src/chain_store.rs

@@ -4,8 +4,9 @@ use acropolis_common::{
     crypto::keyhash,
     messages::{CardanoMessage, Message, StateQuery, StateQueryResponse},
     queries::blocks::{
-        BlockInfo, BlockTransaction, BlockTransactions, BlockTransactionsCBOR, BlocksStateQuery,
-        BlocksStateQueryResponse, NextBlocks, PreviousBlocks, DEFAULT_BLOCKS_QUERY_TOPIC,
+        BlockInfo, BlockKey, BlockTransaction, BlockTransactions, BlockTransactionsCBOR,
+        BlocksStateQuery, BlocksStateQueryResponse, NextBlocks, PreviousBlocks,
+        DEFAULT_BLOCKS_QUERY_TOPIC,
     },
     TxHash,
 };
@@ -107,7 +108,7 @@ impl ChainStore {
                 Ok(BlocksStateQueryResponse::LatestBlockTransactionsCBOR(txs))
             }
             BlocksStateQuery::GetBlockInfo { block_key } => {
-                let Some(block) = store.get_block_by_hash(block_key)? else {
+                let Some(block) = Self::get_block_by_key(store, block_key)? else {
                     return Ok(BlocksStateQueryResponse::NotFound);
                 };
                 let info = Self::to_block_info(block, store, false)?;
@@ -137,7 +138,7 @@ impl ChainStore {
                         blocks: vec![],
                     }));
                 }
-                let Some(block) = store.get_block_by_hash(&block_key)? else {
+                let Some(block) = Self::get_block_by_key(store, &block_key)? else {
                     return Ok(BlocksStateQueryResponse::NotFound);
                 };
                 let number = Self::get_block_number(&block)?;
@@ -159,7 +160,7 @@ impl ChainStore {
                         blocks: vec![],
                     }));
                 }
-                let Some(block) = store.get_block_by_hash(&block_key)? else {
+                let Some(block) = Self::get_block_by_key(store, &block_key)? else {
                     return Ok(BlocksStateQueryResponse::NotFound);
                 };
                 let number = Self::get_block_number(&block)?;
@@ -176,14 +177,14 @@ impl ChainStore {
                 }))
             }
             BlocksStateQuery::GetBlockTransactions { block_key } => {
-                let Some(block) = store.get_block_by_hash(block_key)? else {
+                let Some(block) = Self::get_block_by_key(store, block_key)? else {
                     return Ok(BlocksStateQueryResponse::NotFound);
                 };
                 let txs = Self::to_block_transactions(block)?;
                 Ok(BlocksStateQueryResponse::BlockTransactions(txs))
             }
             BlocksStateQuery::GetBlockTransactionsCBOR { block_key } => {
-                let Some(block) = store.get_block_by_hash(block_key)? else {
+                let Some(block) = Self::get_block_by_key(store, block_key)? else {
                     return Ok(BlocksStateQueryResponse::NotFound);
                 };
                 let txs = Self::to_block_transactions_cbor(block)?;
@@ -194,6 +195,13 @@ impl ChainStore {
         }
     }
 
+    fn get_block_by_key(store: &Arc<dyn Store>, block_key: &BlockKey) -> Result<Option<Block>> {
+        match block_key {
+            BlockKey::Hash(hash) => store.get_block_by_hash(hash),
+            BlockKey::Number(number) => store.get_block_by_number(*number),
+        }
+    }
+
     fn get_block_number(block: &Block) -> Result<u64> {
         Ok(pallas_traverse::MultiEraBlock::decode(&block.bytes)?.number())
     }