Add #[borsh(use_discriminant = <bool>)] that change enum discriminant de- and serialization behavior #148

README.md

@@ -76,6 +76,37 @@ struct A {
 }
 ```
 
+### Enum with explicit discriminant
+
+`#[borsh(use_discriminant=false|true])` is required if you have an enum with explicit discriminant. This settings affects `BorshSerialize` and `BorshDeserialize` behaviour at the same time.
+
+If you don't specify `use_discriminant` option for enum with explicit discriminant, you will get an error:
+
+````bash
+error: You have to specify `#[borsh(use_discriminant=true)]` or `#[borsh(use_discriminant=false)]` for all structs that have enum with explicit discriminant
+```
+```rust
+#[derive(BorshDeserialize, BorshSerialize)]
+#[borsh(use_discriminant=false)]
+enum A {
+    X,
+    Y = 10,
+}
+````
+
+Will keep old behaviour of borsh deserialization and will not use discriminant. This option is left to have backward compatability with previous versions of borsh and to have ability to deserialise data from previous versions of borsh.
+
+```rust
+#[derive(BorshDeserialize, BorshSerialize)]
+#[borsh(use_discriminant=true)]
+enum A {
+    X,
+    Y = 10,
+}
+```
+
+This one will use proper version of serialization of enum with explicit discriminant.
+
 ## Releasing
 
 The versions of all public crates in this repository are collectively managed by a single version in the [workspace manifest](https://github.com/near/borsh-rs/blob/master/Cargo.toml).

borsh-derive-internal/src/attribute_helpers.rs

@@ -1,6 +1,8 @@
 // TODO: remove this unused attribute, when the unsplit is done
 #![allow(unused)]
-use syn::{Attribute, Field, Path, WherePredicate};
+use proc_macro2::TokenStream;
+use quote::ToTokens;
+use syn::{spanned::Spanned, Attribute, DeriveInput, Expr, Field, Path, WherePredicate};
 pub mod parsing_helpers;
 use parsing_helpers::get_where_predicates;
 
@@ -13,6 +15,8 @@ pub struct Symbol(pub &'static str);
 pub const BORSH: Symbol = Symbol("borsh");
 /// bound - sub-borsh nested meta, field-level only, `BorshSerialize` and `BorshDeserialize` contexts
 pub const BOUND: Symbol = Symbol("bound");
+// item level attribute for enums
+pub const USE_DISCRIMINANT: &str = "use_discriminant";
 /// serialize - sub-bound nested meta attribute
 pub const SERIALIZE: Symbol = Symbol("serialize");
 /// deserialize - sub-bound nested meta attribute
@@ -42,6 +46,52 @@ pub(crate) fn contains_skip(attrs: &[Attribute]) -> bool {
     attrs.iter().any(|attr| attr.path() == SKIP)
 }
 
+pub fn check_item_attributes(derive_input: &DeriveInput) -> Result<(), proc_macro2::TokenStream> {
+    for attr in derive_input.attrs.clone() {
+        if attr.path().is_ident(BORSH.0) {
+            if let syn::Data::Struct(ref _data) = derive_input.data {
+                return Err(syn::Error::new(
+                    derive_input.ident.span(),
+                    "borsh (use_discriminant=<bool>) does not support structs",
+                )
+                .to_compile_error());
+            }
+        }
+    }
+    Ok(())
+}
+
+pub(crate) fn contains_use_discriminant(attrs: &[Attribute]) -> Result<Option<bool>, TokenStream> {
+    let mut result = None;
+    for attr in attrs {
+        if attr.path().is_ident("borsh") {
+            attr.parse_nested_meta(|meta| {
+                if meta.path.is_ident(USE_DISCRIMINANT) {
+                    let value_expr: Expr = meta.value()?.parse()?;
+                    let value = value_expr.to_token_stream().to_string();
+                    // this goes to contains_use_discriminant
+                    match value.as_str() {
+                        "true" => {
+                            result = Some(true);
+                        }
+                        "false" => result = Some(false),
+                        _ => {
+                            return Err(syn::Error::new(
+                                value_expr.span(),
+                                "`use_discriminant` accept only `true` or `false`",
+                            ));
+                        }
+                    };
+                }
+
+                Ok(())
+            })
+            .map_err(|err| err.to_compile_error())?;
+        }
+    }
+    Ok(result)
+}
+
 pub(crate) fn contains_initialize_with(attrs: &[Attribute]) -> Option<Path> {
     for attr in attrs.iter() {
         if attr.path() == INIT {
@@ -387,4 +437,96 @@ mod tests {
         let schema_params = parse_schema_attrs(&first_field.attrs).unwrap();
         assert!(schema_params.is_none());
     }
+
+    use super::*;
+    #[test]
+    fn test_check_use_discriminant() {
+        let item_enum: DeriveInput = syn::parse2(quote! {
+            #[derive(BorshDeserialize, Debug)]
+            #[borsh(use_discriminant = false)]
+            enum AWithUseDiscriminantFalse {
+                X,
+                Y,
+            }
+        })
+        .unwrap();
+        let actual = contains_use_discriminant(&item_enum.attrs);
+        assert!(actual.is_ok());
+    }
+
+    #[test]
+    fn test_check_use_discriminant_true() {
+        let item_enum: DeriveInput = syn::parse2(quote! {
+            #[derive(BorshDeserialize, Debug)]
+            #[borsh(use_discriminant = true)]
+            enum AWithUseDiscriminantTrue {
+                X,
+                Y,
+            }
+        })
+        .unwrap();
+        let actual = contains_use_discriminant(&item_enum.attrs);
+        assert!(actual.is_ok());
+    }
+
+    #[test]
+    fn test_check_use_discriminant_wrong_value() {
+        let item_enum: DeriveInput = syn::parse2(quote! {
+            #[derive(BorshDeserialize, Debug)]
+            #[borsh(use_discriminant = 111)]
+            enum AWithUseDiscriminantFalse {
+                X,
+                Y,
+            }
+        })
+        .unwrap();
+        let actual = contains_use_discriminant(&item_enum.attrs);
+        assert!(actual.is_err());
+    }
+    #[test]
+    #[should_panic]
+    fn test_check_use_discriminant_on_struct() {
+        let item_enum: DeriveInput = syn::parse2(quote! {
+            #[derive(BorshDeserialize, Debug)]
+            #[borsh(use_discriminant = false)]
+            struct AWithUseDiscriminantFalse {
+                x: X,
+                y: Y,
+            }
+        })
+        .unwrap();
+        let actual = contains_use_discriminant(&item_enum.attrs);
+        insta::assert_snapshot!(actual.unwrap_err().to_token_stream().to_string());
+    }
+    #[test]
+    #[should_panic]
+    fn test_check_use_borsh_skip_on_whole_struct() {
+        let item_enum: DeriveInput = syn::parse2(quote! {
+            #[derive(BorshDeserialize, Debug)]
+            #[borsh(use_discriminant = false)]
+            #[borsh_skip=x]
+            struct AWithUseDiscriminantFalse {
+                x: X,
+                y: Y,
+            }
+        })
+        .unwrap();
+        let actual = contains_use_discriminant(&item_enum.attrs);
+        insta::assert_snapshot!(actual.unwrap_err().to_token_stream().to_string());
+    }
+    #[test]
+    #[should_panic]
+    fn test_check_use_borsh_invalid_on_whole_struct() {
+        let item_enum: DeriveInput = syn::parse2(quote! {
+            #[derive(BorshDeserialize, Debug)]
+            #[borsh(invalid)]
+            struct AWithUseDiscriminantFalse {
+                x: X,
+                y: Y,
+            }
+        })
+        .unwrap();
+        let actual = contains_use_discriminant(&item_enum.attrs);
+        insta::assert_snapshot!(actual.unwrap_err().to_token_stream().to_string());
+    }
 }

borsh-derive-internal/src/enum_de.rs

@@ -1,14 +1,15 @@
-use proc_macro2::TokenStream as TokenStream2;
-use quote::quote;
-use syn::{Fields, Ident, ItemEnum, Path, WhereClause};
-
 use crate::{
     attribute_helpers::{
-        collect_override_bounds, contains_initialize_with, contains_skip, BoundType,
+        collect_override_bounds, contains_initialize_with, contains_skip,
+        contains_use_discriminant, BoundType,
     },
     enum_discriminant_map::discriminant_map,
     generics::{compute_predicates, without_defaults, FindTyParams},
 };
+use proc_macro2::TokenStream as TokenStream2;
+use quote::quote;
+use std::convert::TryFrom;
+use syn::{Fields, Ident, ItemEnum, Path, WhereClause};
 
 pub fn enum_de(input: &ItemEnum, cratename: Ident) -> syn::Result<TokenStream2> {
     let name = &input.ident;
@@ -27,9 +28,36 @@ pub fn enum_de(input: &ItemEnum, cratename: Ident) -> syn::Result<TokenStream2>
     let mut default_params_visitor = FindTyParams::new(&generics);
 
     let init_method = contains_initialize_with(&input.attrs);
+
+    let use_discriminant = contains_use_discriminant(&input.attrs).map_err(|err| {
+        syn::Error::new(
+            input.ident.span(),
+            format!("error parsing `#[borsh(use_discriminant = ...)]`: {}", err),
+        )
+    })?;
+
+    let has_explicit_discriminants = input
+        .variants
+        .iter()
+        .any(|variant| variant.discriminant.is_some());
+    if has_explicit_discriminants && use_discriminant.is_none() {
+        return Err(syn::Error::new(
+        input.ident.span(),
+        "You have to specify `#[borsh(use_discriminant=true)]` or `#[borsh(use_discriminant=false)]` for all structs that have enum with explicit discriminant",
+    ));
+    }
+    let use_discriminant = use_discriminant.unwrap_or(false);
+
     let mut variant_arms = TokenStream2::new();
     let discriminants = discriminant_map(&input.variants);
-    for variant in input.variants.iter() {
+
+    for (variant_idx, variant) in input.variants.iter().enumerate() {
+        let variant_idx = u8::try_from(variant_idx).map_err(|err| {
+            syn::Error::new(
+                variant.ident.span(),
+                format!("up to 256 enum variants are supported. error{}", err),
+            )
+        })?;
         let variant_ident = &variant.ident;
         let discriminant = discriminants.get(variant_ident).unwrap();
         let mut variant_header = TokenStream2::new();
@@ -85,6 +113,11 @@ pub fn enum_de(input: &ItemEnum, cratename: Ident) -> syn::Result<TokenStream2>
             }
             Fields::Unit => {}
         }
+        let discriminant = if use_discriminant {
+            quote! { #discriminant }
+        } else {
+            quote! { #variant_idx }
+        };
         variant_arms.extend(quote! {
             if variant_tag == #discriminant { #name::#variant_ident #variant_header } else
         });

borsh-derive-internal/src/enum_discriminant_map.rs

@@ -8,14 +8,14 @@ use syn::{punctuated::Punctuated, token::Comma, Variant};
 /// See: https://doc.rust-lang.org/reference/items/enumerations.html#assigning-discriminant-values
 pub fn discriminant_map(variants: &Punctuated<Variant, Comma>) -> HashMap<Ident, TokenStream> {
     let mut map = HashMap::new();
-
     let mut next_discriminant_if_not_specified = quote! {0};
 
     for variant in variants {
         let this_discriminant = variant.discriminant.clone().map_or_else(
             || quote! { #next_discriminant_if_not_specified },
             |(_, e)| quote! { #e },
         );
+
         next_discriminant_if_not_specified = quote! { #this_discriminant + 1 };
         map.insert(variant.ident.clone(), this_discriminant);
     }