Example of Rust attribute macros: data serialization (part 2 - enums)

In the previous article, I demonstrated how serialization function for a struct can be auto-generated.

#[tlb_serializable(u 4 3bit, workchain, hash_high, hash_low)]
pub struct Address {
    workchain: u8,      hash_high: u128,      hash_low: u128
}

{
    let mut result = ::std::vec![];
    result.push("u 4 3bit".to_owned());
    {   let mut s_field = crate::ton::CellSerialize::serialize(&self.workchain);
        result.append(&mut s_field);   }
    {   let mut s_field = crate::ton::CellSerialize::serialize(&self.hash_high);
        result.append(&mut s_field);   }
    {   let mut s_field = crate::ton::CellSerialize::serialize(&self.hash_low);
        result.append(&mut s_field);   }
    result
}

Now, let's do the same for enums!

I feel like this article may be quite hard to understand; please respond in comments whether it was really, or was I too scared?

Designing the solution

Variants of enum are determined by tag, which is stored before data of variant itself. In TON, that tag may actually be any, even non-unique or empty - if you can make your smart contracts parse that.

So, our code needs to get how to represent the tags. Let's use normal Rust attribute #[repr(u**)] if the tag is wanted or custom "attribute" #[tlb_assert_unsafe(items_prefixes_nonoverlap)] if it's not.

After that, we will also need description how to serialize each variant of the enumeration. #[tlb_item_serializable()] will be a logical name for such an attribute.

Coding

Let's start with a simple template,

#[proc_macro_attribute]
pub fn tlb_enum_serializable(_: OldTokenStream, item: OldTokenStream)
        -> OldTokenStream {
    let mut input: ItemEnum = parse_macro_input!(item);
    let name = input.ident.clone();
    ...
}

Parsing enum attributes

For clarity (including error messages), we define our enum determining if user wants automatically added tag in enum serialization, and if he does, then what numeric type is used.

Also, as we don't create procedural attribute #[tlb_assert_unsafe()], we need to remove it so Rust doesn't have to search for it. We can use retain function on input.attrs just for this!

#[derive(Debug)] enum TlbPrefix {Wanted(String), NotWanted}

    ...
    let mut need_prefix: Option<TlbPrefix> = None;
    input.attrs.retain(|attr| {
        if attr.path().is_ident("tlb_assert_unsafe") {
            ...
        } else ...
    });

Attributes can be empty, contain an assignment/comparison (for instance, #[cfg(feature = "std")], or contain arbitrary token sequences. We're interested in the third kind.

...
input.attrs.retain(|attr| {
    if attr.path().is_ident("tlb_assert_unsafe") {
        let Meta::List(MetaList {tokens: ref tokens_assert, ..}) = attr.meta else {
            panic!("#[tlb_assert_unsafe] attribute must have argument with the specific assertion");
        };
        let assertion = tokens_assert.to_string();
        if assertion == "items_prefixes_nonoverlap" {
            assert!(need_prefix.is_none());
            need_prefix = Some(TlbPrefix::NotWanted);
            false    // don't keep the attribute on enum
        } else {
            println!("Unknown assertion {assertion:?}");
            true
        }
    } else ...

The second half is written in pretty much same way:

    ... else if attr.path().is_ident("repr") {
        assert!(need_prefix.is_none(), "Two #[repr] attributes on enum are not supported");
        let Meta::List(MetaList {tokens: ref tokens_type, ..}) = attr.meta else {
            panic!("#[repr] attribute must have argument specifying the type");
        };
        need_prefix = Some(TlbPrefix::Wanted(tokens_type.to_string()));
        true    // we retain #[repr] attribute for Rust's use
    } else {
        true
    }
});

After iterating over enum's attributes, we now have field need_prefix. Let's assert it contains something and go on.

let need_prefix: TlbPrefix = need_prefix.expect(
    "Don't know how to differentiate tags of the enum");

Going over each variant

let mut variant_index = 0;
let variant_generators: Vec<V2TokenStream> = input.variants.iter_mut().map(|variant| {
    let mut store = None;
    variant.attrs.retain(|attr| {
        if !attr.path().is_ident("tlb_item_serializable") {return true;}
        let Meta::List(MetaList {tokens: ref tokens_tlb, ..}) = attr.meta else {
            panic!("#[tlb_item_serializable] attribute must have argument with the specific serialization");
        };
        let tlb = tokens_tlb.to_string();

        assert!(store.is_none(), "multiple serialization definitions found");
        store = Some(create_serialization_code(&tlb, &variant.fields));
        false
    });
    let store = store.expect(&format!("serialization definition for variant {} is required", variant.ident));
    ...

For now, we've determined how to store inner data of each option. But we may also need to store discriminant, and that's why we need variant_index.

    ...
    if let Some((_, Expr::Lit(ref idx))) = variant.discriminant {
        if let Lit::Int(ref discriminant) = idx.lit {
            variant_index = discriminant.base10_parse::<u64>()
                .unwrap();
        }
    };
    let store_tag = match need_prefix {
        TlbPrefix::NotWanted => quote! {},
        TlbPrefix::Wanted(ref discr_type) => {
            let s = &discr_type[1..];    // u64 -> 64
            quote! {
                result.push(::std::format!("u {} {}bit",
                    #variant_index, #s));
            }
        },
    };
    variant_index += 1;

You may think at this point everything is ready to create match branch for each option. Something isn't ready yet: we need to unpack data into variables + tell create_serialization_code function that self is not required. The latter can be done with a simple boolean flag, and the former is written so:

    let vident = &variant.ident;
    let fields_unpacker: Vec<_> = variant.fields.iter().map(|field| {
        let id = field.ident.clone().expect("unnamed field in enum");
        quote!{ #id, }
    }).collect();

    quote! {
        #name::#vident {#(#fields_unpacker)*} => {
            #store_tag
            #store
        }
    }

Combining all the parts

#[proc_macro_attribute]
pub fn tlb_enum_serializable(_: OldTokenStream, item: OldTokenStream)
        -> OldTokenStream {
    let mut input: ItemEnum = parse_macro_input!(item);
    let name = input.ident.clone();

    let mut need_prefix: Option<TlbPrefix> = None;
    input.attrs.retain(|attr| { ... });
    let need_prefix: TlbPrefix = need_prefix.expect(
        "Don't know how to differentiate tags of the enum");

    let mut variant_index = 0;
    let variant_generators: Vec<V2TokenStream> =
        input.variants.iter_mut().map(|variant| {...})
        .collect();

    // we need that because we could've removed some attributes
    let mut result: OldTokenStream = input.to_token_stream().into();
    result.extend(OldTokenStream::from(quote! {
        impl crate::ton::CellSerialize for #name {
            fn serialize(&self) -> ::std::vec::Vec<::std::string::String> {
                let mut result = ::std::vec![];
                match &self {
                    #(#variant_generators)*
                }
                result
            }
        }
    }));
    result
}

The full code is on my Github - ProgramCrafter/tlb-rust-serialization.

Blog

Example of Rust attribute macros: data serialization (part 2 - enums)

ProgramCrafter

Designing the solution

Coding

Parsing enum attributes

Going over each variant

Combining all the parts

Join Our Newsletter. No Spam, Only the good stuff.

Related