flamenco: refactor txn_ctx/instr_ctx account getters API

Author: cali-jumptrading

src/flamenco/runtime/context/fd_exec_instr_ctx.c

@@ -81,54 +81,88 @@ fd_exec_instr_ctx_delete( void * mem ) {
   return mem;
 }
 
+int
+fd_exec_instr_ctx_find_idx_of_instr_account( fd_exec_instr_ctx_t const * ctx,
+                                             fd_pubkey_t const *         pubkey ) {
+  for( int i = 0; i < ctx->instr->acct_cnt; i++ ) {
+    if( memcmp( pubkey->uc, ctx->instr->acct_pubkeys[i].uc, sizeof(fd_pubkey_t) )==0 ) {
+      return i;
+    }
+  }
+  return -1;
+}
+
 int
 fd_exec_instr_ctx_try_borrow_account( fd_exec_instr_ctx_t const * ctx,
                                       ulong                       idx,
+                                      fd_txn_account_t *          txn_account,
                                       fd_borrowed_account_t *     account ) {
   /* TODO this is slightly wrong. Agave returns NotEnoughAccountKeys when the account index is
      out of bounds in the transaction context and MissingAccount when the account index is out of
      bounds in the instruction context. */
 
-  /* Return a NotEnoughAccountKeys error if the idx is out of bounds.
-     https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L603 */
-  if( FD_UNLIKELY( idx >= ctx->instr->acct_cnt ) ) {
-    return FD_EXECUTOR_INSTR_ERR_NOT_ENOUGH_ACC_KEYS;
-  }
-
-  fd_txn_account_t * instr_account = ctx->instr->accounts[idx];
-
   /* Return an AccountBorrowFailed error if the write is not acquirable.
      https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L605 */
-  int acquire_result = fd_txn_account_acquire_write( instr_account );
+  int acquire_result = fd_txn_account_acquire_write( txn_account );
   if( FD_UNLIKELY( !acquire_result ) ) {
     return FD_EXECUTOR_INSTR_ERR_ACC_BORROW_FAILED;
   }
 
   /* Create a BorrowedAccount upon success.
      https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L606 */
-  fd_borrowed_account_init( account, instr_account, ctx, (int)idx );
+  fd_borrowed_account_init( account, txn_account, ctx, (ushort)idx );
   return FD_EXECUTOR_INSTR_SUCCESS;
 }
 
 int
-fd_exec_instr_ctx_try_borrow_account_with_key( fd_exec_instr_ctx_t *   ctx,
-                                               fd_pubkey_t const *     pubkey,
-                                               fd_borrowed_account_t * account ) {
-  for( ulong i = 0; i < ctx->instr->acct_cnt; i++ ) {
-    if( memcmp( pubkey->uc, ctx->instr->acct_pubkeys[i].uc, sizeof(fd_pubkey_t) )==0 ) {
-      return fd_exec_instr_ctx_try_borrow_account( ctx, i, account );
-    }
+fd_exec_instr_ctx_try_borrow_instr_account( fd_exec_instr_ctx_t const * ctx,
+                                            ulong                       idx,
+                                            fd_borrowed_account_t *     account ) {
+  /* Return a NotEnoughAccountKeys error if the idx is out of bounds.
+     https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L603 */
+  if( FD_UNLIKELY( idx >= ctx->instr->acct_cnt ) ) {
+    return FD_EXECUTOR_INSTR_ERR_NOT_ENOUGH_ACC_KEYS;
   }
-  return FD_EXECUTOR_INSTR_ERR_MISSING_ACC;
+
+  fd_txn_account_t * instr_account = ctx->instr->accounts[idx];
+
+  /* Calculate the index_in_instruction by adding the number of program accounts to the instruction account index.
+     https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L654-L655 */
+  ulong idx_in_instr = fd_ulong_sat_add(idx, fd_exec_instr_ctx_get_number_of_program_accounts( ctx ) );
+
+  return fd_exec_instr_ctx_try_borrow_account( ctx,
+                                               idx_in_instr,
+                                               instr_account,
+                                               account );
 }
 
 int
-fd_exec_instr_ctx_find_idx_of_instr_account( fd_exec_instr_ctx_t const * ctx,
-                                             fd_pubkey_t const *         pubkey ) {
-  for( int i = 0; i < ctx->instr->acct_cnt; i++ ) {
+fd_exec_instr_ctx_try_borrow_instr_account_with_key( fd_exec_instr_ctx_t *   ctx,
+                                                     fd_pubkey_t const *     pubkey,
+                                                     fd_borrowed_account_t * account ) {
+  for( ulong i = 0; i < ctx->instr->acct_cnt; i++ ) {
     if( memcmp( pubkey->uc, ctx->instr->acct_pubkeys[i].uc, sizeof(fd_pubkey_t) )==0 ) {
-      return i;
+      return fd_exec_instr_ctx_try_borrow_instr_account( ctx, i, account );
     }
   }
-  return -1;
+
+  /* Return a NotEnoughAccountKeys error if the account is not found
+     in the instruction context to match the error code returned by
+     fd_exec_instr_ctx_try_borrow_instr_account. */
+  return FD_EXECUTOR_INSTR_ERR_NOT_ENOUGH_ACC_KEYS;
+}
+
+int
+fd_exec_instr_ctx_try_borrow_last_program_account( fd_exec_instr_ctx_t const * ctx,
+                                                   fd_borrowed_account_t * account ) {
+  fd_txn_account_t * program_account = NULL;
+  fd_exec_txn_ctx_get_account_at_index( ctx->txn_ctx,
+                                        ctx->instr->program_id,
+                                        &program_account,
+                                        NULL );
+
+  return fd_exec_instr_ctx_try_borrow_account( ctx,
+                                               0UL,
+                                               program_account,
+                                               account );
 }

src/flamenco/runtime/context/fd_exec_instr_ctx.h

@@ -41,7 +41,7 @@ struct __attribute__((aligned(8UL))) fd_exec_instr_ctx {
 /* Be careful when using this macro. There may be places where the error
    will need to be handled differently. */
 #define FD_TRY_BORROW_INSTR_ACCOUNT_DEFAULT_ERR_CHECK( ctx, idx, acc ) do { \
-  int err = fd_exec_instr_ctx_try_borrow_account( ctx, idx, acc );          \
+  int err = fd_exec_instr_ctx_try_borrow_instr_account( ctx, idx, acc );    \
   if( FD_UNLIKELY( err ) ) return err;                                      \
 } while (0)
 
@@ -80,36 +80,59 @@ fd_exec_instr_ctx_check_num_insn_accounts( fd_exec_instr_ctx_t * ctx,
   return FD_EXECUTOR_INSTR_SUCCESS;
 }
 
+/* Mirrors Agave function solana_sdk::transaction_context::InstructionContext::find_index_of_instruction_account.
+
+   Returns the index of the the instruction account given the account pubkey
+   or -1 if the account is not found.
+
+   https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L524-L538 */
+
+int
+fd_exec_instr_ctx_find_idx_of_instr_account( fd_exec_instr_ctx_t const * ctx,
+                                             fd_pubkey_t const *         pubkey );
+
+/* Mirrors Agave function solana_sdk::transaction_context::InstructionContext::get_number_of_program_accounts.
+
+   Strictly returns 1, as we only support one program account per instruction for now.
+
+   https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L480-L482 */
+
+static inline ushort
+fd_exec_instr_ctx_get_number_of_program_accounts( fd_exec_instr_ctx_t const * ctx ) {
+  (void) ctx;
+  return 1U;
+}
+
 /* Mirrors Agave function solana_sdk::transaction_context::InstructionContext::try_borrow_account.
 
    Borrows an account from the instruction context with a given account index.
 
    https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L594 */
 
 int
-fd_exec_instr_ctx_try_borrow_account( fd_exec_instr_ctx_t const * ctx,
-                                      ulong                       idx,
-                                      fd_borrowed_account_t *     account );
+fd_exec_instr_ctx_try_borrow_instr_account( fd_exec_instr_ctx_t const * ctx,
+                                            ulong                       idx,
+                                            fd_borrowed_account_t *     account );
 
 /* A wrapper around fd_exec_instr_ctx_try_borrow_account that accepts an account pubkey.
 
    Borrows an account from the instruction context with a given pubkey. */
 
 int
-fd_exec_instr_ctx_try_borrow_account_with_key( fd_exec_instr_ctx_t *   ctx,
-                                               fd_pubkey_t const *     pubkey,
-                                               fd_borrowed_account_t * account );
+fd_exec_instr_ctx_try_borrow_instr_account_with_key( fd_exec_instr_ctx_t *   ctx,
+                                                     fd_pubkey_t const *     pubkey,
+                                                     fd_borrowed_account_t * account );
 
-/* Mirrors Agave function solana_sdk::transaction_context::InstructionContext::find_index_of_instruction_account.
+/* Mirrors Agave function solana_sdk::transaction_context::InstructionContext::try_borrow_last_program_account
 
-   Returns the index of the the instruction account given the account pubkey
-   or -1 if the account is not found.
+   Borrows the instruction's program account. Since there is only one program account per
+   instruction, this function simply borrows the instruction's only program account, despite the name.
 
-   https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L524-L538 */
+   https://github.com/anza-xyz/agave/blob/v2.1.14/sdk/src/transaction_context.rs#L616 */
 
 int
-fd_exec_instr_ctx_find_idx_of_instr_account( fd_exec_instr_ctx_t const * ctx,
-                                             fd_pubkey_t const *         pubkey );
+fd_exec_instr_ctx_try_borrow_last_program_account( fd_exec_instr_ctx_t const * ctx,
+                                                   fd_borrowed_account_t * account );
 
 FD_PROTOTYPES_END
 

src/flamenco/runtime/context/fd_exec_txn_ctx.c

@@ -84,121 +84,72 @@ fd_exec_txn_ctx_delete( void * mem ) {
 }
 
 int
-fd_exec_txn_ctx_get_account_view_idx( fd_exec_txn_ctx_t * ctx,
-                                      uchar               idx,
-                                      fd_txn_account_t * * account ) {
+fd_exec_txn_ctx_get_account_at_index( fd_exec_txn_ctx_t *  ctx,
+                                      uchar                idx,
+                                      fd_txn_account_t * * account,
+                                      int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                        int                 idx,
+                                                        fd_txn_account_t *  acc) ) {
   if( FD_UNLIKELY( idx>=ctx->accounts_cnt ) ) {
     return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
   }
 
   fd_txn_account_t * txn_account = &ctx->accounts[idx];
   *account = txn_account;
 
-  if( FD_UNLIKELY( !fd_acc_exists( txn_account->const_meta ) ) ) {
-    return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
+  if( condition != NULL ) {
+    if( FD_UNLIKELY( !condition( ctx, idx, *account ) ) )
+      return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
   }
 
   return FD_ACC_MGR_SUCCESS;
 }
 
 int
-fd_exec_txn_ctx_get_account_view( fd_exec_txn_ctx_t *  ctx,
-                                  fd_pubkey_t const *  pubkey,
-                                  fd_txn_account_t * * account ) {
-  for( ulong i = 0; i < ctx->accounts_cnt; i++ ) {
-    if( memcmp( pubkey->uc, ctx->account_keys[i].uc, sizeof(fd_pubkey_t) )==0 ) {
-      // TODO: check if readable???
-      fd_txn_account_t * txn_account = &ctx->accounts[i];
-      *account = txn_account;
-
-      if( FD_UNLIKELY( !fd_acc_exists( txn_account->const_meta ) ) ) {
-        return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
-      }
-
-      return FD_ACC_MGR_SUCCESS;
-    }
+fd_exec_txn_ctx_get_account_with_key( fd_exec_txn_ctx_t *  ctx,
+                                      fd_pubkey_t const *  pubkey,
+                                      fd_txn_account_t * * account,
+                                      int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                        int                 idx,
+                                                        fd_txn_account_t *  acc) ) {
+  int index = fd_exec_txn_ctx_find_index_of_account( ctx, pubkey );
+  if( FD_UNLIKELY( index==-1 ) ) {
+    return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
   }
 
-  return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
+  return fd_exec_txn_ctx_get_account_at_index( ctx,
+                                               (uchar)index,
+                                               account,
+                                               condition );
 }
 
 int
-fd_exec_txn_ctx_get_account_executable_view( fd_exec_txn_ctx_t *  ctx,
-                                             fd_pubkey_t const *  pubkey,
-                                             fd_txn_account_t * * account ) {
+fd_exec_txn_ctx_get_executable_account( fd_exec_txn_ctx_t *  ctx,
+                                        fd_pubkey_t const *  pubkey,
+                                        fd_txn_account_t * * account,
+                                        int (*condition)( fd_exec_txn_ctx_t * ctx,
+                                                          int                 idx,
+                                                          fd_txn_account_t *  acc) ) {
   /* First try to fetch the executable account from the existing borrowed accounts.
      If the pubkey is in the account keys, then we want to re-use that
      borrowed account since it reflects changes from prior instructions. Referencing the
      read-only executable accounts list is incorrect behavior when the program
      data account is written to in a prior instruction (e.g. program upgrade + invoke within the same txn) */
-  int err = fd_exec_txn_ctx_get_account_view( ctx, pubkey, account );
+  int err = fd_exec_txn_ctx_get_account_with_key( ctx, pubkey, account, condition );
   if( FD_UNLIKELY( err==FD_ACC_MGR_SUCCESS ) ) {
     return FD_ACC_MGR_SUCCESS;
   }
 
   for( ulong i = 0; i < ctx->executable_cnt; i++ ) {
     if( memcmp( pubkey->uc, ctx->executable_accounts[i].pubkey->uc, sizeof(fd_pubkey_t) )==0 ) {
-      // TODO: check if readable???
       fd_txn_account_t * txn_account = &ctx->executable_accounts[i];
       *account = txn_account;
 
-      if( FD_UNLIKELY( !fd_acc_exists( txn_account->const_meta ) ) )
-        return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
-
-      return FD_ACC_MGR_SUCCESS;
-    }
-  }
-
-  return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
-}
-
-int
-fd_exec_txn_ctx_get_account_modify_fee_payer( fd_exec_txn_ctx_t *  ctx,
-                                              fd_txn_account_t * * account ) {
-
-  *account = &ctx->accounts[ FD_FEE_PAYER_TXN_IDX ];
-
-  if( FD_UNLIKELY( !fd_txn_is_writable( ctx->txn_descriptor, FD_FEE_PAYER_TXN_IDX ) ) ) {
-    return FD_ACC_MGR_ERR_WRITE_FAILED;
-  }
-  return FD_ACC_MGR_SUCCESS;
-}
-
-int
-fd_exec_txn_ctx_get_account_modify_idx( fd_exec_txn_ctx_t *   ctx,
-                                        uchar                 idx,
-                                        ulong                 min_data_sz,
-                                        fd_txn_account_t * *  account ) {
-  if( idx >= ctx->accounts_cnt ) {
-    return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
-  }
-
-  fd_txn_account_t * txn_account = &ctx->accounts[idx];
-  if( FD_UNLIKELY( !fd_txn_account_is_writable_idx( ctx, (int)idx ) ) ) {
-    return FD_ACC_MGR_ERR_WRITE_FAILED;
-  }
-
-  if( min_data_sz > txn_account->const_meta->dlen ) {
-    fd_txn_account_resize( txn_account, min_data_sz );
-  }
-
-  *account = txn_account;
-  return FD_ACC_MGR_SUCCESS;
-}
-
-int
-fd_exec_txn_ctx_get_account_modify( fd_exec_txn_ctx_t *  ctx,
-                                    fd_pubkey_t const *  pubkey,
-                                    ulong                min_data_sz,
-                                    fd_txn_account_t * * account ) {
-  for( ulong i = 0; i < ctx->accounts_cnt; i++ ) {
-    if( memcmp( pubkey->uc, ctx->account_keys[i].uc, sizeof(fd_pubkey_t) )==0 ) {
-      // TODO: check if writable???
-      fd_txn_account_t * txn_account = &ctx->accounts[i];
-      if( min_data_sz > txn_account->const_meta->dlen ) {
-        fd_txn_account_resize( txn_account, min_data_sz );
+      if( condition != NULL ) {
+        if( FD_UNLIKELY( !condition( ctx, (int)i, *account ) ) )
+          return FD_ACC_MGR_ERR_UNKNOWN_ACCOUNT;
       }
-      *account = txn_account;
+
       return FD_ACC_MGR_SUCCESS;
     }
   }
@@ -335,7 +286,7 @@ fd_txn_account_is_demotion( fd_exec_txn_ctx_t const * txn_ctx, int idx )
 
    https://github.com/anza-xyz/agave/blob/v2.1.11/sdk/program/src/message/sanitized.rs#L38-L47 */
 int
-fd_txn_account_is_writable_idx( fd_exec_txn_ctx_t const * txn_ctx, int idx ) {
+fd_exec_txn_ctx_account_is_writable_idx( fd_exec_txn_ctx_t const * txn_ctx, int idx ) {
 
   /* https://github.com/anza-xyz/agave/blob/v2.1.11/sdk/program/src/message/sanitized.rs#L43 */
   if( !fd_txn_is_writable( txn_ctx->txn_descriptor, idx ) ) {
@@ -359,3 +310,30 @@ fd_txn_account_is_writable_idx( fd_exec_txn_ctx_t const * txn_ctx, int idx ) {
 
   return 1;
 }
+
+/* Account pre-condition filtering functions */
+
+int
+fd_txn_account_exists( fd_exec_txn_ctx_t * ctx,
+                       int                 idx,
+                       fd_txn_account_t * acc ) {
+  (void) ctx;
+  (void) idx;
+  return fd_acc_exists( acc->const_meta );
+}
+
+int
+fd_txn_account_is_writable( fd_exec_txn_ctx_t * ctx,
+                            int                 idx,
+                            fd_txn_account_t * acc ) {
+  (void) acc;
+  return fd_exec_txn_ctx_account_is_writable_idx( ctx, idx );
+}
+
+int
+fd_txn_account_fee_payer_writable( fd_exec_txn_ctx_t * ctx,
+                                   int                 idx,
+                                   fd_txn_account_t * acc ) {
+  (void) acc;
+  return fd_txn_is_writable( ctx->txn_descriptor, idx );
+}