From 97d994ee6f6520a65bcd9ff7641498462c6b8d99 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <93241502+yhchiang-sol@users.noreply.github.com>
Date: Fri, 2 Feb 2024 09:53:13 -0800
Subject: [PATCH] [TieredStorage] Use IndexOffset in TieredStorageMeta and
 get_account() (#35046)

#### Problem
TieredStorageMeta and TieredStorageReader::get_account API uses
u32 to represent IndexOffset.  However, within the TieredStorage scope,
IndexOffset should be used, it is not until working with AccountsFile API
when u32 representation of offset is needed.

#### Summary of Changes
Have TieredStorageMeta and TieredStorageReader to use IndexOffset.

#### Test Plan
Existing unit-tests.
---
 accounts-db/src/account_storage/meta.rs    |  2 +-
 accounts-db/src/tiered_storage/hot.rs      | 10 +++++-----
 accounts-db/src/tiered_storage/readable.rs | 14 +++++++-------
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/accounts-db/src/account_storage/meta.rs b/accounts-db/src/account_storage/meta.rs
index 1442b4845bf604..69c24d7be75f7d 100644
--- a/accounts-db/src/account_storage/meta.rs
+++ b/accounts-db/src/account_storage/meta.rs
@@ -142,7 +142,7 @@ impl<'storage> StoredAccountMeta<'storage> {
     pub fn offset(&self) -> usize {
         match self {
             Self::AppendVec(av) => av.offset(),
-            Self::Hot(hot) => hot.index(),
+            Self::Hot(hot) => hot.index().0 as usize,
         }
     }
 
diff --git a/accounts-db/src/tiered_storage/hot.rs b/accounts-db/src/tiered_storage/hot.rs
index c6e3efdbbb78c6..805b50f2fa3668 100644
--- a/accounts-db/src/tiered_storage/hot.rs
+++ b/accounts-db/src/tiered_storage/hot.rs
@@ -435,7 +435,7 @@ impl HotStorageReader {
     pub fn get_account(
         &self,
         index_offset: IndexOffset,
-    ) -> TieredStorageResult<Option<(StoredAccountMeta<'_>, usize)>> {
+    ) -> TieredStorageResult<Option<(StoredAccountMeta<'_>, IndexOffset)>> {
         if index_offset.0 >= self.footer.account_entry_count {
             return Ok(None);
         }
@@ -452,10 +452,10 @@ impl HotStorageReader {
                 meta,
                 address,
                 owner,
-                index: index_offset.0 as usize,
+                index: index_offset,
                 account_block,
             }),
-            index_offset.0.saturating_add(1) as usize,
+            IndexOffset(index_offset.0.saturating_add(1)),
         )))
     }
 }
@@ -1244,7 +1244,7 @@ pub mod tests {
             );
             assert_eq!(*stored_meta.pubkey(), addresses[i]);
 
-            assert_eq!(i + 1, next);
+            assert_eq!(i + 1, next.0 as usize);
         }
         // Make sure it returns None on NUM_ACCOUNTS to allow termination on
         // while loop in actual accounts-db read case.
@@ -1383,7 +1383,7 @@ pub mod tests {
             let (account, address, account_hash, _write_version) = storable_accounts.get(i);
             verify_account(&stored_meta, account, address, account_hash);
 
-            assert_eq!(i + 1, next);
+            assert_eq!(i + 1, next.0 as usize);
         }
         // Make sure it returns None on NUM_ACCOUNTS to allow termination on
         // while loop in actual accounts-db read case.
diff --git a/accounts-db/src/tiered_storage/readable.rs b/accounts-db/src/tiered_storage/readable.rs
index 647c78d5ca91c1..b6d841b65f4dd8 100644
--- a/accounts-db/src/tiered_storage/readable.rs
+++ b/accounts-db/src/tiered_storage/readable.rs
@@ -25,7 +25,7 @@ pub struct TieredReadableAccount<'accounts_file, M: TieredAccountMeta> {
     /// The address of the account owner
     pub owner: &'accounts_file Pubkey,
     /// The index for accessing the account inside its belonging AccountsFile
-    pub index: usize,
+    pub index: IndexOffset,
     /// The account block that contains this account.  Note that this account
     /// block may be shared with other accounts.
     pub account_block: &'accounts_file [u8],
@@ -43,7 +43,7 @@ impl<'accounts_file, M: TieredAccountMeta> TieredReadableAccount<'accounts_file,
     }
 
     /// Returns the index to this account in its AccountsFile.
-    pub fn index(&self) -> usize {
+    pub fn index(&self) -> IndexOffset {
         self.index
     }
 
@@ -118,10 +118,10 @@ impl TieredStorageReader {
     /// Returns the account located at the specified index offset.
     pub fn get_account(
         &self,
-        index_offset: u32,
-    ) -> TieredStorageResult<Option<(StoredAccountMeta<'_>, usize)>> {
+        index_offset: IndexOffset,
+    ) -> TieredStorageResult<Option<(StoredAccountMeta<'_>, IndexOffset)>> {
         match self {
-            Self::Hot(hot) => hot.get_account(IndexOffset(index_offset)),
+            Self::Hot(hot) => hot.get_account(index_offset),
         }
     }
 
@@ -136,13 +136,13 @@ impl TieredStorageReader {
     /// causes a data overrun.
     pub fn account_matches_owners(
         &self,
-        index_offset: u32,
+        index_offset: IndexOffset,
         owners: &[Pubkey],
     ) -> Result<usize, MatchAccountOwnerError> {
         match self {
             Self::Hot(hot) => {
                 let account_offset = hot
-                    .get_account_offset(IndexOffset(index_offset))
+                    .get_account_offset(index_offset)
                     .map_err(|_| MatchAccountOwnerError::UnableToLoad)?;
                 hot.account_matches_owners(account_offset, owners)
             }