with_subset() -> get_subset_unchecked_mut()

A simpler, safer, and better documented use of unsafe code

src/runtime.rs

@@ -5,6 +5,8 @@ use solana_sdk::pubkey::Pubkey;
 use solana_sdk::system_program;
 use solana_sdk::transaction::Transaction;
 use solana_system_program;
+use std::collections::HashSet;
+use std::iter::FromIterator;
 
 /// Reasons the runtime might have rejected a transaction.
 #[derive(Debug, PartialEq, Eq, Clone)]
@@ -115,23 +117,32 @@ fn execute_instruction(
     Ok(())
 }
 
-/// Execute a function with a subset of accounts as writable references.
-/// Since the subset can point to the same references, in any order there is no way
-/// for the borrow checker to track them with regards to the original set.
-fn with_subset<F, A>(accounts: &mut [Account], ixes: &[u8], func: F) -> A
-where
-    F: FnOnce(&mut [&mut Account]) -> A,
-{
-    let mut subset: Vec<&mut Account> = ixes
+/// Return true if the slice has any duplicate elements
+fn has_duplicates<T: Eq + std::hash::Hash>(xs: &[T]) -> bool {
+    let xs_set: HashSet<&T> = HashSet::from_iter(xs.iter());
+    xs.len() != xs_set.len()
+}
+
+/// Get mut references to a subset of elements.
+fn get_subset_unchecked_mut<'a, T>(xs: &'a mut [T], indexes: &[u8]) -> Vec<&'a mut T> {
+    // Since the compiler doesn't know the indexes are unique, dereferencing
+    // multiple mut elements is assumed to be unsafe. If, however, all
+    // indexes are unique, it's perfectly safe. The returned elements will share
+    // the liftime of the input slice.
+
+    // Make certain there are no duplicate indexes. If there are, panic because we
+    // can't return multiple mut references to the same element.
+    if has_duplicates(indexes) {
+        panic!("duplicate indexes");
+    }
+
+    indexes
         .iter()
-        .map(|ix| {
-            let ptr = &mut accounts[*ix as usize] as *mut Account;
-            // lifetime of this unsafe is only within the scope of the closure
-            // there is no way to reorder them without breaking borrow checker rules
+        .map(|i| {
+            let ptr = &mut xs[*i as usize] as *mut T;
             unsafe { &mut *ptr }
         })
-        .collect();
-    func(&mut subset)
+        .collect()
 }
 
 /// Execute a transaction.
@@ -145,17 +156,49 @@ pub fn execute_transaction(
 ) -> Result<(), RuntimeError> {
     for (instruction_index, instruction) in tx.instructions.iter().enumerate() {
         let executable_accounts = &mut (&mut loaders[instruction.program_ids_index as usize]);
-        with_subset(tx_accounts, &instruction.accounts, |program_accounts| {
-            execute_instruction(
-                tx,
-                instruction_index,
-                executable_accounts,
-                program_accounts,
-                tick_height,
-            )
-            .map_err(|err| RuntimeError::ProgramError(instruction_index as u8, err))?;
-            Ok(())
-        })?;
+        let mut program_accounts = get_subset_unchecked_mut(tx_accounts, &instruction.accounts);
+        execute_instruction(
+            tx,
+            instruction_index,
+            executable_accounts,
+            &mut program_accounts,
+            tick_height,
+        )
+        .map_err(|err| RuntimeError::ProgramError(instruction_index as u8, err))?;
     }
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_has_duplicates() {
+        assert!(!has_duplicates(&[1, 2]));
+        assert!(has_duplicates(&[1, 2, 1]));
+    }
+
+    #[test]
+    fn test_get_subset_unchecked_mut() {
+        assert_eq!(get_subset_unchecked_mut(&mut [7, 8], &[0]), vec![&mut 7]);
+        assert_eq!(
+            get_subset_unchecked_mut(&mut [7, 8], &[0, 1]),
+            vec![&mut 7, &mut 8]
+        );
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_get_subset_unchecked_mut_duplicate_index() {
+        // This panics, because it assumes duplicate detection is done elsewhere.
+        get_subset_unchecked_mut(&mut [7, 8], &[0, 0]);
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_get_subset_unchecked_mut_out_of_bounds() {
+        // This panics, because it assumes bounds validation is done elsewhere.
+        get_subset_unchecked_mut(&mut [7, 8], &[2]);
+    }
+}