chore(versionable): add '#[versionize(dispatch = T)]' in macro

This allows to add new attributes without arguments to the proc-macro
zama-ai · Oct 2, 2024 · fc1ce85 · fc1ce85
1 parent 77df5e4
commit fc1ce85
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 4 deletions.
diff --git a/utils/tfhe-versionable-derive/src/versionize_attribute.rs b/utils/tfhe-versionable-derive/src/versionize_attribute.rs
@@ -106,7 +106,14 @@ impl VersionizeAttribute {
                             attribute_builder.into =
                                 Some(parse_path_ignore_quotes(&name_value.value)?);
                         }
-                        // parse versionize(bound = "Type: Bound")
+                        // parse versionize(dispatch = "Type")
+                    } else if name_value.path.is_ident("dispatch") {
+                        if attribute_builder.dispatch_enum.is_some() {
+                            return Err(Self::default_error(meta.span()));
+                        } else {
+                            attribute_builder.dispatch_enum =
+                                Some(parse_path_ignore_quotes(&name_value.value)?);
+                        }
                     } else {
                         return Err(Self::default_error(meta.span()));
                     }

diff --git a/utils/tfhe-versionable/examples/simple.rs b/utils/tfhe-versionable/examples/simple.rs
@@ -6,8 +6,9 @@ use tfhe_versionable::{Unversionize, Upgrade, Version, Versionize, VersionsDispa
 
 // The structure that should be versioned, as defined in your code
 #[derive(Versionize)]
-#[versionize(MyStructVersions)] // Link to the enum type that will holds all the versions of this
-                                // type
+// We have to link to the enum type that will holds all the versions of this
+// type. This can also be written `#[versionize(dispatch = MyStructVersions)]`.
+#[versionize(MyStructVersions)]
 struct MyStruct<T> {
     attr: T,
     builtin: u32,

diff --git a/utils/tfhe-versionable/tests/testcases/struct.rs b/utils/tfhe-versionable/tests/testcases/struct.rs
@@ -37,6 +37,18 @@ pub struct MyStruct2<T, U> {
     field1: U,
 }
 
+#[derive(Versionize)]
+#[versionize(dispatch = MyStruct3Versions)]
+pub struct MyStruct3<T> {
+    field0: u64,
+    field1: T,
+}
+
+#[derive(VersionsDispatch)]
+pub enum MyStruct3Versions<T> {
+    V0(MyStruct3<T>),
+}
+
 fn main() {
     assert_impl_all!(MyEmptyStruct: Version);
     assert_impl_all!(MyEmptyStruct2: Version);
@@ -47,7 +59,9 @@ fn main() {
 
     assert_impl_all!(MyAnonStruct3<u64>: Version);
 
-    assert_impl_all!(MyStruct<u32>: Version);
+    assert_impl_all!(MyStruct<u32>: Versionize);
 
     assert_impl_all!(MyStruct2<usize, String>: Version);
+
+    assert_impl_all!(MyStruct3<u32>: Versionize);
 }