From 634f925d750c7a16b49bb7318fb9d24326a27d96 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <93241502+yhchiang-sol@users.noreply.github.com>
Date: Thu, 8 Jun 2023 12:31:27 -0700
Subject: [PATCH] [TieredStorage] AccountMetaOptionalFields (#32005)

#### Summary of Changes
This PR introduces AccountMetaOptionalFields, the in-memory struct for
handling optional fields.

#### Test Plan
Unit tests are included in this PR.
Tested in mnb w/ the prototype implementation of the tiered storage (#30626).
---
 runtime/src/tiered_storage/meta.rs | 98 +++++++++++++++++++++++++++++-
 1 file changed, 96 insertions(+), 2 deletions(-)

diff --git a/runtime/src/tiered_storage/meta.rs b/runtime/src/tiered_storage/meta.rs
index 1c25756871..124fd1d5fd 100644
--- a/runtime/src/tiered_storage/meta.rs
+++ b/runtime/src/tiered_storage/meta.rs
@@ -1,7 +1,11 @@
 #![allow(dead_code)]
 //! The account meta and related structs for the tiered storage.
 
-use modular_bitfield::prelude::*;
+use {
+    crate::account_storage::meta::StoredMetaWriteVersion,
+    modular_bitfield::prelude::*,
+    solana_sdk::{hash::Hash, stake_history::Epoch},
+};
 
 /// The struct that handles the account meta flags.
 #[bitfield(bits = 32)]
@@ -18,9 +22,45 @@ pub struct AccountMetaFlags {
     reserved: B29,
 }
 
+impl AccountMetaFlags {
+    fn new_from(optional_fields: &AccountMetaOptionalFields) -> Self {
+        let mut flags = AccountMetaFlags::default();
+        flags.set_has_rent_epoch(optional_fields.rent_epoch.is_some());
+        flags.set_has_account_hash(optional_fields.account_hash.is_some());
+        flags.set_has_write_version(optional_fields.write_version.is_some());
+        flags
+    }
+}
+
+/// The in-memory struct for the optional fields for tiered account meta.
+///
+/// Note that the storage representation of the optional fields might be
+/// different from its in-memory representation.
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct AccountMetaOptionalFields {
+    /// the epoch at which its associated account will next owe rent
+    pub rent_epoch: Option<Epoch>,
+    /// the hash of its associated account
+    pub account_hash: Option<Hash>,
+    /// Order of stores of its associated account to an accounts file will
+    /// determine 'latest' account data per pubkey.
+    pub write_version: Option<StoredMetaWriteVersion>,
+}
+
+impl AccountMetaOptionalFields {
+    /// The size of the optional fields in bytes (excluding the boolean flags).
+    pub fn size(&self) -> usize {
+        self.rent_epoch.map_or(0, |_| std::mem::size_of::<Epoch>())
+            + self.account_hash.map_or(0, |_| std::mem::size_of::<Hash>())
+            + self
+                .write_version
+                .map_or(0, |_| std::mem::size_of::<StoredMetaWriteVersion>())
+    }
+}
+
 #[cfg(test)]
 pub mod tests {
-    use crate::tiered_storage::meta::*;
+    use super::*;
 
     #[test]
     fn test_account_meta_flags_new() {
@@ -69,4 +109,58 @@ pub mod tests {
         // make sure the reserved bits are untouched.
         assert_eq!(flags.reserved(), 0u32);
     }
+
+    fn update_and_verify_flags(opt_fields: &AccountMetaOptionalFields) {
+        let flags: AccountMetaFlags = AccountMetaFlags::new_from(opt_fields);
+        assert_eq!(flags.has_rent_epoch(), opt_fields.rent_epoch.is_some());
+        assert_eq!(flags.has_account_hash(), opt_fields.account_hash.is_some());
+        assert_eq!(
+            flags.has_write_version(),
+            opt_fields.write_version.is_some()
+        );
+        assert_eq!(flags.reserved(), 0u32);
+    }
+
+    #[test]
+    fn test_optional_fields_update_flags() {
+        let test_epoch = 5432312;
+        let test_write_version = 231;
+
+        for rent_epoch in [None, Some(test_epoch)] {
+            for account_hash in [None, Some(Hash::new_unique())] {
+                for write_version in [None, Some(test_write_version)] {
+                    update_and_verify_flags(&AccountMetaOptionalFields {
+                        rent_epoch,
+                        account_hash,
+                        write_version,
+                    });
+                }
+            }
+        }
+    }
+
+    #[test]
+    fn test_optional_fields_size() {
+        let test_epoch = 5432312;
+        let test_write_version = 231;
+
+        for rent_epoch in [None, Some(test_epoch)] {
+            for account_hash in [None, Some(Hash::new_unique())] {
+                for write_version in [None, Some(test_write_version)] {
+                    let opt_fields = AccountMetaOptionalFields {
+                        rent_epoch,
+                        account_hash,
+                        write_version,
+                    };
+                    assert_eq!(
+                        opt_fields.size(),
+                        rent_epoch.map_or(0, |_| std::mem::size_of::<Epoch>())
+                            + account_hash.map_or(0, |_| std::mem::size_of::<Hash>())
+                            + write_version
+                                .map_or(0, |_| std::mem::size_of::<StoredMetaWriteVersion>())
+                    );
+                }
+            }
+        }
+    }
 }