lang: Optimize sized fields when using LazyAccount (#3218)

coral-xyz · Sep 2, 2024 · 06527e5 · 06527e5
1 parent 879601e
commit 06527e5
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 5 deletions.
diff --git a/lang/attribute/account/src/lazy.rs b/lang/attribute/account/src/lazy.rs
@@ -57,10 +57,14 @@ pub fn gen_lazy(strct: &syn::ItemStruct) -> syn::Result<TokenStream> {
             });
 
             let ty = &field.ty;
+            let ty_as_lazy = quote! { <#ty as anchor_lang::__private::Lazy> };
             let size = quote! {
-                <#ty as anchor_lang::__private::Lazy>::size_of(
-                    &self.__info.data.borrow()[self.#offset_of_ident()..]
-                )
+                // Calculating the offset is highly wasteful if the type is sized.
+                if #ty_as_lazy::SIZED {
+                    #ty_as_lazy::size_of(&[])
+                } else {
+                    #ty_as_lazy::size_of(&self.__info.data.borrow()[self.#offset_of_ident()..])
+                }
             };
 
             let signatures = quote! {

diff --git a/lang/derive/serde/src/lazy.rs b/lang/derive/serde/src/lazy.rs
@@ -4,8 +4,18 @@ use syn::{spanned::Spanned, Fields, Item};
 
 pub fn gen_lazy(input: proc_macro::TokenStream) -> syn::Result<proc_macro2::TokenStream> {
     let item = syn::parse::<Item>(input)?;
-    let (name, generics, size) = match &item {
-        Item::Struct(strct) => (&strct.ident, &strct.generics, sum_fields(&strct.fields)),
+    let (name, generics, size, sized) = match &item {
+        Item::Struct(strct) => (
+            &strct.ident,
+            &strct.generics,
+            sum_fields(&strct.fields),
+            strct
+                .fields
+                .iter()
+                .map(|field| &field.ty)
+                .map(|ty| quote! { <#ty as anchor_lang::__private::Lazy>::SIZED })
+                .fold(quote!(true), |acc, sized| quote! { #acc && #sized }),
+        ),
         Item::Enum(enm) => {
             let arms = enm
                 .variants
@@ -24,6 +34,7 @@ pub fn gen_lazy(input: proc_macro::TokenStream) -> syn::Result<proc_macro2::Toke
                         _ => unreachable!(),
                     }
                 },
+                quote!(false),
             )
         }
         Item::Union(_) => return Err(syn::Error::new(item.span(), "Unions are not supported")),
@@ -34,6 +45,8 @@ pub fn gen_lazy(input: proc_macro::TokenStream) -> syn::Result<proc_macro2::Toke
 
     Ok(quote! {
         impl #impl_generics anchor_lang::__private::Lazy for #name #ty_generics #where_clause {
+            const SIZED: bool = #sized;
+
             #[inline(always)]
             fn size_of(buf: &[u8]) -> usize {
                 #size

diff --git a/lang/src/lazy.rs b/lang/src/lazy.rs
@@ -19,6 +19,9 @@ use crate::{AnchorDeserialize, Pubkey};
 /// }
 /// ```
 pub trait Lazy: AnchorDeserialize {
+    /// Whether the type is a fixed-size type.
+    const SIZED: bool = false;
+
     /// Get the serialized size of the type from the given buffer.
     ///
     /// For performance reasons, this method does not verify the validity of the data, and should
@@ -35,6 +38,8 @@ pub trait Lazy: AnchorDeserialize {
 macro_rules! impl_sized {
     ($ty: ty) => {
         impl Lazy for $ty {
+            const SIZED: bool = true;
+
             #[inline(always)]
             fn size_of(_buf: &[u8]) -> usize {
                 ::core::mem::size_of::<$ty>()
@@ -59,20 +64,26 @@ impl_sized!(f64);
 impl_sized!(Pubkey);
 
 impl<T: Lazy, const N: usize> Lazy for [T; N] {
+    const SIZED: bool = T::SIZED;
+
     #[inline(always)]
     fn size_of(buf: &[u8]) -> usize {
         N * T::size_of(buf)
     }
 }
 
 impl Lazy for String {
+    const SIZED: bool = false;
+
     #[inline(always)]
     fn size_of(buf: &[u8]) -> usize {
         LEN + get_len(buf)
     }
 }
 
 impl<T: Lazy> Lazy for Option<T> {
+    const SIZED: bool = false;
+
     #[inline(always)]
     fn size_of(buf: &[u8]) -> usize {
         1 + match buf.first() {
@@ -84,6 +95,8 @@ impl<T: Lazy> Lazy for Option<T> {
 }
 
 impl<T: Lazy> Lazy for Vec<T> {
+    const SIZED: bool = false;
+
     #[inline(always)]
     fn size_of(buf: &[u8]) -> usize {
         (0..get_len(buf)).fold(LEN, |acc, _| acc + T::size_of(&buf[acc..]))
@@ -166,6 +179,7 @@ mod tests {
                 c: Some(String::from("a"))
             })
         );
+        assert!(!MyStruct::SIZED);
 
         // Enum
         #[derive(AnchorSerialize, AnchorDeserialize)]
@@ -181,6 +195,7 @@ mod tests {
             MyEnum::size_of(&[2, 1, 2, 1, 2]),
             len!(MyEnum::Unnamed(1, 2))
         );
+        assert!(!MyEnum::SIZED);
     }
 
     #[test]
@@ -194,9 +209,12 @@ mod tests {
             GenericStruct::<i64>::size_of(&[1, 2, 3, 4, 5, 6, 7, 8]),
             len!(GenericStruct { t: 1i64 })
         );
+        assert!(GenericStruct::<i64>::SIZED);
+
         assert_eq!(
             GenericStruct::<Vec<u8>>::size_of(&[8, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8]),
             len!(GenericStruct { t: vec![0u8; 8] })
         );
+        assert!(!GenericStruct::<Vec<u8>>::SIZED);
     }
 }