fix: Support deriving Print for enums

crates/dojo-lang/src/introspect.rs

@@ -99,7 +99,7 @@ pub fn handle_introspect_struct(db: &dyn SyntaxGroup, struct_ast: ItemStruct) ->
 /// A handler for Dojo code derives Introspect for an enum
 /// Parameters:
 /// * db: The semantic database.
-/// * struct_ast: The AST of the struct.
+/// * enum_ast: The AST of the enum.
 /// Returns:
 /// * A RewriteNode containing the generated code.
 pub fn handle_introspect_enum(

crates/dojo-lang/src/plugin.rs

@@ -33,7 +33,7 @@ use crate::inline_macros::get::GetMacro;
 use crate::inline_macros::set::SetMacro;
 use crate::introspect::{handle_introspect_enum, handle_introspect_struct};
 use crate::model::handle_model_struct;
-use crate::print::derive_print;
+use crate::print::{handle_print_enum, handle_print_struct};
 
 const DOJO_CONTRACT_ATTR: &str = "dojo::contract";
 
@@ -303,6 +303,7 @@ impl MacroPlugin for BuiltinDojoPlugin {
                                     enum_ast.clone(),
                                 ));
                             }
+                            "Print" => rewrite_nodes.push(handle_print_enum(db, enum_ast.clone())),
                             _ => continue,
                         }
                     }
@@ -379,7 +380,7 @@ impl MacroPlugin for BuiltinDojoPlugin {
                                 diagnostics.extend(model_diagnostics);
                             }
                             "Print" => {
-                                rewrite_nodes.push(derive_print(db, struct_ast.clone()));
+                                rewrite_nodes.push(handle_print_struct(db, struct_ast.clone()));
                             }
                             "Introspect" => {
                                 rewrite_nodes

crates/dojo-lang/src/plugin_test_data/print

@@ -32,6 +32,13 @@ struct Player {
     name: felt252, 
 }
 
+#[derive(Print, Copy, Drop, Serde)]
+enum Enemy {
+    Unknown,
+    Bot: felt252,
+    OtherPlayer: ContractAddress,
+}
+
 //! > generated_cairo_code
 use serde::Serde;
 
@@ -84,6 +91,7 @@ struct Player {
 
     name: felt252, 
 }
+
 #[cfg(test)]
 impl PlayerPrintImpl of debug::PrintTrait<Player> {
     fn print(self: Player) {
@@ -126,6 +134,13 @@ struct Player {
 
     name: felt252, 
 }
+
+#[derive(Print, Copy, Drop, Serde)]
+enum Enemy {
+    Unknown,
+    Bot: felt252,
+    OtherPlayer: ContractAddress,
+}
 impl PositionCopy of Copy::<Position>;
 impl PositionDrop of Drop::<Position>;
 impl PositionSerde of Serde::<Position> {
@@ -168,3 +183,23 @@ impl PlayerSerde of Serde::<Player> {
         })
     }
 }
+impl EnemyCopy of Copy::<Enemy>;
+impl EnemyDrop of Drop::<Enemy>;
+impl EnemySerde of Serde::<Enemy> {
+    fn serialize(self: @Enemy, ref output: array::Array<felt252>) {
+        match self {
+            Enemy::Unknown(x) => { serde::Serde::serialize(@0, ref output); serde::Serde::serialize(x, ref output); },
+            Enemy::Bot(x) => { serde::Serde::serialize(@1, ref output); serde::Serde::serialize(x, ref output); },
+            Enemy::OtherPlayer(x) => { serde::Serde::serialize(@2, ref output); serde::Serde::serialize(x, ref output); },
+        }
+    }
+    fn deserialize(ref serialized: array::Span<felt252>) -> Option<Enemy> {
+        let idx: felt252 = serde::Serde::deserialize(ref serialized)?;
+        Option::Some(
+            if idx == 0 { Enemy::Unknown(serde::Serde::deserialize(ref serialized)?) }
+            else if idx == 1 { Enemy::Bot(serde::Serde::deserialize(ref serialized)?) }
+            else if idx == 2 { Enemy::OtherPlayer(serde::Serde::deserialize(ref serialized)?) }
+            else { return Option::None; }
+        )
+    }
+}

crates/dojo-lang/src/print.rs

@@ -1,5 +1,5 @@
 use cairo_lang_defs::patcher::RewriteNode;
-use cairo_lang_syntax::node::ast::ItemStruct;
+use cairo_lang_syntax::node::ast::{ItemEnum, ItemStruct};
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use cairo_lang_syntax::node::{Terminal, TypedSyntaxNode};
 use cairo_lang_utils::unordered_hash_map::UnorderedHashMap;
@@ -10,7 +10,7 @@ use cairo_lang_utils::unordered_hash_map::UnorderedHashMap;
 /// * struct_ast: The AST of the model struct.
 /// Returns:
 /// * A RewriteNode containing the generated code.
-pub fn derive_print(db: &dyn SyntaxGroup, struct_ast: ItemStruct) -> RewriteNode {
+pub fn handle_print_struct(db: &dyn SyntaxGroup, struct_ast: ItemStruct) -> RewriteNode {
     let prints: Vec<_> = struct_ast
         .members(db)
         .elements(db)
@@ -26,7 +26,7 @@ pub fn derive_print(db: &dyn SyntaxGroup, struct_ast: ItemStruct) -> RewriteNode
 
     RewriteNode::interpolate_patched(
         "#[cfg(test)]
-            impl $type_name$PrintImpl of debug::PrintTrait<$type_name$> {
+            impl $type_name$StructPrintImpl of debug::PrintTrait<$type_name$> {
                 fn print(self: $type_name$) {
                     $print$
                 }
@@ -40,3 +40,42 @@ pub fn derive_print(db: &dyn SyntaxGroup, struct_ast: ItemStruct) -> RewriteNode
         ]),
     )
 }
+
+/// Derives PrintTrait for an enum.
+/// Parameters:
+/// * db: The semantic database.
+/// * enum_ast: The AST of the model enum.
+/// Returns:
+/// * A RewriteNode containing the generated code.
+pub fn handle_print_enum(db: &dyn SyntaxGroup, enum_ast: ItemEnum) -> RewriteNode {
+    let enum_name = enum_ast.name(db).text(db);
+    let prints: Vec<_> = enum_ast
+        .variants(db)
+        .elements(db)
+        .iter()
+        .map(|m| {
+            format!(
+                "{}::{}(value) => {{ debug::PrintTrait::print('{}'); \
+                 debug::PrintTrait::print(value); }}",
+                enum_name,
+                m.name(db).text(db).to_string(),
+                m.name(db).text(db).to_string()
+            )
+        })
+        .collect();
+
+    RewriteNode::interpolate_patched(
+        "#[cfg(test)]
+            impl $type_name$EnumPrintImpl of debug::PrintTrait<$type_name$> {
+                fn print(self: $type_name$) {
+                    match self {
+                        $print$
+                    }
+                }
+            }",
+        &UnorderedHashMap::from([
+            ("type_name".to_string(), RewriteNode::new_trimmed(enum_ast.name(db).as_syntax_node())),
+            ("print".to_string(), RewriteNode::Text(prints.join("\n"))),
+        ]),
+    )
+}