lang: Allow multiple with targets for associated accounts (#197)

coral-xyz · Apr 15, 2021 · 512604b · 512604b
1 parent 0d87886
commit 512604b
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 31 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,10 @@ incremented for features.
 
 ## [Unreleased]
 
+## Features
+
+* lang: Allows one to specify multiple `with` targets when creating associated acconts ([#197](https://github.com/project-serum/anchor/pull/197)).
+
 ## [0.4.3] - 2021-04-13
 
 ## Features

diff --git a/examples/misc/programs/misc/src/lib.rs b/examples/misc/programs/misc/src/lib.rs
@@ -98,11 +98,12 @@ pub struct TestStateCpi<'info> {
 // the program.
 #[derive(Accounts)]
 pub struct TestAssociatedAccount<'info> {
-    #[account(associated = authority, with = state)]
+    #[account(associated = authority, with = state, with = data)]
     my_account: ProgramAccount<'info, TestData>,
     #[account(mut, signer)]
     authority: AccountInfo<'info>,
     state: ProgramState<'info, MyState>,
+    data: ProgramAccount<'info, Data>,
     rent: Sysvar<'info, Rent>,
     system_program: AccountInfo<'info>,
 }

diff --git a/examples/misc/tests/misc.js b/examples/misc/tests/misc.js
@@ -125,6 +125,7 @@ describe("misc", () => {
         Buffer.from([97, 110, 99, 104, 111, 114]), // b"anchor".
         program.provider.wallet.publicKey.toBuffer(),
         state.toBuffer(),
+        data.publicKey.toBuffer(),
       ],
       program.programId
     );
@@ -145,14 +146,16 @@ describe("misc", () => {
         myAccount: associatedAccount,
         authority: program.provider.wallet.publicKey,
         state,
+        data: data.publicKey,
         rent: anchor.web3.SYSVAR_RENT_PUBKEY,
         systemProgram: anchor.web3.SystemProgram.programId,
       },
     });
     // Try out the generated associated method.
     const account = await program.account.testData.associated(
       program.provider.wallet.publicKey,
-      state
+      state,
+      data.publicKey,
     );
     assert.ok(account.data.toNumber() === 1234);
   });

diff --git a/lang/syn/src/codegen/accounts.rs b/lang/syn/src/codegen/accounts.rs
@@ -540,37 +540,43 @@ pub fn generate_constraint_associated(
         },
     };
 
-    let seeds_no_nonce = match &f.associated_seed {
-        None => quote! {
-            [
-                &b"anchor"[..],
-                #associated_target.to_account_info().key.as_ref(),
-            ]
-        },
-        Some(seed) => quote! {
+    let seeds_no_nonce = match f.associated_seeds.len() {
+        0 => quote! {
             [
                 &b"anchor"[..],
                 #associated_target.to_account_info().key.as_ref(),
-                #seed.to_account_info().key.as_ref(),
             ]
         },
+        _ => {
+            let seeds = to_seeds_tts(&f.associated_seeds);
+            quote! {
+                [
+                    &b"anchor"[..],
+                    #associated_target.to_account_info().key.as_ref(),
+                    #seeds
+                ]
+            }
+        }
     };
-    let seeds_with_nonce = match &f.associated_seed {
-        None => quote! {
-            [
-                &b"anchor"[..],
-                #associated_target.to_account_info().key.as_ref(),
-                &[nonce],
-            ]
-        },
-        Some(seed) => quote! {
+    let seeds_with_nonce = match f.associated_seeds.len() {
+        0 => quote! {
             [
                 &b"anchor"[..],
                 #associated_target.to_account_info().key.as_ref(),
-                #seed.to_account_info().key.as_ref(),
                 &[nonce],
             ]
         },
+        _ => {
+            let seeds = to_seeds_tts(&f.associated_seeds);
+            quote! {
+                [
+                    &b"anchor"[..],
+                    #associated_target.to_account_info().key.as_ref(),
+                    #seeds
+                    &[nonce],
+                ]
+            }
+        }
     };
 
     quote! {
@@ -619,3 +625,19 @@ pub fn generate_constraint_associated(
         };
     }
 }
+
+// Returns the inner part of the seeds slice as a token stream.
+fn to_seeds_tts(seeds: &[syn::Ident]) -> proc_macro2::TokenStream {
+    assert!(seeds.len() > 0);
+    let seed_0 = &seeds[0];
+    let mut tts = quote! {
+        #seed_0.to_account_info().key.as_ref(),
+    };
+    for seed in &seeds[1..] {
+        tts = quote! {
+            #tts
+            #seed.to_account_info().key.as_ref(),
+        };
+    }
+    tts
+}
diff --git a/lang/syn/src/lib.rs b/lang/syn/src/lib.rs
@@ -177,7 +177,7 @@ pub struct Field {
     // Used by the associated attribute only.
     pub space: Option<proc_macro2::TokenStream>,
     // Used by the associated attribute only.
-    pub associated_seed: Option<syn::Ident>,
+    pub associated_seeds: Vec<syn::Ident>,
 }
 
 impl Field {

diff --git a/lang/syn/src/parser/accounts.rs b/lang/syn/src/parser/accounts.rs
@@ -41,8 +41,8 @@ fn parse_account_attr(f: &syn::Field) -> Option<&syn::Attribute> {
 
 fn parse_field(f: &syn::Field, anchor: Option<&syn::Attribute>) -> AccountField {
     let ident = f.ident.clone().unwrap();
-    let (constraints, is_mut, is_signer, is_init, payer, space, associated_seed) = match anchor {
-        None => (vec![], false, false, false, None, None, None),
+    let (constraints, is_mut, is_signer, is_init, payer, space, associated_seeds) = match anchor {
+        None => (vec![], false, false, false, None, None, Vec::new()),
         Some(anchor) => parse_constraints(anchor),
     };
     match is_field_primitive(f) {
@@ -57,7 +57,7 @@ fn parse_field(f: &syn::Field, anchor: Option<&syn::Attribute>) -> AccountField
                 is_init,
                 payer,
                 space,
-                associated_seed,
+                associated_seeds,
             })
         }
         false => AccountField::AccountsStruct(CompositeField {
@@ -186,7 +186,7 @@ fn parse_constraints(
     bool,
     Option<syn::Ident>,
     Option<proc_macro2::TokenStream>,
-    Option<syn::Ident>,
+    Vec<syn::Ident>,
 ) {
     let mut tts = anchor.tokens.clone().into_iter();
     let g_stream = match tts.next().expect("Must have a token group") {
@@ -202,7 +202,7 @@ fn parse_constraints(
     let mut payer = None;
     let mut space = None;
     let mut is_associated = false;
-    let mut associated_seed = None;
+    let mut associated_seeds = Vec::new();
 
     let mut inner_tts = g_stream.into_iter();
     while let Some(token) = inner_tts.next() {
@@ -333,10 +333,10 @@ fn parse_constraints(
                         }
                         _ => panic!("invalid syntax"),
                     };
-                    associated_seed = match inner_tts.next().unwrap() {
-                        proc_macro2::TokenTree::Ident(ident) => Some(ident),
+                    associated_seeds.push(match inner_tts.next().unwrap() {
+                        proc_macro2::TokenTree::Ident(ident) => ident,
                         _ => panic!("invalid syntax"),
-                    };
+                    });
                 }
                 "payer" => {
                     match inner_tts.next().unwrap() {
@@ -408,6 +408,6 @@ fn parse_constraints(
         is_init,
         payer,
         space,
-        associated_seed,
+        associated_seeds,
     )
 }