Example of Rust attribute macros: data serialization (part 2 - enums)
Tags: tutorial, rust, blockchain
Published: 2024-02-27 18:43:39

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

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

```rust
{
    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,

```rust
#[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!

```rust
#[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.

```rust
...
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:

```rust
    ... 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.

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

### Going over each variant

```rust
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`.

```rust
    ...
    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:

```rust
    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

```rust
#[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](https://github.com/ProgramCrafter/tlb-rust-serialization/blob/c7dc4111658c178b1d7262ee12b32e66e17ae3ce/tlb_macro/src/lib.rs).