Derive StorageAccess impl for enums

corelib/src/integer.cairo

@@ -1059,7 +1059,7 @@ impl U64BitOr of BitOr<u64> {
     }
 }
 
-#[derive(Copy, Drop, PartialEq, Serde, storage_access::StorageAccess)]
+#[derive(Copy, Drop, PartialEq, Serde, starknet::StorageAccess)]
 struct u256 {
     low: u128,
     high: u128,

corelib/src/starknet/eth_address.cairo

@@ -6,7 +6,7 @@ use traits::{Into, TryInto};
 use zeroable::Zeroable;
 
 // An Ethereum address (160 bits).
-#[derive(Copy, Drop, storage_access::StorageAccess)]
+#[derive(Copy, Drop, starknet::StorageAccess)]
 struct EthAddress {
     address: felt252, 
 }

corelib/src/starknet/secp256_trait.cairo

@@ -7,7 +7,7 @@ use traits::{Into, TryInto};
 use integer::U256TryIntoNonZero;
 
 /// Secp256{k/r}1 ECDSA signature.
-#[derive(Copy, Drop, PartialEq, Serde, storage_access::StorageAccess)]
+#[derive(Copy, Drop, PartialEq, Serde, starknet::StorageAccess)]
 struct Signature {
     r: u256,
     s: u256,

crates/cairo-lang-starknet/cairo_level_tests/storage_access.cairo

@@ -19,14 +19,21 @@ impl StorageAddressPartialEq of PartialEq<StorageAddress> {
     }
 }
 
-#[derive(Drop, Serde, PartialEq, Copy, storage_access::StorageAccess)]
+#[derive(Drop, Serde, PartialEq, Copy, starknet::StorageAccess)]
 struct Abc {
     a: u8,
     b: u16,
     c: u32,
 }
 
-#[derive(Drop, Serde, PartialEq, Copy, storage_access::StorageAccess)]
+#[derive(Drop, Serde, PartialEq, Copy, starknet::StorageAccess)]
+enum Efg {
+    E: (),
+    F: (),
+    G: u256
+}
+
+#[derive(Drop, Serde, PartialEq, Copy, starknet::StorageAccess)]
 struct AbcEtc {
     a: u8,
     b: u16,
@@ -40,6 +47,8 @@ struct AbcEtc {
     j: bool,
     k: EthAddress,
     abc: Abc,
+    efg1: Efg,
+    efg2: Efg,
 }
 
 
@@ -64,7 +73,7 @@ mod test_contract {
 }
 
 #[test]
-#[available_gas(900000)]
+#[available_gas(10000000)]
 fn write_read_struct() {
     let x = AbcEtc {
         a: 1_u8,
@@ -80,7 +89,7 @@ fn write_read_struct() {
         k: 123_felt252.try_into().unwrap(),
         abc: Abc {
             a: 1_u8, b: 2_u16, c: 3_u32, 
-        }
+        }, efg1: Efg::E(()), efg2: Efg::G(123_u256)
     };
 
     assert(test_contract::__external::set_data(serialized_element(*@x)).is_empty(), 'Not empty');

crates/cairo-lang-starknet/src/plugin/dispatcher.rs

@@ -205,7 +205,7 @@ pub fn handle_trait(db: &dyn SyntaxGroup, trait_ast: ast::ItemTrait) -> PluginRe
             $dispatcher_signatures$
             }}
 
-            #[derive(Copy, Drop, storage_access::StorageAccess, Serde)]
+            #[derive(Copy, Drop, starknet::StorageAccess, Serde)]
             struct {contract_caller_name} {{
                 contract_address: starknet::ContractAddress,
             }}
@@ -214,7 +214,7 @@ pub fn handle_trait(db: &dyn SyntaxGroup, trait_ast: ast::ItemTrait) -> PluginRe
             $contract_caller_method_impls$
             }}
 
-            #[derive(Copy, Drop, storage_access::StorageAccess, Serde)]
+            #[derive(Copy, Drop, starknet::StorageAccess, Serde)]
             struct {library_caller_name} {{
                 class_hash: starknet::ClassHash,
             }}
@@ -227,7 +227,7 @@ pub fn handle_trait(db: &dyn SyntaxGroup, trait_ast: ast::ItemTrait) -> PluginRe
             $safe_dispatcher_signatures$
             }}
 
-            #[derive(Copy, Drop, storage_access::StorageAccess, Serde)]
+            #[derive(Copy, Drop, starknet::StorageAccess, Serde)]
             struct {safe_library_caller_name} {{
                 class_hash: starknet::ClassHash,
             }}
@@ -238,7 +238,7 @@ pub fn handle_trait(db: &dyn SyntaxGroup, trait_ast: ast::ItemTrait) -> PluginRe
             }}
 
 
-            #[derive(Copy, Drop, storage_access::StorageAccess, Serde)]
+            #[derive(Copy, Drop, starknet::StorageAccess, Serde)]
             struct {safe_contract_caller_name} {{
                 contract_address: starknet::ContractAddress,
             }}

crates/cairo-lang-starknet/src/plugin/events.rs

@@ -194,10 +194,6 @@ fn get_field_kind_for_variant(
 
 /// Derive the `Event` trait for enums annotated with `derive(starknet::Event)`.
 pub fn handle_enum(db: &dyn SyntaxGroup, enum_ast: ast::ItemEnum) -> PluginResult {
-    if !derive_event_needed(&enum_ast, db) {
-        return PluginResult::default();
-    }
-
     let mut builder = PatchBuilder::new(db);
     let mut diagnostics = vec![];
     let enum_name = RewriteNode::new_trimmed(enum_ast.name(db).as_syntax_node());

crates/cairo-lang-starknet/src/plugin/mod.rs

@@ -25,7 +25,6 @@ use events::derive_event_needed;
 use storage_access::derive_storage_access_needed;
 
 use self::contract::{handle_contract_by_storage, handle_module};
-use self::events::handle_enum;
 
 #[derive(Debug, Default)]
 #[non_exhaustive]
@@ -45,7 +44,12 @@ impl MacroPlugin for StarkNetPlugin {
             ast::Item::Struct(struct_ast) if struct_ast.has_attr(db, STORAGE_ATTR) => {
                 handle_contract_by_storage(db, struct_ast).unwrap_or_default()
             }
-            ast::Item::Enum(enum_ast) => handle_enum(db, enum_ast),
+            ast::Item::Enum(enum_ast) if derive_storage_access_needed(&enum_ast, db) => {
+                storage_access::handle_enum(db, enum_ast)
+            }
+            ast::Item::Enum(enum_ast) if derive_event_needed(&enum_ast, db) => {
+                events::handle_enum(db, enum_ast)
+            }
             // Nothing to do for other items.
             _ => PluginResult::default(),
         }

crates/cairo-lang-starknet/src/plugin/plugin_test_data/dispatcher

@@ -36,7 +36,7 @@ trait IContractDispatcherTrait<T> {
     fn empty(self: T, );
 }
 
-#[derive(Copy, Drop, storage_access::StorageAccess, Serde)]
+#[derive(Copy, Drop, starknet::StorageAccess, Serde)]
 struct IContractDispatcher {
     contract_address: starknet::ContractAddress,
 }
@@ -73,7 +73,7 @@ impl IContractDispatcherImpl of IContractDispatcherTrait<IContractDispatcher> {
 
 }
 
-#[derive(Copy, Drop, storage_access::StorageAccess, Serde)]
+#[derive(Copy, Drop, starknet::StorageAccess, Serde)]
 struct IContractLibraryDispatcher {
     class_hash: starknet::ClassHash,
 }
@@ -115,7 +115,7 @@ trait IContractSafeDispatcherTrait<T> {
     fn empty(self: T, ) -> starknet::SyscallResult<()>;
 }
 
-#[derive(Copy, Drop, storage_access::StorageAccess, Serde)]
+#[derive(Copy, Drop, starknet::StorageAccess, Serde)]
 struct IContractSafeLibraryDispatcher {
     class_hash: starknet::ClassHash,
 }
@@ -155,7 +155,7 @@ impl IContractSafeLibraryDispatcherImpl of IContractSafeDispatcherTrait<IContrac
 }
 
 
-#[derive(Copy, Drop, storage_access::StorageAccess, Serde)]
+#[derive(Copy, Drop, starknet::StorageAccess, Serde)]
 struct IContractSafeDispatcher {
     contract_address: starknet::ContractAddress,
 }

crates/cairo-lang-starknet/src/plugin/storage_access.rs

@@ -7,7 +7,7 @@ use cairo_lang_syntax::attribute::structured::{
 };
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use cairo_lang_syntax::node::helpers::QueryAttrs;
-use cairo_lang_syntax::node::{ast, TypedSyntaxNode};
+use cairo_lang_syntax::node::{ast, Terminal, TypedSyntaxNode};
 use indoc::formatdoc;
 
 /// Derive the `StorageAccess` trait for structs annotated with `derive(starknet::StorageAccess)`.
@@ -165,6 +165,152 @@ pub fn handle_struct(db: &dyn SyntaxGroup, struct_ast: ast::ItemStruct) -> Plugi
     }
 }
 
+/// Derive the `StorageAccess` trait for structs annotated with `derive(starknet::StorageAccess)`.
+pub fn handle_enum(db: &dyn SyntaxGroup, enum_ast: ast::ItemEnum) -> PluginResult {
+    let enum_name = enum_ast.name(db).as_syntax_node().get_text_without_trivia(db);
+    let mut match_idx = Vec::new();
+    let mut match_idx_at_offset = Vec::new();
+
+    let mut match_value = Vec::new();
+    let mut match_value_at_offset = Vec::new();
+
+    let mut match_size = Vec::new();
+
+    for (i, variant) in enum_ast.variants(db).elements(db).iter().enumerate() {
+        let variant_name = variant.name(db).text(db);
+        let variant_type = match variant.type_clause(db) {
+            ast::OptionTypeClause::Empty(_) => "()".to_string(),
+            ast::OptionTypeClause::TypeClause(tc) => {
+                tc.ty(db).as_syntax_node().get_text_without_trivia(db)
+            }
+        };
+
+        if variant_type == "()" {
+            match_idx.push(format!(
+                "if idx == {i} {{ starknet::SyscallResult::Ok({enum_name}::{variant_name}(())) }}",
+            ));
+            match_idx_at_offset.push(format!(
+                "if idx == {i} {{ starknet::SyscallResult::Ok({enum_name}::{variant_name}(())) }}",
+            ));
+            match_value.push(format!(
+                "{enum_name}::{variant_name}(x) => {{ \
+                 starknet::StorageAccess::write(address_domain, base, {i})?; }}"
+            ));
+            match_value_at_offset.push(format!(
+                "{enum_name}::{variant_name}(x) => {{ \
+                 starknet::StorageAccess::write_at_offset_internal(address_domain, base, offset, \
+                 {i})?; }}"
+            ));
+            match_size.push(format!("{enum_name}::{variant_name}(x) => {{ 1_u8 }}"));
+        } else {
+            match_idx.push(format!(
+                "if idx == {i} {{ \
+                    starknet::SyscallResult::Ok(\
+                        {enum_name}::{variant_name}(\
+                            starknet::StorageAccess::read_at_offset_internal(address_domain, base, 1_u8)?\
+                        )\
+                    ) \
+                }}",
+            ));
+            match_idx_at_offset.push(format!(
+                "if idx == {i} {{ \
+                    starknet::SyscallResult::Ok(\
+                        {enum_name}::{variant_name}(\
+                            starknet::StorageAccess::read_at_offset_internal(address_domain, base, offset + 1_u8)?\
+                        )\
+                    ) \
+                }}",
+            ));
+            match_value.push(format!(
+                "{enum_name}::{variant_name}(x) => {{ \
+                 starknet::StorageAccess::write(address_domain, base, {i})?; \
+                 starknet::StorageAccess::write_at_offset_internal(address_domain, base, 1_u8, \
+                 x)?; }}"
+            ));
+            match_value_at_offset.push(format!(
+                "{enum_name}::{variant_name}(x) => {{ \
+                 starknet::StorageAccess::write_at_offset_internal(address_domain, base, offset, \
+                 {i})?; starknet::StorageAccess::write_at_offset_internal(address_domain, base, \
+                 offset + 1_u8, x)?; }}"
+            ));
+
+            match_size.push(format!(
+                "{enum_name}::{variant_name}(x) => {{ 1_u8 + \
+                 starknet::StorageAccess::size_internal(x) }}"
+            ));
+        }
+    }
+
+    let sa_impl = formatdoc!(
+        "
+        impl StorageAccess{enum_name} of starknet::StorageAccess::<{enum_name}> {{
+            fn read(address_domain: u32, base: starknet::StorageBaseAddress) -> \
+         starknet::SyscallResult<{enum_name}> {{
+                let idx = starknet::StorageAccess::<felt252>::read(address_domain, base)?;
+                {match_idx}
+                else {{
+                    let mut message = Default::default();
+                    message.append('Incorrect index:');
+                    message.append(idx);
+                    starknet::SyscallResult::Err(message)
+                }}
+            }}
+            fn write(address_domain: u32, base: starknet::StorageBaseAddress, value: {enum_name}) \
+         -> starknet::SyscallResult<()> {{
+                match value {{
+                    {match_value}
+                }};
+                starknet::SyscallResult::Ok(())
+            }}
+            fn read_at_offset_internal(address_domain: u32, base: starknet::StorageBaseAddress, \
+         offset: u8) -> starknet::SyscallResult<{enum_name}> {{
+                let idx = \
+         starknet::StorageAccess::<felt252>::read_at_offset_internal(address_domain, base, \
+         offset)?;
+                {match_idx_at_offset}
+                else {{
+                    let mut message = Default::default();
+                    message.append('Incorrect index:');
+                    message.append(idx);
+                    starknet::SyscallResult::Err(message)
+                }}
+            }}
+            #[inline(always)]
+            fn write_at_offset_internal(address_domain: u32, base: starknet::StorageBaseAddress, \
+         offset: u8, value: {enum_name}) -> starknet::SyscallResult<()> {{
+                match value {{
+                    {match_value_at_offset}
+                }};
+                starknet::SyscallResult::Ok(())
+            }}
+            #[inline(always)]
+            fn size_internal(value: {enum_name}) -> u8 {{
+                match value {{
+                    {match_size}
+                }}
+            }}
+        }}",
+        enum_name = enum_name,
+        match_idx = match_idx.join("\n        else "),
+        match_idx_at_offset = match_idx_at_offset.join("\n        else "),
+        match_value = match_value.join(",\n            "),
+        match_value_at_offset = match_value_at_offset.join(",\n            "),
+        match_size = match_size.join(",\n            ")
+    );
+
+    let diagnostics = vec![];
+
+    PluginResult {
+        code: Some(PluginGeneratedFile {
+            name: "storage_access_impl".into(),
+            content: sa_impl,
+            aux_data: DynGeneratedFileAuxData(Arc::new(TrivialPluginAuxData {})),
+        }),
+        diagnostics,
+        remove_original_item: false,
+    }
+}
+
 /// Returns true if the type should be derived as a storage_access.
 pub fn derive_storage_access_needed<T: QueryAttrs>(with_attrs: &T, db: &dyn SyntaxGroup) -> bool {
     with_attrs.query_attr(db, "derive").into_iter().any(|attr| {
@@ -177,8 +323,7 @@ pub fn derive_storage_access_needed<T: QueryAttrs>(with_attrs: &T, db: &dyn Synt
             else {
                 continue;
             };
-            if path.as_syntax_node().get_text_without_trivia(db) == "storage_access::StorageAccess"
-            {
+            if path.as_syntax_node().get_text_without_trivia(db) == "starknet::StorageAccess" {
                 return true;
             }
         }

tests/bug_samples/issue3345.cairo

@@ -1,4 +1,4 @@
-#[derive(Copy, Drop, storage_access::StorageAccess)]
+#[derive(Copy, Drop, starknet::StorageAccess)]
 struct Node {
     left: u128,
     right: u128