refactor: use ordered publisher list

program/rust/src/accounts/price.rs

@@ -1,4 +1,6 @@
 pub use price_pythnet::*;
+#[cfg(test)]
+use quickcheck::Arbitrary;
 use {
     super::{
         AccountHeader,
@@ -188,14 +190,29 @@ mod price_pythnet {
 }
 
 #[repr(C)]
+#[cfg_attr(test, derive(Debug, PartialEq))]
 #[derive(Copy, Clone, Pod, Zeroable)]
 pub struct PriceComponent {
     pub pub_:    Pubkey,
     pub agg_:    PriceInfo,
     pub latest_: PriceInfo,
 }
 
+#[cfg(test)]
+impl Arbitrary for PriceComponent {
+    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
+        let mut key = [0u8; 32];
+        key.iter_mut().for_each(|item| *item = u8::arbitrary(g));
+        PriceComponent {
+            pub_:    Pubkey::new_from_array(key),
+            agg_:    PriceInfo::arbitrary(g),
+            latest_: PriceInfo::arbitrary(g),
+        }
+    }
+}
+
 #[repr(C)]
+#[cfg_attr(test, derive(Debug, PartialEq))]
 #[derive(Copy, Clone, Pod, Zeroable)]
 pub struct PriceInfo {
     pub price_:           i64,
@@ -205,6 +222,19 @@ pub struct PriceInfo {
     pub pub_slot_:        u64,
 }
 
+#[cfg(test)]
+impl Arbitrary for PriceInfo {
+    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
+        PriceInfo {
+            price_:           i64::arbitrary(g),
+            conf_:            u64::arbitrary(g),
+            status_:          u32::arbitrary(g),
+            corp_act_status_: u32::arbitrary(g),
+            pub_slot_:        u64::arbitrary(g),
+        }
+    }
+}
+
 #[repr(C)]
 #[derive(Copy, Clone, Pod, Zeroable)]
 pub struct PriceEma {

program/rust/src/processor/add_publisher.rs

@@ -24,7 +24,10 @@ use {
         account_info::AccountInfo,
         entrypoint::ProgramResult,
         program_error::ProgramError,
-        program_memory::sol_memset,
+        program_memory::{
+            sol_memcmp,
+            sol_memset,
+        },
         pubkey::Pubkey,
     },
     std::mem::size_of,
@@ -41,8 +44,7 @@ pub fn add_publisher(
     let cmd_args = load::<AddPublisherArgs>(instruction_data)?;
 
     pyth_assert(
-        instruction_data.len() == size_of::<AddPublisherArgs>()
-            && cmd_args.publisher != Pubkey::default(),
+        instruction_data.len() == size_of::<AddPublisherArgs>(),
         ProgramError::InvalidArgument,
     )?;
 
@@ -63,6 +65,14 @@ pub fn add_publisher(
 
     let mut price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
 
+    // Use the call with the default pubkey (000..) as a trigger to sort the publishers as a
+    // migration step from unsorted list to sorted list.
+    if cmd_args.publisher == Pubkey::default() {
+        let num_comps = try_convert::<u32, usize>(price_data.num_)?;
+        sort_price_comps(&mut price_data.comp_, num_comps)?;
+        return Ok(());
+    }
+
     if price_data.num_ >= PC_NUM_COMP {
         return Err(ProgramError::InvalidArgument);
     }
@@ -81,6 +91,141 @@ pub fn add_publisher(
     );
     price_data.comp_[current_index].pub_ = cmd_args.publisher;
     price_data.num_ += 1;
+
+    // Sort the publishers in the list
+    {
+        let num_comps = try_convert::<u32, usize>(price_data.num_)?;
+        sort_price_comps(&mut price_data.comp_, num_comps)?;
+    }
+
     price_data.header.size = try_convert::<_, u32>(PriceAccount::INITIAL_SIZE)?;
     Ok(())
 }
+
+/// A copy of rust slice/sort.rs heapsort implementation which is small and fast. We couldn't use
+/// the sort directly because it was only accessible behind a unstable feature flag at the time of
+/// writing this code.
+fn heapsort(v: &mut [(Pubkey, usize)]) {
+    // This binary heap respects the invariant `parent >= child`.
+    let sift_down = |v: &mut [(Pubkey, usize)], mut node: usize| {
+        loop {
+            // Children of `node`.
+            let mut child = 2 * node + 1;
+            if child >= v.len() {
+                break;
+            }
+
+            // Choose the greater child.
+            if child + 1 < v.len()
+                && sol_memcmp(v[child].0.as_ref(), v[child + 1].0.as_ref(), 32) < 0
+            {
+                child += 1;
+            }
+
+            // Stop if the invariant holds at `node`.
+            if sol_memcmp(v[node].0.as_ref(), v[child].0.as_ref(), 32) >= 0 {
+                break;
+            }
+
+            // Swap `node` with the greater child, move one step down, and continue sifting.
+            v.swap(node, child);
+            node = child;
+        }
+    };
+
+    // Build the heap in linear time.
+    for i in (0..v.len() / 2).rev() {
+        sift_down(v, i);
+    }
+
+    // Pop maximal elements from the heap.
+    for i in (1..v.len()).rev() {
+        v.swap(0, i);
+        sift_down(&mut v[..i], 0);
+    }
+}
+
+/// Sort the publishers price component list in place by performing minimal swaps.
+/// This code is inspired by the sort_by_cached_key implementation in the Rust stdlib.
+/// The rust stdlib implementation is not used because it uses a fast sort variant that has
+/// a large code size.
+///
+/// num_publishers is the number of publishers in the list that should be sorted. It is explicitly
+/// passed to avoid callers mistake of passing the full slice which may contain uninitialized values.
+fn sort_price_comps(comps: &mut [PriceComponent], num_comps: usize) -> Result<(), ProgramError> {
+    let comps = comps
+        .get_mut(..num_comps)
+        .ok_or(ProgramError::InvalidArgument)?;
+
+    // Publishers are likely sorted in ascending order but
+    // heapsorts creates a max-heap so we reverse the order
+    // of the keys to make the heapify step faster.
+    let mut keys = comps
+        .iter()
+        .enumerate()
+        .map(|(i, x)| (x.pub_, i))
+        .rev()
+        .collect::<Vec<_>>();
+
+    heapsort(&mut keys);
+
+    for i in 0..num_comps {
+        // We know that the publisher with key[i].0 should be at index i in the sorted array and
+        // want to swap them in-place in O(n). Normally, the publisher at key[i].0 should be at
+        // key[i].1 but if it is swapped, we need to find the correct index by following the chain
+        // of swaps.
+        let mut index = keys[i].1;
+
+        while index < i {
+            index = keys[index].1;
+        }
+        // Setting the final index here is important to make the code linear as we won't
+        // loop over from i to index again when we reach i again.
+        keys[i].1 = index;
+        comps.swap(i, index);
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod test {
+    use {
+        super::*,
+        quickcheck_macros::quickcheck,
+    };
+
+    #[quickcheck]
+    pub fn test_sort_price_comps(mut comps: Vec<PriceComponent>) {
+        let mut rust_std_sorted_comps = comps.clone();
+        rust_std_sorted_comps.sort_by_key(|x| x.pub_);
+
+        let num_comps = comps.len();
+        assert_eq!(
+            sort_price_comps(&mut comps, num_comps + 1),
+            Err(ProgramError::InvalidArgument)
+        );
+
+        assert_eq!(sort_price_comps(&mut comps, num_comps), Ok(()));
+        assert_eq!(comps, rust_std_sorted_comps);
+    }
+
+    #[quickcheck]
+    pub fn test_sort_price_comps_smaller_slice(
+        mut comps: Vec<PriceComponent>,
+        mut num_comps: usize,
+    ) {
+        num_comps = if comps.is_empty() {
+            0
+        } else {
+            num_comps % comps.len()
+        };
+
+        let mut rust_std_sorted_comps = comps.get(..num_comps).unwrap().to_vec();
+        rust_std_sorted_comps.sort_by_key(|x| x.pub_);
+
+
+        assert_eq!(sort_price_comps(&mut comps, num_comps), Ok(()));
+        assert_eq!(comps.get(..num_comps).unwrap(), rust_std_sorted_comps);
+    }
+}

program/rust/src/processor/upd_price.rs

@@ -2,6 +2,7 @@ use {
     crate::{
         accounts::{
             PriceAccount,
+            PriceComponent,
             PriceInfo,
             PythOracleSerialize,
             UPD_PRICE_WRITE_SEED,
@@ -31,6 +32,7 @@ use {
         },
         program::invoke_signed,
         program_error::ProgramError,
+        program_memory::sol_memcmp,
         pubkey::Pubkey,
         sysvar::Sysvar,
     },
@@ -127,7 +129,7 @@ pub fn upd_price(
     // Check clock
     let clock = Clock::from_account_info(clock_account)?;
 
-    let mut publisher_index: usize = 0;
+    let publisher_index: usize;
     let latest_aggregate_price: PriceInfo;
 
     // The price_data borrow happens in a scope because it must be
@@ -137,17 +139,15 @@ pub fn upd_price(
         // Verify that symbol account is initialized
         let price_data = load_checked::<PriceAccount>(price_account, cmd_args.header.version)?;
 
-        // Verify that publisher is authorized
-        while publisher_index < try_convert::<u32, usize>(price_data.num_)? {
-            if price_data.comp_[publisher_index].pub_ == *funding_account.key {
-                break;
+        publisher_index = match find_publisher_index(
+            &price_data.comp_[..try_convert::<u32, usize>(price_data.num_)?],
+            funding_account.key,
+        ) {
+            Some(index) => index,
+            None => {
+                return Err(OracleError::PermissionViolation.into());
             }
-            publisher_index += 1;
-        }
-        pyth_assert(
-            publisher_index < try_convert::<u32, usize>(price_data.num_)?,
-            OracleError::PermissionViolation.into(),
-        )?;
+        };
 
         latest_aggregate_price = price_data.agg_;
         let latest_publisher_price = price_data.comp_[publisher_index].latest_;
@@ -281,6 +281,64 @@ pub fn upd_price(
     Ok(())
 }
 
+/// Find the index of the publisher in the list of components.
+///
+/// This method first tries to binary search for the publisher's key in the list of components
+/// to get the result faster if the list is sorted. If the list is not sorted, it falls back to
+/// a linear search.
+#[inline(always)]
+fn find_publisher_index(comps: &[PriceComponent], key: &Pubkey) -> Option<usize> {
+    // Verify that publisher is authorized by initially binary searching
+    // for the publisher's component in the price account. The binary
+    // search might not work if the publisher list is not sorted; therefore
+    // we fall back to a linear search.
+    let mut binary_search_result = None;
+
+    {
+        // Binary search to find the publisher key. Rust std binary search is not used because
+        // they guarantee valid outcome only if the array is sorted whereas we want to rely on
+        // a Equal match if it is a result on an unsorted array. Currently the rust
+        // implementation behaves the same but we do not want to rely on api internals.
+        let mut left = 0;
+        let mut right = comps.len();
+        while left < right {
+            let mid = left + (right - left) / 2;
+            // sol_memcmp is much faster than rust default comparison of Pubkey. It costs
+            // 10CU whereas rust default comparison costs a few times more.
+            match sol_memcmp(comps[mid].pub_.as_ref(), key.as_ref(), 32) {
+                i if i < 0 => {
+                    left = mid + 1;
+                }
+                i if i > 0 => {
+                    right = mid;
+                }
+                _ => {
+                    binary_search_result = Some(mid);
+                    break;
+                }
+            }
+        }
+    }
+
+    match binary_search_result {
+        Some(index) => Some(index),
+        None => {
+            let mut index = 0;
+            while index < comps.len() {
+                if sol_memcmp(comps[index].pub_.as_ref(), key.as_ref(), 32) == 0 {
+                    break;
+                }
+                index += 1;
+            }
+            if index == comps.len() {
+                None
+            } else {
+                Some(index)
+            }
+        }
+    }
+}
+
 #[allow(dead_code)]
 // Wrapper struct for the accounts required to add data to the accumulator program.
 struct MessageBufferAccounts<'a, 'b: 'a> {
@@ -289,3 +347,45 @@ struct MessageBufferAccounts<'a, 'b: 'a> {
     oracle_auth_pda:     &'a AccountInfo<'b>,
     message_buffer_data: &'a AccountInfo<'b>,
 }
+
+#[cfg(test)]
+mod test {
+    use {
+        super::*,
+        crate::accounts::PriceComponent,
+        quickcheck_macros::quickcheck,
+        solana_program::pubkey::Pubkey,
+    };
+
+    /// Test the find_publisher_index method works with an unordered list of components.
+    #[quickcheck]
+    pub fn test_find_publisher_index_unordered_comp(comps: Vec<PriceComponent>) {
+        comps.iter().enumerate().for_each(|(idx, comp)| {
+            assert_eq!(find_publisher_index(&comps, &comp.pub_), Some(idx));
+        });
+
+        let mut key_not_in_list = Pubkey::new_unique();
+        while comps.iter().any(|comp| comp.pub_ == key_not_in_list) {
+            key_not_in_list = Pubkey::new_unique();
+        }
+
+        assert_eq!(find_publisher_index(&comps, &key_not_in_list), None);
+    }
+
+    /// Test the find_publisher_index method works with a sorted list of components.
+    #[quickcheck]
+    pub fn test_find_publisher_index_ordered_comp(mut comps: Vec<PriceComponent>) {
+        comps.sort_by_key(|comp| comp.pub_);
+
+        comps.iter().enumerate().for_each(|(idx, comp)| {
+            assert_eq!(find_publisher_index(&comps, &comp.pub_), Some(idx));
+        });
+
+        let mut key_not_in_list = Pubkey::new_unique();
+        while comps.iter().any(|comp| comp.pub_ == key_not_in_list) {
+            key_not_in_list = Pubkey::new_unique();
+        }
+
+        assert_eq!(find_publisher_index(&comps, &key_not_in_list), None);
+    }
+}