1use std::collections::HashMap;
5use std::sync::Arc;
6
7use sui_sdk_types::{Address, Object, Version};
8use sui_sdk_types::{CheckpointSequenceNumber, EpochId, SignedTransaction, ValidatorCommittee};
9use sui_types::balance_change::BalanceChange;
10use sui_types::base_types::{ObjectID, ObjectType};
11use sui_types::storage::ObjectKey;
12use sui_types::storage::RpcStateReader;
13use sui_types::storage::error::{Error as StorageError, Result};
14use sui_types::storage::{ObjectStore, TransactionInfo};
15use tap::Pipe;
16
17use crate::Direction;
18
19#[derive(Clone)]
20pub struct StateReader {
21 inner: Arc<dyn RpcStateReader>,
22}
23
24impl StateReader {
25 pub fn new(inner: Arc<dyn RpcStateReader>) -> Self {
26 Self { inner }
27 }
28
29 pub fn inner(&self) -> &Arc<dyn RpcStateReader> {
30 &self.inner
31 }
32
33 #[tracing::instrument(skip(self))]
34 pub fn get_object(&self, object_id: Address) -> crate::Result<Option<Object>> {
35 self.inner
36 .get_object(&object_id.into())
37 .map(TryInto::try_into)
38 .transpose()
39 .map_err(Into::into)
40 }
41
42 #[tracing::instrument(skip(self))]
43 pub fn get_object_with_version(
44 &self,
45 object_id: Address,
46 version: Version,
47 ) -> crate::Result<Option<Object>> {
48 self.inner
49 .get_object_by_key(&object_id.into(), version.into())
50 .map(TryInto::try_into)
51 .transpose()
52 .map_err(Into::into)
53 }
54
55 #[tracing::instrument(skip(self))]
56 pub fn get_committee(&self, epoch: EpochId) -> Option<ValidatorCommittee> {
57 self.inner
58 .get_committee(epoch)
59 .map(|committee| (*committee).clone().into())
60 }
61
62 #[tracing::instrument(skip(self))]
63 pub fn get_system_state(&self) -> Result<sui_types::sui_system_state::SuiSystemState> {
64 sui_types::sui_system_state::get_sui_system_state(self.inner())
65 .map_err(StorageError::custom)
66 .map_err(StorageError::custom)
67 }
68
69 #[tracing::instrument(skip(self))]
70 pub fn get_system_state_summary(
71 &self,
72 ) -> Result<sui_types::sui_system_state::sui_system_state_summary::SuiSystemStateSummary> {
73 use sui_types::sui_system_state::SuiSystemStateTrait;
74
75 let system_state = self.get_system_state()?;
76 let summary = system_state.into_sui_system_state_summary();
77
78 Ok(summary)
79 }
80
81 pub fn get_authenticator_state(
82 &self,
83 ) -> Result<Option<sui_types::authenticator_state::AuthenticatorStateInner>> {
84 sui_types::authenticator_state::get_authenticator_state(self.inner())
85 .map_err(StorageError::custom)
86 }
87
88 #[tracing::instrument(skip(self))]
89 pub fn get_transaction(
90 &self,
91 digest: sui_sdk_types::Digest,
92 ) -> crate::Result<(
93 sui_sdk_types::SignedTransaction,
94 sui_types::effects::TransactionEffects,
95 Option<sui_types::effects::TransactionEvents>,
96 )> {
97 use sui_types::effects::TransactionEffectsAPI;
98
99 let transaction_digest = digest.into();
100
101 let transaction = (*self
102 .inner()
103 .get_transaction(&transaction_digest)
104 .ok_or(TransactionNotFoundError(digest))?)
105 .clone()
106 .into_inner();
107 let effects = self
108 .inner()
109 .get_transaction_effects(&transaction_digest)
110 .ok_or(TransactionNotFoundError(digest))?;
111 let events = if effects.events_digest().is_some() {
112 self.inner()
113 .get_events(effects.transaction_digest())
114 .ok_or(TransactionNotFoundError(digest))?
115 .pipe(Some)
116 } else {
117 None
118 };
119
120 Ok((transaction.try_into()?, effects, events))
121 }
122
123 #[tracing::instrument(skip(self))]
124 pub fn get_transaction_info(
125 &self,
126 digest: &sui_types::digests::TransactionDigest,
127 ) -> Option<TransactionInfo> {
128 self.inner()
129 .indexes()?
130 .get_transaction_info(digest)
131 .ok()
132 .flatten()
133 }
134
135 #[tracing::instrument(skip(self))]
136 pub fn get_transaction_read(
137 &self,
138 digest: sui_sdk_types::Digest,
139 ) -> crate::Result<TransactionRead> {
140 let (
141 SignedTransaction {
142 transaction,
143 signatures,
144 },
145 effects,
146 events,
147 ) = self.get_transaction(digest)?;
148
149 let (checkpoint, balance_changes, object_types) =
150 if let Some(info) = self.get_transaction_info(&(digest.into())) {
151 (
152 Some(info.checkpoint),
153 Some(info.balance_changes),
154 Some(info.object_types),
155 )
156 } else {
157 let checkpoint = self.inner().get_transaction_checkpoint(&(digest.into()));
158 (checkpoint, None, None)
159 };
160 let timestamp_ms = if let Some(checkpoint) = checkpoint {
161 self.inner()
162 .get_checkpoint_by_sequence_number(checkpoint)
163 .map(|checkpoint| checkpoint.timestamp_ms)
164 } else {
165 None
166 };
167
168 let unchanged_loaded_runtime_objects = self
169 .inner()
170 .get_unchanged_loaded_runtime_objects(&(digest.into()));
171
172 Ok(TransactionRead {
173 digest: transaction.digest(),
174 transaction,
175 signatures,
176 effects,
177 events,
178 checkpoint,
179 timestamp_ms,
180 balance_changes,
181 object_types,
182 unchanged_loaded_runtime_objects,
183 })
184 }
185
186 #[allow(unused)]
187 pub fn checkpoint_iter(
188 &self,
189 direction: Direction,
190 start: CheckpointSequenceNumber,
191 ) -> CheckpointIter {
192 CheckpointIter::new(self.clone(), direction, start)
193 }
194
195 #[allow(unused)]
196 pub fn transaction_iter(
197 &self,
198 direction: Direction,
199 cursor: (CheckpointSequenceNumber, Option<usize>),
200 ) -> CheckpointTransactionsIter {
201 CheckpointTransactionsIter::new(self.clone(), direction, cursor)
202 }
203}
204
205#[derive(Debug)]
206pub struct TransactionRead {
207 pub digest: sui_sdk_types::Digest,
208 pub transaction: sui_sdk_types::Transaction,
209 pub signatures: Vec<sui_sdk_types::UserSignature>,
210 pub effects: sui_types::effects::TransactionEffects,
211 pub events: Option<sui_types::effects::TransactionEvents>,
212 pub checkpoint: Option<u64>,
213 pub timestamp_ms: Option<u64>,
214 pub balance_changes: Option<Vec<BalanceChange>>,
215 pub object_types: Option<HashMap<ObjectID, ObjectType>>,
216 pub unchanged_loaded_runtime_objects: Option<Vec<ObjectKey>>,
217}
218
219pub struct CheckpointTransactionsIter {
220 reader: StateReader,
221 direction: Direction,
222
223 next_cursor: Option<(CheckpointSequenceNumber, Option<usize>)>,
224 checkpoint: Option<(
225 sui_types::messages_checkpoint::CheckpointSummary,
226 sui_types::messages_checkpoint::CheckpointContents,
227 )>,
228}
229
230impl CheckpointTransactionsIter {
231 #[allow(unused)]
232 pub fn new(
233 reader: StateReader,
234 direction: Direction,
235 start: (CheckpointSequenceNumber, Option<usize>),
236 ) -> Self {
237 Self {
238 reader,
239 direction,
240 next_cursor: Some(start),
241 checkpoint: None,
242 }
243 }
244}
245
246impl Iterator for CheckpointTransactionsIter {
247 type Item = Result<(CursorInfo, sui_types::digests::TransactionDigest)>;
248
249 fn next(&mut self) -> Option<Self::Item> {
250 loop {
251 let (current_checkpoint, transaction_index) = self.next_cursor?;
252
253 let (checkpoint, contents) = if let Some(checkpoint) = &self.checkpoint {
254 if checkpoint.0.sequence_number != current_checkpoint {
255 self.checkpoint = None;
256 continue;
257 } else {
258 checkpoint
259 }
260 } else {
261 let checkpoint = self
262 .reader
263 .inner()
264 .get_checkpoint_by_sequence_number(current_checkpoint)?;
265 let contents = self
266 .reader
267 .inner()
268 .get_checkpoint_contents_by_sequence_number(checkpoint.sequence_number)?;
269
270 self.checkpoint = Some((checkpoint.into_inner().into_data(), contents));
271 self.checkpoint.as_ref().unwrap()
272 };
273
274 let index = transaction_index
275 .map(|idx| idx.clamp(0, contents.size().saturating_sub(1)))
276 .unwrap_or_else(|| match self.direction {
277 Direction::Ascending => 0,
278 Direction::Descending => contents.size().saturating_sub(1),
279 });
280
281 self.next_cursor = {
282 let next_index = match self.direction {
283 Direction::Ascending => {
284 let next_index = index + 1;
285 if next_index >= contents.size() {
286 None
287 } else {
288 Some(next_index)
289 }
290 }
291 Direction::Descending => index.checked_sub(1),
292 };
293
294 let next_checkpoint = if next_index.is_some() {
295 Some(current_checkpoint)
296 } else {
297 match self.direction {
298 Direction::Ascending => current_checkpoint.checked_add(1),
299 Direction::Descending => current_checkpoint.checked_sub(1),
300 }
301 };
302
303 next_checkpoint.map(|checkpoint| (checkpoint, next_index))
304 };
305
306 if contents.size() == 0 {
307 continue;
308 }
309
310 let digest = contents.inner()[index].transaction;
311
312 let cursor_info = CursorInfo {
313 checkpoint: checkpoint.sequence_number,
314 timestamp_ms: checkpoint.timestamp_ms,
315 index: index as u64,
316 next_cursor: self.next_cursor,
317 };
318
319 return Some(Ok((cursor_info, digest)));
320 }
321 }
322}
323
324#[allow(unused)]
325pub struct CursorInfo {
326 pub checkpoint: CheckpointSequenceNumber,
327 pub timestamp_ms: u64,
328 #[allow(unused)]
329 pub index: u64,
330
331 pub next_cursor: Option<(CheckpointSequenceNumber, Option<usize>)>,
333}
334
335pub struct CheckpointIter {
336 reader: StateReader,
337 direction: Direction,
338
339 next_cursor: Option<CheckpointSequenceNumber>,
340}
341
342impl CheckpointIter {
343 #[allow(unused)]
344 pub fn new(reader: StateReader, direction: Direction, start: CheckpointSequenceNumber) -> Self {
345 Self {
346 reader,
347 direction,
348 next_cursor: Some(start),
349 }
350 }
351}
352
353impl Iterator for CheckpointIter {
354 type Item = Result<(
355 sui_types::messages_checkpoint::CertifiedCheckpointSummary,
356 sui_types::messages_checkpoint::CheckpointContents,
357 )>;
358
359 fn next(&mut self) -> Option<Self::Item> {
360 let current_checkpoint = self.next_cursor?;
361
362 let checkpoint = self
363 .reader
364 .inner()
365 .get_checkpoint_by_sequence_number(current_checkpoint)?
366 .into_inner();
367 let contents = self
368 .reader
369 .inner()
370 .get_checkpoint_contents_by_sequence_number(checkpoint.sequence_number)?;
371
372 self.next_cursor = match self.direction {
373 Direction::Ascending => current_checkpoint.checked_add(1),
374 Direction::Descending => current_checkpoint.checked_sub(1),
375 };
376
377 Some(Ok((checkpoint, contents)))
378 }
379}
380
381#[derive(Debug)]
382pub struct TransactionNotFoundError(pub sui_sdk_types::Digest);
383
384impl std::fmt::Display for TransactionNotFoundError {
385 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
386 write!(f, "Transaction {} not found", self.0)
387 }
388}
389
390impl std::error::Error for TransactionNotFoundError {}
391
392impl From<TransactionNotFoundError> for crate::RpcError {
393 fn from(value: TransactionNotFoundError) -> Self {
394 Self::new(tonic::Code::NotFound, value.to_string())
395 }
396}