From 2296a667c508afc7ca9ccd4f28740e16633e9a17 Mon Sep 17 00:00:00 2001
From: Jonathan <94441036+zeapoz@users.noreply.github.com>
Date: Tue, 24 Oct 2023 12:14:04 +0200
Subject: [PATCH] ref: include batch number in root hash query (#36)

---
 src/main.rs                      |  6 ++----
 src/processor/tree/query_tree.rs | 30 ++++++++++++++++++++++++++----
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 180df80..ff6c02e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -20,7 +20,7 @@ use std::{
 };
 
 use clap::Parser;
-use cli::{Cli, Command, Query, ReconstructSource};
+use cli::{Cli, Command, ReconstructSource};
 use constants::storage;
 use eyre::Result;
 use snapshot::StateSnapshot;
@@ -131,9 +131,7 @@ async fn main() -> Result<()> {
             };
 
             let tree = QueryTree::new(&db_path);
-            let result = match query {
-                Query::RootHash => tree.latest_root_hash(),
-            };
+            let result = tree.query(&query);
 
             if json {
                 println!("{}", serde_json::to_string(&result)?);
diff --git a/src/processor/tree/query_tree.rs b/src/processor/tree/query_tree.rs
index 340e060..b3a4433 100644
--- a/src/processor/tree/query_tree.rs
+++ b/src/processor/tree/query_tree.rs
@@ -1,8 +1,21 @@
-use std::path::Path;
+use std::{fmt, path::Path};
 
+use serde::Serialize;
 use zksync_merkle_tree::{MerkleTree, RocksDBWrapper};
 
-use super::RootHash;
+use crate::cli::Query;
+
+#[derive(Serialize)]
+pub struct RootHashQuery {
+    pub batch: u64,
+    pub root_hash: String,
+}
+
+impl fmt::Display for RootHashQuery {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "Batch: {}\nRoot Hash: {}", self.batch, self.root_hash)
+    }
+}
 
 pub struct QueryTree<'a>(MerkleTree<'a, RocksDBWrapper>);
 
@@ -16,7 +29,16 @@ impl QueryTree<'static> {
         Self(tree)
     }
 
-    pub fn latest_root_hash(&self) -> RootHash {
-        self.0.latest_root_hash()
+    pub fn query(&self, query: &Query) -> RootHashQuery {
+        match query {
+            Query::RootHash => self.query_root_hash(),
+        }
+    }
+
+    fn query_root_hash(&self) -> RootHashQuery {
+        RootHashQuery {
+            batch: self.0.latest_version().unwrap_or_default(),
+            root_hash: hex::encode(self.0.latest_root_hash()),
+        }
     }
 }