Use u16 for deduplication map (#9132)

* Use u16 for deduplication map

* Do not limit account number for tests

program-runtime/src/invoke_context.rs

@@ -297,9 +297,7 @@ impl<'a, 'ix_data> InvokeContext<'a, 'ix_data> {
         signers: &[Pubkey],
     ) -> Result<(), InstructionError> {
         // We reference accounts by an u8 index, so we have a total of 256 accounts.
-        // This algorithm allocates the array on the stack for speed.
-        // On AArch64 in release mode, this function only consumes 640 bytes of stack.
-        let mut transaction_callee_map: Vec<u8> = vec![u8::MAX; MAX_ACCOUNTS_PER_TRANSACTION];
+        let mut transaction_callee_map: Vec<u16> = vec![u16::MAX; MAX_ACCOUNTS_PER_TRANSACTION];
         let mut instruction_accounts: Vec<InstructionAccount> =
             Vec::with_capacity(instruction.accounts.len());
 
@@ -308,7 +306,6 @@ impl<'a, 'ix_data> InvokeContext<'a, 'ix_data> {
         // function, we must borrow it again as mutable.
         let program_account_index = {
             let instruction_context = self.transaction_context.get_current_instruction_context()?;
-            debug_assert!(instruction.accounts.len() <= transaction_callee_map.len());
 
             for account_meta in instruction.accounts.iter() {
                 let index_in_transaction = self
@@ -343,7 +340,7 @@ impl<'a, 'ix_data> InvokeContext<'a, 'ix_data> {
                     };
                     instruction_accounts.push(cloned_account);
                 } else {
-                    *index_in_callee = instruction_accounts.len() as u8;
+                    *index_in_callee = instruction_accounts.len() as u16;
                     instruction_accounts.push(InstructionAccount::new(
                         index_in_transaction,
                         account_meta.is_signer,
@@ -445,11 +442,7 @@ impl<'a, 'ix_data> InvokeContext<'a, 'ix_data> {
         data: &'ix_data [u8],
     ) -> Result<(), InstructionError> {
         // We reference accounts by an u8 index, so we have a total of 256 accounts.
-        // This algorithm allocates the array on the stack for speed.
-        // On AArch64 in release mode, this function only consumes 464 bytes of stack (when it is
-        // not inlined).
-        let mut transaction_callee_map: Vec<u8> = vec![u8::MAX; MAX_ACCOUNTS_PER_TRANSACTION];
-        debug_assert!(instruction.accounts.len() <= transaction_callee_map.len());
+        let mut transaction_callee_map: Vec<u16> = vec![u16::MAX; MAX_ACCOUNTS_PER_TRANSACTION];
 
         let mut instruction_accounts: Vec<InstructionAccount> =
             Vec::with_capacity(instruction.accounts.len());
@@ -461,7 +454,7 @@ impl<'a, 'ix_data> InvokeContext<'a, 'ix_data> {
                 .unwrap();
 
             if (*index_in_callee as usize) > instruction_accounts.len() {
-                *index_in_callee = instruction_accounts.len() as u8;
+                *index_in_callee = instruction_accounts.len() as u16;
             }
 
             let index_in_transaction = *index_in_transaction as usize;
@@ -1393,9 +1386,10 @@ mod tests {
 
     #[test]
     fn test_prepare_instruction_maximum_accounts() {
+        const MAX_ACCOUNTS_REFERENCED: usize = u16::MAX as usize;
         let mut transaction_accounts: Vec<KeyedAccountSharedData> =
             Vec::with_capacity(MAX_ACCOUNTS_PER_TRANSACTION);
-        let mut account_metas: Vec<AccountMeta> = Vec::with_capacity(MAX_ACCOUNTS_PER_INSTRUCTION);
+        let mut account_metas: Vec<AccountMeta> = Vec::with_capacity(MAX_ACCOUNTS_REFERENCED);
 
         // Fee-payer
         let fee_payer = Keypair::new();
@@ -1411,10 +1405,20 @@ mod tests {
         transaction_accounts.push((program_id, program_account));
         account_metas.push(AccountMeta::new_readonly(program_id, false));
 
-        for _ in 2..MAX_ACCOUNTS_PER_INSTRUCTION {
-            let key = Pubkey::new_unique();
-            transaction_accounts.push((key, AccountSharedData::new(1, 1, &Pubkey::new_unique())));
-            account_metas.push(AccountMeta::new_readonly(key, false));
+        for i in 2..MAX_ACCOUNTS_REFERENCED {
+            // Let's reference 256 unique accounts, and the rest is repeated.
+            if i < MAX_ACCOUNTS_PER_TRANSACTION {
+                let key = Pubkey::new_unique();
+                transaction_accounts
+                    .push((key, AccountSharedData::new(1, 1, &Pubkey::new_unique())));
+                account_metas.push(AccountMeta::new_readonly(key, false));
+            } else {
+                let repeated_key = transaction_accounts
+                    .get(i % MAX_ACCOUNTS_PER_TRANSACTION)
+                    .unwrap()
+                    .0;
+                account_metas.push(AccountMeta::new_readonly(repeated_key, false));
+            }
         }
 
         with_mock_invoke_context!(invoke_context, transaction_context, transaction_accounts);
@@ -1441,15 +1445,26 @@ mod tests {
                 .transaction_context
                 .get_next_instruction_context()
                 .unwrap();
-            for index_in_transaction in 0..MAX_ACCOUNTS_PER_INSTRUCTION as IndexOfAccount {
-                let index_in_instruction = instruction_context
-                    .get_index_of_account_in_instruction(index_in_transaction as IndexOfAccount)
-                    .unwrap();
-                let other_transaction = instruction_context
+            for index_in_instruction in 0..MAX_ACCOUNTS_REFERENCED as IndexOfAccount {
+                let index_in_transaction = instruction_context
                     .get_index_of_instruction_account_in_transaction(index_in_instruction)
                     .unwrap();
-                assert_eq!(index_in_transaction, other_transaction);
-                assert_eq!(index_in_transaction, index_in_instruction);
+                let other_ix_index = instruction_context
+                    .get_index_of_account_in_instruction(index_in_transaction)
+                    .unwrap();
+                if (index_in_instruction as usize) < MAX_ACCOUNTS_PER_TRANSACTION {
+                    assert_eq!(index_in_instruction, index_in_transaction);
+                    assert_eq!(index_in_instruction, other_ix_index);
+                } else {
+                    assert_eq!(
+                        index_in_instruction as usize % MAX_ACCOUNTS_PER_TRANSACTION,
+                        index_in_transaction as usize
+                    );
+                    assert_eq!(
+                        index_in_instruction as usize % MAX_ACCOUNTS_PER_TRANSACTION,
+                        other_ix_index as usize
+                    );
+                }
             }
         }
 
@@ -1458,20 +1473,29 @@ mod tests {
                 .transaction_context
                 .get_next_instruction_context()
                 .unwrap();
-            for index_in_transaction in 0..MAX_ACCOUNTS_PER_INSTRUCTION as IndexOfAccount {
-                let index_in_instruction = instruction_context
-                    .get_index_of_account_in_instruction(index_in_transaction as IndexOfAccount)
-                    .unwrap();
-                let other_transaction = instruction_context
+            for index_in_instruction in 0..MAX_ACCOUNTS_REFERENCED as IndexOfAccount {
+                let index_in_transaction = instruction_context
                     .get_index_of_instruction_account_in_transaction(index_in_instruction)
                     .unwrap();
+                let other_ix_index = instruction_context
+                    .get_index_of_account_in_instruction(index_in_transaction)
+                    .unwrap();
                 assert_eq!(
-                    index_in_instruction,
-                    (MAX_ACCOUNTS_PER_INSTRUCTION as IndexOfAccount)
-                        .saturating_sub(index_in_transaction)
+                    index_in_transaction,
+                    (MAX_ACCOUNTS_REFERENCED as u16)
+                        .saturating_sub(index_in_instruction)
                         .saturating_sub(1)
+                        .overflowing_rem(MAX_ACCOUNTS_PER_TRANSACTION as u16)
+                        .0
                 );
-                assert_eq!(index_in_transaction, other_transaction);
+                if (index_in_instruction as usize) < MAX_ACCOUNTS_PER_TRANSACTION {
+                    assert_eq!(index_in_instruction, other_ix_index);
+                } else {
+                    assert_eq!(
+                        index_in_instruction as usize % MAX_ACCOUNTS_PER_TRANSACTION,
+                        other_ix_index as usize
+                    );
+                }
             }
         }
 

program-runtime/src/serialization.rs

@@ -786,7 +786,7 @@ mod tests {
                     // Special case implementation of configure_next_instruction_for_tests()
                     // which avoids the overflow when constructing the dedup_map
                     // by simply not filling it.
-                    let dedup_map = vec![u8::MAX; MAX_ACCOUNTS_PER_TRANSACTION];
+                    let dedup_map = vec![u16::MAX; MAX_ACCOUNTS_PER_TRANSACTION];
                     invoke_context
                         .transaction_context
                         .configure_next_instruction(

transaction-context/src/instruction.rs

@@ -18,7 +18,7 @@ pub struct InstructionFrame<'ix_data> {
     /// This is an account deduplication map that maps index_in_transaction to index_in_instruction
     /// Usage: dedup_map[index_in_transaction] = index_in_instruction
     /// This is a vector of u8s to save memory, since many entries may be unused.
-    pub(crate) dedup_map: Vec<u8>,
+    pub(crate) dedup_map: Vec<u16>,
     pub instruction_data: Cow<'ix_data, [u8]>,
 }
 
@@ -31,7 +31,7 @@ pub struct InstructionContext<'a, 'ix_data> {
     pub(crate) nesting_level: usize,
     pub(crate) program_account_index_in_tx: IndexOfAccount,
     pub(crate) instruction_accounts: &'a [InstructionAccount],
-    pub(crate) dedup_map: &'a [u8],
+    pub(crate) dedup_map: &'a [u16],
     pub(crate) instruction_data: &'ix_data [u8],
 }
 

transaction-context/src/lib.rs

@@ -244,7 +244,7 @@ impl<'ix_data> TransactionContext<'ix_data> {
         &mut self,
         program_index: IndexOfAccount,
         instruction_accounts: Vec<InstructionAccount>,
-        deduplication_map: Vec<u8>,
+        deduplication_map: Vec<u16>,
         instruction_data: Cow<'ix_data, [u8]>,
     ) -> Result<(), InstructionError> {
         debug_assert_eq!(deduplication_map.len(), MAX_ACCOUNTS_PER_TRANSACTION);
@@ -266,14 +266,14 @@ impl<'ix_data> TransactionContext<'ix_data> {
         instruction_accounts: Vec<InstructionAccount>,
         instruction_data: Vec<u8>,
     ) -> Result<(), InstructionError> {
-        debug_assert!(instruction_accounts.len() <= u8::MAX as usize);
-        let mut dedup_map = vec![u8::MAX; MAX_ACCOUNTS_PER_TRANSACTION];
+        debug_assert!(instruction_accounts.len() <= u16::MAX as usize);
+        let mut dedup_map = vec![u16::MAX; MAX_ACCOUNTS_PER_TRANSACTION];
         for (idx, account) in instruction_accounts.iter().enumerate() {
             let index_in_instruction = dedup_map
                 .get_mut(account.index_in_transaction as usize)
                 .unwrap();
-            if *index_in_instruction == u8::MAX {
-                *index_in_instruction = idx as u8;
+            if *index_in_instruction == u16::MAX {
+                *index_in_instruction = idx as u16;
             }
         }
         self.configure_next_instruction(