lang: Add LazyAccount (#3194)

.github/workflows/reusable-tests.yaml

@@ -463,6 +463,8 @@ jobs:
             path: tests/bench
           - cmd: cd tests/idl && ./test.sh
             path: tests/idl
+          - cmd: cd tests/lazy-account && anchor test
+            path: tests/lazy-account
           # TODO: Enable when `solang` becomes compatible with the new IDL spec
           # - cmd: cd tests/solang && anchor test
           #   path: tests/solang

CHANGELOG.md

@@ -45,6 +45,7 @@ The minor version will be incremented upon a breaking change and the patch versi
 - cli: Make `clean` command also remove the `.anchor` directory ([#3192](https://github.com/coral-xyz/anchor/pull/3192)).
 - lang: Deprecate `#[interface]` attribute ([#3195](https://github.com/coral-xyz/anchor/pull/3195)).
 - ts: Include unresolved accounts in the resolution error message ([#3207](https://github.com/coral-xyz/anchor/pull/3207)).
+- lang: Add `LazyAccount` ([#3194](https://github.com/coral-xyz/anchor/pull/3194)).
 
 ### Fixes
 

lang/Cargo.toml

@@ -36,6 +36,7 @@ idl-build = [
 ]
 init-if-needed = ["anchor-derive-accounts/init-if-needed"]
 interface-instructions = ["anchor-attribute-program/interface-instructions"]
+lazy-account = ["anchor-attribute-account/lazy-account", "anchor-derive-serde/lazy-account"]
 
 [dependencies]
 anchor-attribute-access-control = { path = "./attribute/access-control", version = "0.30.1" }

lang/attribute/account/Cargo.toml

@@ -13,6 +13,7 @@ proc-macro = true
 [features]
 anchor-debug = ["anchor-syn/anchor-debug"]
 idl-build = ["anchor-syn/idl-build"]
+lazy-account = []
 
 [dependencies]
 anchor-syn = { path = "../../syn", version = "0.30.1", features = ["hash"] }

lang/attribute/account/src/lazy.rs

@@ -0,0 +1,302 @@
+use proc_macro2::{Literal, TokenStream};
+use quote::{format_ident, quote, ToTokens};
+
+pub fn gen_lazy(strct: &syn::ItemStruct) -> syn::Result<TokenStream> {
+    let ident = &strct.ident;
+    let lazy_ident = format_ident!("Lazy{}", ident);
+    let load_common_ident = to_private_ident("load_common");
+    let initialize_fields = to_private_ident("initialize_fields");
+    let lazy_acc_ty = quote! { anchor_lang::accounts::lazy_account::LazyAccount };
+    let disc_len = quote! { <#ident as anchor_lang::Discriminator>::DISCRIMINATOR.len() };
+
+    let load_common_docs = quote! {
+        /// The deserialized value is cached for future uses i.e. all subsequent calls to this
+        /// method do not deserialize the data again, instead, they return the cached value.
+        ///
+        /// To reload the data from the underlying account info (e.g. after a CPI call), run
+        /// [`LazyAccount::unload`] before running this method.
+        ///
+        /// See [`LazyAccount`]'s documentation for more information.
+    };
+    let load_panic_docs = quote! {
+        /// # Panics
+        ///
+        /// If there is an existing mutable reference crated by any of the `load_mut` methods.
+    };
+    let load_mut_panic_docs = quote! {
+        /// # Panics
+        ///
+        /// If there is an existing reference (mutable or not) created by any of the `load` methods.
+    };
+
+    let (loader_signatures, loader_impls) = strct
+        .fields
+        .iter()
+        .enumerate()
+        .map(|(i, field)| {
+            let field_ident = to_field_ident(field, i);
+            let load_ident = format_ident!("load_{field_ident}");
+            let load_mut_ident = format_ident!("load_mut_{field_ident}");
+            let load_common_ident = to_private_ident(format!("load_common_{field_ident}"));
+            let offset_of_ident = to_private_ident(format!("offset_of_{field_ident}"));
+            let size_of_ident = to_private_ident(format!("size_of_{field_ident}"));
+
+            let offset = i.eq(&0).then(|| quote!(#disc_len)).unwrap_or_else(|| {
+                // Current offset is the previous field's offset + size
+                strct
+                    .fields
+                    .iter()
+                    .nth(i - 1)
+                    .map(|field| {
+                        let field_ident = to_field_ident(field, i - 1);
+                        let offset_of_ident = to_private_ident(format!("offset_of_{field_ident}"));
+                        let size_of_ident = to_private_ident(format!("size_of_{field_ident}"));
+                        quote! { self.#offset_of_ident() + self.#size_of_ident() }
+                    })
+                    .expect("Previous field should always exist when i > 0")
+            });
+
+            let ty = &field.ty;
+            let size = quote! {
+                <#ty as anchor_lang::__private::Lazy>::size_of(
+                    &self.__info.data.borrow()[self.#offset_of_ident()..]
+                )
+            };
+
+            let signatures = quote! {
+                /// Load a reference to the field.
+                ///
+                #load_common_docs
+                ///
+                #load_panic_docs
+                fn #load_ident(&self) -> anchor_lang::Result<::core::cell::Ref<'_, #ty>>;
+
+                /// Load a mutable reference to the field.
+                ///
+                #load_common_docs
+                ///
+                #load_mut_panic_docs
+                fn #load_mut_ident(&self) -> anchor_lang::Result<::core::cell::RefMut<'_, #ty>>;
+
+                 #[doc(hidden)]
+                fn #load_common_ident<R>(&self, f: impl FnOnce() -> R) -> anchor_lang::Result<R>;
+
+                 #[doc(hidden)]
+                fn #offset_of_ident(&self) -> usize;
+
+                 #[doc(hidden)]
+                fn #size_of_ident(&self) -> usize;
+            };
+
+            let impls = quote! {
+                fn #load_ident(&self) -> anchor_lang::Result<::core::cell::Ref<'_, #ty>> {
+                    self.#load_common_ident(|| {
+                        // SAFETY: The common load method makes sure the field is initialized.
+                        ::core::cell::Ref::map(self.__account.borrow(), |acc| unsafe {
+                            &*::core::ptr::addr_of!((*acc.as_ptr()).#field_ident)
+                        })
+                    })
+                }
+
+                fn #load_mut_ident(&self) -> anchor_lang::Result<::core::cell::RefMut<'_, #ty>> {
+                    self.#load_common_ident(|| {
+                        // SAFETY: The common load method makes sure the field is initialized.
+                        ::core::cell::RefMut::map(self.__account.borrow_mut(), |acc| unsafe {
+                            &mut *::core::ptr::addr_of_mut!((*acc.as_mut_ptr()).#field_ident)
+                        })
+                    })
+                }
+
+                #[inline(never)]
+                fn #load_common_ident<R>(&self, f: impl FnOnce() -> R) -> anchor_lang::Result<R> {
+                    self.#initialize_fields();
+
+                    // Return early if initialized
+                    if self.__fields.borrow().as_ref().unwrap()[#i] {
+                        return Ok(f());
+                    }
+
+                    // Deserialize and write
+                    let offset = self.#offset_of_ident();
+                    let size = self.#size_of_ident();
+                    let data = self.__info.data.borrow();
+                    let val = anchor_lang::AnchorDeserialize::try_from_slice(
+                        &data[offset..offset + size]
+                    )?;
+                    unsafe {
+                        ::core::ptr::addr_of_mut!(
+                            (*self.__account.borrow_mut().as_mut_ptr()).#field_ident
+                        ).write(val)
+                     };
+
+                    // Set initialized
+                    self.__fields.borrow_mut().as_mut().unwrap()[#i] = true;
+
+                    Ok(f())
+                }
+
+                // If this method gets inlined when there are >= 12 fields, compilation breaks with
+                // `LLVM ERROR: Branch target out of insn range`
+                #[inline(never)]
+                fn #offset_of_ident(&self) -> usize {
+                    #offset
+                }
+
+                #[inline(always)]
+                fn #size_of_ident(&self) -> usize {
+                    #size
+                }
+            };
+
+            Ok((signatures, impls))
+        })
+        .collect::<syn::Result<Vec<_>>>()?
+        .into_iter()
+        .unzip::<_, _, Vec<_>, Vec<_>>();
+
+    let load_idents = strct
+        .fields
+        .iter()
+        .enumerate()
+        .map(|(i, field)| to_field_ident(field, i))
+        .map(|field| format_ident!("load_{field}"));
+    let total_fields = strct.fields.len();
+
+    Ok(quote! {
+        pub trait #lazy_ident {
+            /// Load a reference to the entire account.
+            ///
+            #load_common_docs
+            ///
+            #load_panic_docs
+            fn load(&self) -> anchor_lang::Result<::core::cell::Ref<'_, #ident>>;
+
+            /// Load a mutable reference to the entire account.
+            ///
+            #load_common_docs
+            ///
+            #load_mut_panic_docs
+            fn load_mut(&self) -> anchor_lang::Result<::core::cell::RefMut<'_, #ident>>;
+
+            #[doc(hidden)]
+            fn #load_common_ident<R>(&self, f: impl FnOnce() -> R) -> anchor_lang::Result<R>;
+
+            #(#loader_signatures)*
+
+            #[doc(hidden)]
+            fn #initialize_fields(&self);
+
+            /// Run the exit routine of the account, similar to [`AccountsExit`] but implemented
+            /// as a regular method because we can't implement external traits for external structs.
+            fn exit(&self, program_id: &anchor_lang::prelude::Pubkey) -> anchor_lang::Result<()>;
+        }
+
+        impl<'info> #lazy_ident for #lazy_acc_ty<'info, #ident> {
+            fn load(&self) -> anchor_lang::Result<::core::cell::Ref<'_, #ident>> {
+                self.#load_common_ident(|| {
+                    // SAFETY: The common load method makes sure all fields are initialized.
+                    ::core::cell::Ref::map(self.__account.borrow(), |acc| unsafe {
+                        acc.assume_init_ref()
+                    })
+                })
+            }
+
+            fn load_mut(&self) -> anchor_lang::Result<::core::cell::RefMut<'_, #ident>> {
+                self.#load_common_ident(|| {
+                    // SAFETY: The common load method makes sure all fields are initialized.
+                    ::core::cell::RefMut::map(self.__account.borrow_mut(), |acc| unsafe {
+                        acc.assume_init_mut()
+                    })
+                })
+            }
+
+            #[inline(never)]
+            fn #load_common_ident<R>(&self, f: impl FnOnce() -> R) -> anchor_lang::Result<R> {
+                self.#initialize_fields();
+
+                // Create a scope to drop the `__fields` borrow
+                let all_uninit = {
+                    // Return early if all fields are initialized
+                    let fields = self.__fields.borrow();
+                    let fields = fields.as_ref().unwrap();
+                    if !fields.contains(&false) {
+                        return Ok(f());
+                    }
+
+                    !fields.contains(&true)
+                };
+
+                if all_uninit {
+                    // Nothing is initialized, initialize all
+                    let offset = #disc_len;
+                    let mut data = self.__info.data.borrow();
+                    let val = anchor_lang::AnchorDeserialize::deserialize(&mut &data[offset..])?;
+                    unsafe { self.__account.borrow_mut().as_mut_ptr().write(val) };
+
+                    // Set fields to initialized
+                    let mut fields = self.__fields.borrow_mut();
+                    let fields = fields.as_mut().unwrap();
+                    for field in fields {
+                        *field = true;
+                    }
+                } else {
+                    // Only initialize uninitialized fields (`load` methods already do this).
+                    //
+                    // This is not exactly efficient because `load` methods have a bit of
+                    // runtime ownership overhead. This could be optimized further, but it
+                    // requires some refactoring and also makes the code harder to reason about.
+                    //
+                    // We can return back to this if benchmarks show this is a bottleneck.
+                    #(self.#load_idents()?;)*
+                }
+
+                Ok(f())
+            }
+
+            #(#loader_impls)*
+
+            #[inline(always)]
+            fn #initialize_fields(&self) {
+                if self.__fields.borrow().is_none() {
+                    *self.__fields.borrow_mut() = Some(vec![false; #total_fields]);
+                }
+            }
+
+            // TODO: This method can be optimized to *only* serialize the fields that we have
+            // initialized rather than deserializing the whole account, and then serializing it
+            // back, which consumes a lot more CUs than it should for most accounts.
+            fn exit(&self, program_id: &anchor_lang::prelude::Pubkey) -> anchor_lang::Result<()> {
+                // Only persist if the owner is the current program and the account is not closed
+                if &<#ident as anchor_lang::Owner>::owner() == program_id
+                    && !anchor_lang::__private::is_closed(self.__info)
+                {
+                    // Make sure all fields are initialized
+                    let acc = self.load()?;
+                    let mut data = self.__info.try_borrow_mut_data()?;
+                    let dst: &mut [u8] = &mut data;
+                    let mut writer = anchor_lang::__private::BpfWriter::new(dst);
+                    acc.try_serialize(&mut writer)?;
+                }
+
+                Ok(())
+            }
+        }
+    })
+}
+
+/// Get the field's ident and if the ident doesn't exist (e.g. for tuple structs), default to the
+/// given index.
+fn to_field_ident(field: &syn::Field, i: usize) -> TokenStream {
+    field
+        .ident
+        .as_ref()
+        .map(ToTokens::to_token_stream)
+        .unwrap_or_else(|| Literal::usize_unsuffixed(i).to_token_stream())
+}
+
+/// Convert to private ident.
+///
+/// This is used to indicate to the users that they shouldn't use this identifier.
+fn to_private_ident<S: AsRef<str>>(ident: S) -> syn::Ident {
+    format_ident!("__{}", ident.as_ref())
+}

lang/attribute/account/src/lib.rs

@@ -12,6 +12,9 @@ use syn::{
 
 mod id;
 
+#[cfg(feature = "lazy-account")]
+mod lazy;
+
 /// An attribute for a data structure representing a Solana account.
 ///
 /// `#[account]` generates trait implementations for the following traits:
@@ -207,6 +210,17 @@ pub fn account(
                 #owner_impl
             }
         } else {
+            let lazy = {
+                #[cfg(feature = "lazy-account")]
+                match namespace.is_empty().then(|| lazy::gen_lazy(&account_strct)) {
+                    Some(Ok(lazy)) => lazy,
+                    // If lazy codegen fails for whatever reason, return empty tokenstream which
+                    // will make the account unusable with `LazyAccount<T>`
+                    _ => Default::default(),
+                }
+                #[cfg(not(feature = "lazy-account"))]
+                proc_macro2::TokenStream::default()
+            };
             quote! {
                 #[derive(AnchorSerialize, AnchorDeserialize, Clone)]
                 #account_strct
@@ -251,6 +265,8 @@ pub fn account(
                 }
 
                 #owner_impl
+
+                #lazy
             }
         }
     })

lang/derive/serde/Cargo.toml

@@ -12,6 +12,7 @@ proc-macro = true
 
 [features]
 idl-build = ["anchor-syn/idl-build"]
+lazy-account = []
 
 [dependencies]
 anchor-syn = { path = "../../syn", version = "0.30.1" }