Simplistic trait

program/rust/src/bindings.h

@@ -11,4 +11,8 @@ typedef unsigned int uint32_t;
 typedef signed long int int64_t;
 typedef unsigned long int uint64_t;
 
+#include <stddef.h>
 #include "../../c/src/oracle/oracle.h"
+
+const size_t PC_PRICE_T_COMP_OFFSET = offsetof(struct pc_price, comp_);
+const size_t PC_MAP_TABLE_T_PROD_OFFSET = offsetof(struct pc_map_table, prod_);

program/rust/src/c_oracle_header.rs

@@ -7,11 +7,35 @@ use bytemuck::{
     Pod,
     Zeroable,
 };
-
+use std::mem::size_of;
 //bindings.rs is generated by build.rs to include
 //things defined in bindings.h
 include!("../bindings.rs");
 
+/// The PythAccount trait's purpose is to attach constants to the 3 types of accounts that Pyth has
+/// (mapping, price, product). This allows less duplicated code, because now we can create generic
+/// functions to perform common checks on the accounts and to load and initialize the accounts.
+pub trait PythAccount: Pod {
+    const ACCOUNT_TYPE: u32;
+    const INITIAL_SIZE: u32;
+}
+
+impl PythAccount for pc_map_table_t {
+    const ACCOUNT_TYPE: u32 = PC_ACCTYPE_MAPPING;
+    const INITIAL_SIZE: u32 = PC_MAP_TABLE_T_PROD_OFFSET as u32;
+}
+
+impl PythAccount for pc_prod_t {
+    const ACCOUNT_TYPE: u32 = PC_ACCTYPE_PRODUCT;
+    const INITIAL_SIZE: u32 = size_of::<pc_prod_t>() as u32;
+}
+
+impl PythAccount for pc_price_t {
+    const ACCOUNT_TYPE: u32 = PC_ACCTYPE_PRICE;
+    const INITIAL_SIZE: u32 = PC_PRICE_T_COMP_OFFSET as u32;
+}
+
+
 #[cfg(target_endian = "little")]
 unsafe impl Zeroable for pc_acc {
 }

program/rust/src/rust_oracle.rs

@@ -32,9 +32,7 @@ use crate::c_oracle_header::{
     pc_price_t,
     pc_prod_t,
     pc_pub_key_t,
-    PC_ACCTYPE_MAPPING,
-    PC_ACCTYPE_PRICE,
-    PC_ACCTYPE_PRODUCT,
+    PythAccount,
     PC_COMP_SIZE,
     PC_MAGIC,
     PC_MAP_TABLE_SIZE,
@@ -95,7 +93,7 @@ pub fn init_mapping(
 
     // Initialize by setting to zero again (just in case) and populating the account header
     let hdr = load::<cmd_hdr_t>(instruction_data)?;
-    initialize_mapping_account(fresh_mapping_account, hdr.ver_)?;
+    initialize_checked::<pc_map_table_t>(fresh_mapping_account, hdr.ver_)?;
 
     Ok(SUCCESS)
 }
@@ -116,13 +114,13 @@ pub fn add_mapping(
     check_valid_fresh_account(next_mapping)?;
 
     let hdr = load::<cmd_hdr_t>(instruction_data)?;
-    let mut cur_mapping = load_mapping_account_mut(cur_mapping, hdr.ver_)?;
+    let mut cur_mapping = load_checked::<pc_map_table_t>(cur_mapping, hdr.ver_)?;
     pyth_assert(
         cur_mapping.num_ == PC_MAP_TABLE_SIZE && pubkey_is_zero(&cur_mapping.next_),
         ProgramError::InvalidArgument,
     )?;
 
-    initialize_mapping_account(next_mapping, hdr.ver_)?;
+    initialize_checked::<pc_map_table_t>(next_mapping, hdr.ver_)?;
     pubkey_assign(&mut cur_mapping.next_, &next_mapping.key.to_bytes());
 
     Ok(SUCCESS)
@@ -156,15 +154,9 @@ pub fn add_price(
     check_valid_signable_account(program_id, price_account, size_of::<pc_price_t>())?;
     check_valid_fresh_account(price_account)?;
 
-    let mut product_data = load_product_account_mut(product_account, cmd_args.ver_)?;
+    let mut product_data = load_checked::<pc_prod_t>(product_account, cmd_args.ver_)?;
 
-    clear_account(price_account)?;
-
-    let mut price_data = load_account_as_mut::<pc_price_t>(price_account)?;
-    price_data.magic_ = PC_MAGIC;
-    price_data.ver_ = cmd_args.ver_;
-    price_data.type_ = PC_ACCTYPE_PRICE;
-    price_data.size_ = (size_of::<pc_price_t>() - size_of_val(&price_data.comp_)) as u32;
+    let mut price_data = initialize_checked::<pc_price_t>(price_account, cmd_args.ver_)?;
     price_data.expo_ = cmd_args.expo_;
     price_data.ptype_ = cmd_args.ptype_;
     pubkey_assign(&mut price_data.prod_, &product_account.key.to_bytes());
@@ -198,7 +190,7 @@ pub fn add_publisher(
     check_valid_funding_account(funding_account)?;
     check_valid_signable_account(program_id, price_account, size_of::<pc_price_t>())?;
 
-    let mut price_data = load_price_account_mut(price_account, cmd_args.ver_)?;
+    let mut price_data = load_checked::<pc_price_t>(price_account, cmd_args.ver_)?;
 
     if price_data.num_ >= PC_COMP_SIZE {
         return Err(ProgramError::InvalidArgument);
@@ -246,14 +238,14 @@ pub fn add_product(
     check_valid_fresh_account(new_product_account)?;
 
     let hdr = load::<cmd_hdr_t>(instruction_data)?;
-    let mut mapping_data = load_mapping_account_mut(tail_mapping_account, hdr.ver_)?;
+    let mut mapping_data = load_checked::<pc_map_table_t>(tail_mapping_account, hdr.ver_)?;
     // The mapping account must have free space to add the product account
     pyth_assert(
         mapping_data.num_ < PC_MAP_TABLE_SIZE,
         ProgramError::InvalidArgument,
     )?;
 
-    initialize_product_account(new_product_account, hdr.ver_)?;
+    initialize_checked::<pc_prod_t>(new_product_account, hdr.ver_)?;
 
     let current_index: usize = try_convert(mapping_data.num_)?;
     pubkey_assign(
@@ -322,94 +314,40 @@ pub fn clear_account(account: &AccountInfo) -> Result<(), ProgramError> {
     Ok(())
 }
 
-
-/// Mutably borrow the data in `account` as a mapping account, validating that the account
-/// is properly formatted. Any mutations to the returned value will be reflected in the
-/// account data. Use this to read already-initialized accounts.
-pub fn load_mapping_account_mut<'a>(
+pub fn load_checked<'a, T: PythAccount>(
     account: &'a AccountInfo,
-    expected_version: u32,
-) -> Result<RefMut<'a, pc_map_table_t>, ProgramError> {
-    let mapping_data = load_account_as_mut::<pc_map_table_t>(account)?;
-
-    pyth_assert(
-        mapping_data.magic_ == PC_MAGIC
-            && mapping_data.ver_ == expected_version
-            && mapping_data.type_ == PC_ACCTYPE_MAPPING,
-        ProgramError::InvalidArgument,
-    )?;
-
-    Ok(mapping_data)
-}
-
-/// Initialize account as a new mapping account. This function will zero out any existing data in
-/// the account.
-pub fn initialize_mapping_account(account: &AccountInfo, version: u32) -> Result<(), ProgramError> {
-    clear_account(account)?;
-
-    let mut mapping_account = load_account_as_mut::<pc_map_table_t>(account)?;
-    mapping_account.magic_ = PC_MAGIC;
-    mapping_account.ver_ = version;
-    mapping_account.type_ = PC_ACCTYPE_MAPPING;
-    mapping_account.size_ =
-        try_convert(size_of::<pc_map_table_t>() - size_of_val(&mapping_account.prod_))?;
-
-    Ok(())
-}
-
-/// Initialize account as a new product account. This function will zero out any existing data in
-/// the account.
-pub fn initialize_product_account(account: &AccountInfo, version: u32) -> Result<(), ProgramError> {
-    clear_account(account)?;
-
-    let mut prod_account = load_account_as_mut::<pc_prod_t>(account)?;
-    prod_account.magic_ = PC_MAGIC;
-    prod_account.ver_ = version;
-    prod_account.type_ = PC_ACCTYPE_PRODUCT;
-    prod_account.size_ = try_convert(size_of::<pc_prod_t>())?;
-
-    Ok(())
-}
-
-/// Mutably borrow the data in `account` as a product account, validating that the account
-/// is properly formatted. Any mutations to the returned value will be reflected in the
-/// account data. Use this to read already-initialized accounts.
-pub fn load_product_account_mut<'a>(
-    account: &'a AccountInfo,
-    expected_version: u32,
-) -> Result<RefMut<'a, pc_prod_t>, ProgramError> {
-    let product_data = load_account_as_mut::<pc_prod_t>(account)?;
-
-    pyth_assert(
-        product_data.magic_ == PC_MAGIC
-            && product_data.ver_ == expected_version
-            && product_data.type_ == PC_ACCTYPE_PRODUCT,
-        ProgramError::InvalidArgument,
-    )?;
+    version: u32,
+) -> Result<RefMut<'a, T>, ProgramError> {
+    {
+        let account_header = load_account_as::<pc_acc>(account)?;
+        pyth_assert(
+            account_header.magic_ == PC_MAGIC
+                && account_header.ver_ == version
+                && account_header.type_ == T::ACCOUNT_TYPE,
+            ProgramError::InvalidArgument,
+        )?;
+    }
 
-    Ok(product_data)
+    load_account_as_mut::<T>(account)
 }
 
-/// Mutably borrow the data in `account` as a price account, validating that the account
-/// is properly formatted. Any mutations to the returned value will be reflected in the
-/// account data. Use this to read already-initialized accounts.
-fn load_price_account_mut<'a>(
+pub fn initialize_checked<'a, T: PythAccount>(
     account: &'a AccountInfo,
-    expected_version: u32,
-) -> Result<RefMut<'a, pc_price_t>, ProgramError> {
-    let price_data = load_account_as_mut::<pc_price_t>(account)?;
+    version: u32,
+) -> Result<RefMut<'a, T>, ProgramError> {
+    clear_account(account)?;
 
-    pyth_assert(
-        price_data.magic_ == PC_MAGIC
-            && price_data.ver_ == expected_version
-            && price_data.type_ == PC_ACCTYPE_PRICE,
-        ProgramError::InvalidArgument,
-    )?;
+    {
+        let mut account_header = load_account_as_mut::<pc_acc>(account)?;
+        account_header.magic_ = PC_MAGIC;
+        account_header.ver_ = version;
+        account_header.type_ = T::ACCOUNT_TYPE;
+        account_header.size_ = T::INITIAL_SIZE;
+    }
 
-    Ok(price_data)
+    load_account_as_mut::<T>(account)
 }
 
-
 // Assign pubkey bytes from source to target, fails if source is not 32 bytes
 pub fn pubkey_assign(target: &mut pc_pub_key_t, source: &[u8]) {
     unsafe { target.k1_.copy_from_slice(source) }

program/rust/src/tests/test_add_mapping.rs

@@ -10,8 +10,8 @@ use crate::deserialize::load_account_as_mut;
 use crate::rust_oracle::{
     add_mapping,
     clear_account,
-    initialize_mapping_account,
-    load_mapping_account_mut,
+    initialize_checked,
+    load_checked,
     pubkey_assign,
     pubkey_equal,
     pubkey_is_zero,
@@ -67,10 +67,11 @@ fn test_add_mapping() {
         Epoch::default(),
     );
 
-    initialize_mapping_account(&cur_mapping, PC_VERSION).unwrap();
+    initialize_checked::<pc_map_table_t>(&cur_mapping, PC_VERSION).unwrap();
 
     {
-        let mut cur_mapping_data = load_mapping_account_mut(&cur_mapping, PC_VERSION).unwrap();
+        let mut cur_mapping_data =
+            load_checked::<pc_map_table_t>(&cur_mapping, PC_VERSION).unwrap();
         cur_mapping_data.num_ = PC_MAP_TABLE_SIZE;
     }
 
@@ -101,8 +102,9 @@ fn test_add_mapping() {
     .is_ok());
 
     {
-        let next_mapping_data = load_mapping_account_mut(&next_mapping, PC_VERSION).unwrap();
-        let mut cur_mapping_data = load_mapping_account_mut(&cur_mapping, PC_VERSION).unwrap();
+        let next_mapping_data = load_checked::<pc_map_table_t>(&next_mapping, PC_VERSION).unwrap();
+        let mut cur_mapping_data =
+            load_checked::<pc_map_table_t>(&cur_mapping, PC_VERSION).unwrap();
 
         assert!(pubkey_equal(
             &cur_mapping_data.next_,
@@ -129,7 +131,8 @@ fn test_add_mapping() {
     );
 
     {
-        let mut cur_mapping_data = load_mapping_account_mut(&cur_mapping, PC_VERSION).unwrap();
+        let mut cur_mapping_data =
+            load_checked::<pc_map_table_t>(&cur_mapping, PC_VERSION).unwrap();
         assert!(pubkey_is_zero(&cur_mapping_data.next_));
         cur_mapping_data.num_ = PC_MAP_TABLE_SIZE;
         cur_mapping_data.magic_ = 0;

program/rust/src/tests/test_add_product.rs

@@ -27,8 +27,8 @@ use crate::deserialize::load_account_as;
 use crate::rust_oracle::{
     add_product,
     clear_account,
-    initialize_mapping_account,
-    load_mapping_account_mut,
+    initialize_checked,
+    load_checked,
     pubkey_equal,
 };
 
@@ -116,7 +116,7 @@ fn test_add_product() {
 
     {
         let product_data = load_account_as::<pc_prod_t>(&product_account).unwrap();
-        let mapping_data = load_mapping_account_mut(&mapping_account, PC_VERSION).unwrap();
+        let mapping_data = load_checked::<pc_map_table_t>(&mapping_account, PC_VERSION).unwrap();
 
         assert_eq!(product_data.magic_, PC_MAGIC);
         assert_eq!(product_data.ver_, PC_VERSION);
@@ -140,7 +140,7 @@ fn test_add_product() {
     )
     .is_ok());
     {
-        let mapping_data = load_mapping_account_mut(&mapping_account, PC_VERSION).unwrap();
+        let mapping_data = load_checked::<pc_map_table_t>(&mapping_account, PC_VERSION).unwrap();
         assert_eq!(mapping_data.num_, 2);
         assert!(pubkey_equal(
             &mapping_data.prod_[1],
@@ -175,7 +175,7 @@ fn test_add_product() {
 
     // test fill up of mapping table
     clear_account(&mapping_account).unwrap();
-    initialize_mapping_account(&mapping_account, PC_VERSION).unwrap();
+    initialize_checked::<pc_map_table_t>(&mapping_account, PC_VERSION).unwrap();
 
     for i in 0..PC_MAP_TABLE_SIZE {
         clear_account(&product_account).unwrap();
@@ -190,7 +190,7 @@ fn test_add_product() {
             instruction_data
         )
         .is_ok());
-        let mapping_data = load_mapping_account_mut(&mapping_account, PC_VERSION).unwrap();
+        let mapping_data = load_checked::<pc_map_table_t>(&mapping_account, PC_VERSION).unwrap();
         assert_eq!(mapping_data.num_, i + 1);
     }
 
@@ -207,6 +207,6 @@ fn test_add_product() {
     )
     .is_err());
 
-    let mapping_data = load_mapping_account_mut(&mapping_account, PC_VERSION).unwrap();
+    let mapping_data = load_checked::<pc_map_table_t>(&mapping_account, PC_VERSION).unwrap();
     assert_eq!(mapping_data.num_, PC_MAP_TABLE_SIZE);
 }