From 5aebb3228501da2f08726809e61202a30c54f4ce Mon Sep 17 00:00:00 2001
From: Preston Evans <preston.b.evans@gmail.com>
Date: Tue, 25 Apr 2023 10:42:22 -0500
Subject: [PATCH] Add ledger RPC endpoint

---
 src/rpc.rs | 58 +++++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 46 insertions(+), 12 deletions(-)

diff --git a/src/rpc.rs b/src/rpc.rs
index f5c4bb6..2c040b1 100644
--- a/src/rpc.rs
+++ b/src/rpc.rs
@@ -4,11 +4,13 @@ use bank::{
     query::{BankRpcImpl, BankRpcServer},
     Bank,
 };
-use jsonrpsee::server::ServerHandle;
+use jsonrpsee::{server::ServerHandle, RpcModule};
 use sov_modules_api::mocks::MockContext;
 use sov_state::{mocks::MockStorageSpec, ProverStorage, WorkingSet};
 use sovereign_db::ledger_db::LedgerDB;
-use sovereign_sdk::rpc::{LedgerRpcProvider, QueryMode, SlotIdentifier};
+use sovereign_sdk::rpc::{
+    BatchIdentifier, EventIdentifier, LedgerRpcProvider, QueryMode, SlotIdentifier, TxIdentifier,
+};
 
 use crate::{runtime::Runtime, tx_verifier_impl::DemoAppTxVerifier, Spec};
 
@@ -49,22 +51,50 @@ impl RpcProvider {
             _tx_verifier: tx_verifier,
         };
 
-        let ledger_rpc = bank_rpc.clone();
+        // let module = RpcModule::new(ledger_rpc);
 
         let mut bank_rpc = bank_rpc.into_rpc();
-        bank_rpc.register_method("chain_getSlots", move |params, _| {
-            let slot_ids: Vec<SlotIdentifier>;
-            let query_mode: QueryMode;
-            (slot_ids, query_mode) = params.parse()?;
-            ledger_rpc
-                .get_slots(&slot_ids, query_mode)
-                .map_err(|e| e.into())
-        })?;
+
+        register_ledger_rpc(&mut bank_rpc)?;
 
         server.start(bank_rpc)
     }
 }
 
+fn register_ledger_rpc(rpc: &mut RpcModule<RpcProvider>) -> Result<(), jsonrpsee::core::Error> {
+    rpc.register_method("ledger_head", move |_, db| {
+        db.get_head().map_err(|e| e.into())
+    })?;
+
+    rpc.register_method("ledger_getSlots", move |params, db| {
+        let ids: Vec<SlotIdentifier>;
+        let query_mode: QueryMode;
+        (ids, query_mode) = params.parse()?;
+        db.get_slots(&ids, query_mode).map_err(|e| e.into())
+    })?;
+
+    rpc.register_method("ledger_getBatches", move |params, db| {
+        let ids: Vec<BatchIdentifier>;
+        let query_mode: QueryMode;
+        (ids, query_mode) = params.parse()?;
+        db.get_batches(&ids, query_mode).map_err(|e| e.into())
+    })?;
+
+    rpc.register_method("ledger_getTransactions", move |params, db| {
+        let ids: Vec<TxIdentifier>;
+        let query_mode: QueryMode;
+        (ids, query_mode) = params.parse()?;
+        db.get_transactions(&ids, query_mode).map_err(|e| e.into())
+    })?;
+
+    rpc.register_method("ledger_getEvents", move |params, db| {
+        let ids: Vec<EventIdentifier> = params.parse()?;
+        db.get_events(&ids).map_err(|e| e.into())
+    })?;
+
+    Ok(())
+}
+
 // TODO: Re-implement
 // impl TransactionRpcProvider for RpcProvider {
 //     type Transaction = RawTx;
@@ -88,6 +118,10 @@ impl LedgerRpcProvider for RpcProvider {
 
     type EventResponse = <LedgerDB<Spec> as LedgerRpcProvider>::EventResponse;
 
+    fn get_head(&self) -> Result<Option<Self::SlotResponse>, anyhow::Error> {
+        self.ledger_db.get_head()
+    }
+
     fn get_slots(
         &self,
         slot_ids: &[sovereign_sdk::rpc::SlotIdentifier],
@@ -115,7 +149,7 @@ impl LedgerRpcProvider for RpcProvider {
     fn get_events(
         &self,
         event_ids: &[sovereign_sdk::rpc::EventIdentifier],
-    ) -> Result<Option<Vec<Self::EventResponse>>, anyhow::Error> {
+    ) -> Result<Vec<Option<Self::EventResponse>>, anyhow::Error> {
         self.ledger_db.get_events(event_ids)
     }
 }