diff --git a/lang-v2/Cargo.toml b/lang-v2/Cargo.toml index 69c9e4aedc..fe05fa3069 100644 --- a/lang-v2/Cargo.toml +++ b/lang-v2/Cargo.toml @@ -169,6 +169,9 @@ required-features = ["testing"] [[test]] name = "address_idl_surface" +[[test]] +name = "init_space" + [[test]] name = "update_hook_order" required-features = ["testing"] diff --git a/lang-v2/derive/src/init_space.rs b/lang-v2/derive/src/init_space.rs index 02a5bb61c0..42af9c6d40 100644 --- a/lang-v2/derive/src/init_space.rs +++ b/lang-v2/derive/src/init_space.rs @@ -164,49 +164,55 @@ fn len_from_type(ty: Type, attrs: &mut Option>) -> TokenS quote!((#array_len * #type_len)) } Type::Path(ty_path) => { - let path_segment = ty_path - .path - .segments - .last() - .expect("syn::TypePath always has at least one segment"); - let ident = &path_segment.ident; - let type_name = ident.to_string(); - let first_ty = get_first_ty_arg(&path_segment.arguments); + if let Some(type_name) = builtin_type_name(&ty_path) { + let path_segment = ty_path + .path + .segments + .last() + .expect("syn::TypePath always has at least one segment"); + let ident = &path_segment.ident; + let first_ty = get_first_ty_arg(&path_segment.arguments); - match type_name.as_str() { - "i8" | "u8" | "bool" => quote!(1), - "i16" | "u16" => quote!(2), - "i32" | "u32" | "f32" => quote!(4), - "i64" | "u64" | "f64" => quote!(8), - "i128" | "u128" => quote!(16), - "String" => { - let max_len = get_next_arg(ident, attrs); - quote!((4 + #max_len)) - } - "Pubkey" | "Address" => quote!(32), - "Option" => { - if let Some(ty) = first_ty { - let type_len = len_from_type(ty, attrs); + match type_name { + "i8" | "u8" | "bool" => quote!(1), + "i16" | "u16" => quote!(2), + "i32" | "u32" | "f32" => quote!(4), + "i64" | "u64" | "f64" => quote!(8), + "i128" | "u128" => quote!(16), + "String" => { + let max_len = get_next_arg(ident, attrs); + quote!((4 + #max_len)) + } + "Pubkey" | "Address" => quote!(32), + "Option" => { + if let Some(ty) = first_ty { + let type_len = len_from_type(ty, attrs); - quote!((1 + #type_len)) - } else { - quote_spanned!(ident.span() => compile_error!("Invalid argument in Option")) + quote!((1 + #type_len)) + } else { + quote_spanned!(ident.span() => compile_error!("Invalid argument in Option")) + } } - } - "Vec" => { - if let Some(ty) = first_ty { - let max_len = get_next_arg(ident, attrs); - let type_len = len_from_type(ty, attrs); + "Vec" => { + if let Some(ty) = first_ty { + let max_len = get_next_arg(ident, attrs); + let type_len = len_from_type(ty, attrs); - quote!((4 + #type_len * #max_len)) - } else { - quote_spanned!(ident.span() => compile_error!("Invalid argument in Vec")) + quote!((4 + #type_len * #max_len)) + } else { + quote_spanned!(ident.span() => compile_error!("Invalid argument in Vec")) + } } + _ => unreachable!("all builtin type names should be covered"), } - _ => { - let ty = &ty_path.path; - quote!(<#ty as anchor_lang::Space>::INIT_SPACE) - } + } else { + // Keep the full TypePath so `::Assoc` retains + // its qself; quoting only `.path` would emit `Trait::Assoc`. + // + // Arbitrary qualified paths such as `custom::Address` must not + // match built-in shortcuts keyed on the final segment name + // alone. + quote!(<#ty_path as anchor_lang::Space>::INIT_SPACE) } } Type::Tuple(ty_tuple) => { @@ -231,6 +237,50 @@ fn len_from_type(ty: Type, attrs: &mut Option>) -> TokenS } } +fn builtin_type_name(ty_path: &syn::TypePath) -> Option<&'static str> { + const PRIMITIVES_AND_ALIASES: &[(&str, &[&str])] = &[ + ("i8", &["i8"]), + ("u8", &["u8"]), + ("bool", &["bool"]), + ("i16", &["i16"]), + ("u16", &["u16"]), + ("i32", &["i32"]), + ("u32", &["u32"]), + ("f32", &["f32"]), + ("i64", &["i64"]), + ("u64", &["u64"]), + ("f64", &["f64"]), + ("i128", &["i128"]), + ("u128", &["u128"]), + ("String", &["String"]), + ("String", &["alloc", "string", "String"]), + ("String", &["std", "string", "String"]), + ("Pubkey", &["Pubkey"]), + ("Address", &["Address"]), + ("Option", &["Option"]), + ("Option", &["core", "option", "Option"]), + ("Option", &["std", "option", "Option"]), + ("Vec", &["Vec"]), + ("Vec", &["alloc", "vec", "Vec"]), + ("Vec", &["std", "vec", "Vec"]), + ]; + + PRIMITIVES_AND_ALIASES + .iter() + .find_map(|(name, segments)| path_matches(ty_path, segments).then_some(*name)) +} + +fn path_matches(ty_path: &syn::TypePath, segments: &[&str]) -> bool { + ty_path.qself.is_none() + && ty_path.path.segments.len() == segments.len() + && ty_path + .path + .segments + .iter() + .zip(segments.iter()) + .all(|(segment, expected)| segment.ident == *expected) +} + fn get_first_ty_arg(args: &PathArguments) -> Option { match args { PathArguments::AngleBracketed(bracket) => bracket.args.iter().find_map(|el| match el { diff --git a/lang-v2/tests/init_space.rs b/lang-v2/tests/init_space.rs index d1cd177fad..c7ccd43446 100644 --- a/lang-v2/tests/init_space.rs +++ b/lang-v2/tests/init_space.rs @@ -9,6 +9,7 @@ #![allow(dead_code)] use anchor_lang::{prelude::*, InitSpace, Space}; +extern crate alloc; #[derive(InitSpace)] struct Primitives { @@ -66,6 +67,16 @@ fn option_adds_one_byte_discriminator() { assert_eq!(WithOption::INIT_SPACE, 9); } +#[derive(InitSpace)] +struct WithQualifiedCoreOption { + _maybe: core::option::Option, // 1 + 8 +} + +#[test] +fn qualified_core_option_uses_option_layout() { + assert_eq!(WithQualifiedCoreOption::INIT_SPACE, 9); +} + #[derive(InitSpace)] struct WithString { #[max_len(32)] @@ -77,6 +88,17 @@ fn string_reserves_max_len_plus_length_prefix() { assert_eq!(WithString::INIT_SPACE, 36); } +#[derive(InitSpace)] +struct WithQualifiedAllocString { + #[max_len(32)] + _name: alloc::string::String, // 4 + 32 +} + +#[test] +fn qualified_alloc_string_uses_string_layout() { + assert_eq!(WithQualifiedAllocString::INIT_SPACE, 36); +} + #[derive(InitSpace)] struct WithVec { #[max_len(10)] @@ -88,6 +110,17 @@ fn vec_reserves_max_len_times_element_plus_prefix() { assert_eq!(WithVec::INIT_SPACE, 84); } +#[derive(InitSpace)] +struct WithQualifiedAllocVec { + #[max_len(10)] + _xs: alloc::vec::Vec, // 4 + 8 * 10 +} + +#[test] +fn qualified_alloc_vec_uses_vec_layout() { + assert_eq!(WithQualifiedAllocVec::INIT_SPACE, 84); +} + #[derive(InitSpace)] struct WithVecOfStrings { #[max_len(4, 16)] @@ -126,6 +159,46 @@ fn nested_struct_uses_inner_init_space() { assert_eq!(Outer::INIT_SPACE, 8 + 24); } +trait Schema { + type Value: Space; +} + +struct SchemaHost; + +impl Schema for SchemaHost { + type Value = u64; +} + +#[derive(InitSpace)] +struct WithQualifiedAssoc { + _value: ::Value, +} + +#[test] +fn qualified_associated_type_preserves_qself() { + assert_eq!(WithQualifiedAssoc::INIT_SPACE, 8); +} + +mod custom { + use super::Space; + + pub struct Address; + + impl Space for Address { + const INIT_SPACE: usize = 8; + } +} + +#[derive(InitSpace)] +struct WithCustomAddress { + _addr: custom::Address, +} + +#[test] +fn qualified_path_does_not_use_builtin_address_size() { + assert_eq!(WithCustomAddress::INIT_SPACE, 8); +} + #[derive(InitSpace)] enum Variant { A, // 0