Merge 003a298cd7 into 81c8c556e8

CHANGELOG.md

@@ -13,6 +13,7 @@ The minor version will be incremented upon a breaking change and the patch versi
 ### Features
 
 - avm: Support customizing the installation location using `AVM_HOME` environment variable ([#2917](https://github.com/coral-xyz/anchor/pull/2917))
+- idl, ts: Add accounts resolution for associated token accounts ([#2927](https://github.com/coral-xyz/anchor/pull/2927))
 
 ### Fixes
 

lang/derive/space/src/lib.rs

@@ -4,8 +4,8 @@ use proc_macro::TokenStream;
 use proc_macro2::{Ident, TokenStream as TokenStream2, TokenTree};
 use quote::{quote, quote_spanned, ToTokens};
 use syn::{
-    parse::ParseStream, parse2, parse_macro_input, Attribute, DeriveInput, Fields, GenericArgument,
-    LitInt, PathArguments, Type, TypeArray,
+    parse::ParseStream, parse2, parse_macro_input, spanned::Spanned, DeriveInput, Field, Fields,
+    GenericArgument, LitInt, PathArguments, Type, TypeArray,
 };
 
 /// Implements a [`Space`](./trait.Space.html) trait on the given
@@ -13,6 +13,7 @@ use syn::{
 ///
 /// For types that have a variable size like String and Vec, it is necessary to indicate the size by the `max_len` attribute.
 /// For nested types, it is necessary to specify a size for each variable type (see example).
+/// A raw space can be specified with the `space` attribute.
 ///
 /// # Example
 /// ```ignore
@@ -24,6 +25,8 @@ use syn::{
 ///     pub string_one: String,
 ///     #[max_len(10, 5)]
 ///     pub nested: Vec<Vec<u8>>,
+///     #[space(200)]
+///     pub string_two: String,
 /// }
 ///
 /// #[derive(Accounts)]
@@ -35,7 +38,7 @@ use syn::{
 ///    pub data: Account<'info, ExampleAccount>,
 /// }
 /// ```
-#[proc_macro_derive(InitSpace, attributes(max_len))]
+#[proc_macro_derive(InitSpace, attributes(max_len, space))]
 pub fn derive_init_space(item: TokenStream) -> TokenStream {
     let input = parse_macro_input!(item as DeriveInput);
     let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
@@ -44,10 +47,7 @@ pub fn derive_init_space(item: TokenStream) -> TokenStream {
     let expanded: TokenStream2 = match input.data {
         syn::Data::Struct(strct) => match strct.fields {
             Fields::Named(named) => {
-                let recurse = named.named.into_iter().map(|f| {
-                    let mut max_len_args = get_max_len_args(&f.attrs);
-                    len_from_type(f.ty, &mut max_len_args)
-                });
+                let recurse = named.named.into_iter().map(get_len);
 
                 quote! {
                     #[automatically_derived]
@@ -60,10 +60,7 @@ pub fn derive_init_space(item: TokenStream) -> TokenStream {
         },
         syn::Data::Enum(enm) => {
             let variants = enm.variants.into_iter().map(|v| {
-                let len = v.fields.into_iter().map(|f| {
-                    let mut max_len_args = get_max_len_args(&f.attrs);
-                    len_from_type(f.ty, &mut max_len_args)
-                });
+                let len = v.fields.into_iter().map(get_len);
 
                 quote! {
                     0 #(+ #len)*
@@ -174,21 +171,45 @@ fn parse_len_arg(item: ParseStream) -> Result<VecDeque<TokenStream2>, syn::Error
     Ok(result)
 }
 
-fn get_max_len_args(attributes: &[Attribute]) -> Option<VecDeque<TokenStream2>> {
-    attributes
-        .iter()
-        .find(|a| a.path.is_ident("max_len"))
-        .and_then(|a| a.parse_args_with(parse_len_arg).ok())
+fn get_next_arg(ident: &Ident, args: &mut Option<VecDeque<TokenStream2>>) -> TokenStream2 {
+    let Some(arg_list) = args else {
+        return quote_spanned!(ident.span() => compile_error!("The number of lengths are invalid (Please use `#[max_len(your_len)]`)."));
+    };
+
+    let Some(arg) = arg_list.pop_back() else {
+        return quote_spanned!(ident.span() => compile_error!("Expected max_len attribute."));
+    };
+
+    quote!(#arg)
 }
 
-fn get_next_arg(ident: &Ident, args: &mut Option<VecDeque<TokenStream2>>) -> TokenStream2 {
-    if let Some(arg_list) = args {
-        if let Some(arg) = arg_list.pop_back() {
-            quote!(#arg)
-        } else {
-            quote_spanned!(ident.span() => compile_error!("The number of lengths are invalid."))
-        }
+fn get_len(field: Field) -> TokenStream2 {
+    let mut len_attrs = field
+        .attrs
+        .iter()
+        .filter(|a| a.path.is_ident("max_len") || a.path.is_ident("space"));
+
+    let Some(len_attr) = len_attrs.next() else {
+        return len_from_type(field.ty, &mut None);
+    };
+
+    if len_attrs.next().is_some() {
+        return quote_spanned!(field.span() => compile_error!("Cannot have space and max_len or duplicate at the same time."));
+    }
+
+    let mut args = len_attr.parse_args_with(parse_len_arg).ok();
+
+    if len_attr.path.is_ident("max_len") {
+        len_from_type(field.ty, &mut args)
     } else {
-        quote_spanned!(ident.span() => compile_error!("Expected max_len attribute."))
+        let Some(mut space_arg) = args else {
+            return quote_spanned!(field.span() => compile_error!("Can't parse the arg len token."));
+        };
+
+        if space_arg.len() != 1 {
+            return quote_spanned!(field.span() => compile_error!("Need to have only one space specified."));
+        }
+
+        space_arg.pop_front().unwrap()
     }
 }

lang/syn/src/idl/accounts.rs

@@ -3,7 +3,7 @@ use proc_macro2::TokenStream;
 use quote::{quote, ToTokens};
 
 use super::common::{get_idl_module_path, get_no_docs};
-use crate::{AccountField, AccountsStruct, Field, Ty};
+use crate::{AccountField, AccountsStruct, Field, InitKind, Ty};
 
 /// Generate the IDL build impl for the Accounts struct.
 pub fn gen_idl_build_impl_accounts_struct(accounts: &AccountsStruct) -> TokenStream {
@@ -168,26 +168,87 @@ fn get_address(acc: &Field) -> TokenStream {
 
 fn get_pda(acc: &Field, accounts: &AccountsStruct) -> TokenStream {
     let idl = get_idl_module_path();
+    let parse_default = |expr: &syn::Expr| parse_seed(expr, accounts);
+
+    // Seeds
     let seed_constraints = acc.constraints.seeds.as_ref();
-    let seeds = seed_constraints
-        .map(|seed| seed.seeds.iter().map(|seed| parse_seed(seed, accounts)))
-        .and_then(|seeds| seeds.collect::<Result<Vec<_>>>().ok());
-    let program = seed_constraints
-        .and_then(|seed| seed.program_seed.as_ref())
-        .and_then(|program| parse_seed(program, accounts).ok())
-        .map(|program| quote! { Some(#program) })
-        .unwrap_or_else(|| quote! { None });
-    match seeds {
-        Some(seeds) => quote! {
-            Some(
-                #idl::IdlPda {
-                    seeds: vec![#(#seeds),*],
-                    program: #program,
-                }
-            )
-        },
-        _ => quote! { None },
+    let pda = seed_constraints
+        .map(|seed| seed.seeds.iter().map(parse_default))
+        .and_then(|seeds| seeds.collect::<Result<Vec<_>>>().ok())
+        .map(|seeds| {
+            let program = seed_constraints
+                .and_then(|seed| seed.program_seed.as_ref())
+                .and_then(|program| parse_default(program).ok())
+                .map(|program| quote! { Some(#program) })
+                .unwrap_or_else(|| quote! { None });
+
+            quote! {
+                Some(
+                    #idl::IdlPda {
+                        seeds: vec![#(#seeds),*],
+                        program: #program,
+                    }
+                )
+            }
+        });
+    if let Some(pda) = pda {
+        return pda;
     }
+
+    // Associated token
+    let pda = acc
+        .constraints
+        .init
+        .as_ref()
+        .and_then(|init| match &init.kind {
+            InitKind::AssociatedToken {
+                owner,
+                mint,
+                token_program,
+            } => Some((owner, mint, token_program)),
+            _ => None,
+        })
+        .or_else(|| {
+            acc.constraints
+                .associated_token
+                .as_ref()
+                .map(|ata| (&ata.wallet, &ata.mint, &ata.token_program))
+        })
+        .and_then(|(wallet, mint, token_program)| {
+            // ATA constraints have implicit `.key()` call
+            let parse_expr = |ts| parse_default(&syn::parse2(ts).unwrap()).ok();
+            let parse_ata = |expr| parse_expr(quote! { #expr.key().as_ref() });
+
+            let wallet = parse_ata(wallet);
+            let mint = parse_ata(mint);
+            let token_program = token_program
+                .as_ref()
+                .and_then(parse_ata)
+                .or_else(|| parse_expr(quote!(anchor_spl::token::ID)));
+
+            let seeds = match (wallet, mint, token_program) {
+                (Some(w), Some(m), Some(tp)) => quote! { vec![#w, #tp, #m] },
+                _ => return None,
+            };
+
+            let program = parse_expr(quote!(anchor_spl::associated_token::ID))
+                .map(|program| quote! { Some(#program) })
+                .unwrap();
+
+            Some(quote! {
+                Some(
+                    #idl::IdlPda {
+                        seeds: #seeds,
+                        program: #program,
+                    }
+                )
+            })
+        });
+    if let Some(pda) = pda {
+        return pda;
+    }
+
+    quote! { None }
 }
 
 /// Parse a seeds constraint, extracting the `IdlSeed` types.

lang/tests/space.rs

@@ -97,6 +97,12 @@ pub struct TestConst {
     pub test_array: [u8; MAX_LEN as usize],
 }
 
+#[derive(InitSpace)]
+pub struct RawSpace {
+    #[space(100)]
+    pub test_string: Vec<u8>,
+}
+
 #[test]
 fn test_empty_struct() {
     assert_eq!(TestEmptyAccount::INIT_SPACE, 0);
@@ -147,3 +153,8 @@ fn test_full_path() {
 fn test_const() {
     assert_eq!(TestConst::INIT_SPACE, (4 + 10) + 10)
 }
+
+#[test]
+fn test_raw_space() {
+    assert_eq!(RawSpace::INIT_SPACE, 100)
+}

tests/pda-derivation/programs/pda-derivation/Cargo.toml

@@ -14,7 +14,8 @@ no-entrypoint = []
 no-idl = []
 cpi = ["no-entrypoint"]
 default = []
-idl-build = ["anchor-lang/idl-build"]
+idl-build = ["anchor-lang/idl-build", "anchor-spl/idl-build"]
 
 [dependencies]
 anchor-lang = { path = "../../../../lang" }
+anchor-spl = { path = "../../../../spl" }

tests/pda-derivation/programs/pda-derivation/src/lib.rs

@@ -4,6 +4,10 @@
 mod other;
 
 use anchor_lang::prelude::*;
+use anchor_spl::{
+    associated_token::AssociatedToken,
+    token::{Mint, Token, TokenAccount},
+};
 
 declare_id!("Fg6PaFpoGXkYsidMpWTK6W2BeZ7FEfcYkg476zPFsLnS");
 
@@ -34,6 +38,10 @@ pub mod pda_derivation {
         ctx.accounts.account.data = 1337;
         Ok(())
     }
+
+    pub fn associated_token_resolution(_ctx: Context<AssociatedTokenResolution>) -> Result<()> {
+        Ok(())
+    }
 }
 
 #[derive(Accounts)]
@@ -115,6 +123,29 @@ pub struct Nested<'info> {
     account_nested: AccountInfo<'info>,
 }
 
+#[derive(Accounts)]
+pub struct AssociatedTokenResolution<'info> {
+    #[account(
+        init,
+        payer = payer,
+        mint::authority = payer,
+        mint::decimals = 9,
+    )]
+    pub mint: Account<'info, Mint>,
+    #[account(
+        init,
+        payer = payer,
+        associated_token::authority = payer,
+        associated_token::mint = mint,
+    )]
+    pub ata: Account<'info, TokenAccount>,
+    #[account(mut)]
+    pub payer: Signer<'info>,
+    pub system_program: Program<'info, System>,
+    pub token_program: Program<'info, Token>,
+    pub associated_token_program: Program<'info, AssociatedToken>,
+}
+
 #[account]
 pub struct MyAccount {
     data: u64,

tests/pda-derivation/tests/typescript.spec.ts

@@ -103,4 +103,13 @@ describe("typescript", () => {
 
     expect(called).is.true;
   });
+
+  it("Can resolve associated token accounts", async () => {
+    const mintKp = anchor.web3.Keypair.generate();
+    await program.methods
+      .associatedTokenResolution()
+      .accounts({ mint: mintKp.publicKey })
+      .signers([mintKp])
+      .rpc();
+  });
 });