Support deriving Introspect for generics (#1082)

dojoengine · Nov 23, 2023 · d8a9a3f · d8a9a3f
1 parent b3f4997
commit d8a9a3f
Show file tree

Hide file tree

Showing 15 changed files with 547 additions and 1,178 deletions.
diff --git a/crates/dojo-core/src/database.cairo b/crates/dojo-core/src/database.cairo
@@ -7,7 +7,9 @@ use poseidon::poseidon_hash_span;
 mod index;
 #[cfg(test)]
 mod index_test;
-mod schema;
+mod introspect;
+#[cfg(test)]
+mod introspect_test;
 mod storage;
 #[cfg(test)]
 mod storage_test;

diff --git a/crates/dojo-core/src/database/schema.cairo → ...s/dojo-core/src/database/introspect.cairo b/crates/dojo-core/src/database/schema.cairo → ...s/dojo-core/src/database/introspect.cairo
@@ -41,7 +41,7 @@ fn serialize_member_type(m: @Ty) -> Span<felt252> {
     serialized.span()
 }
 
-trait SchemaIntrospection<T> {
+trait Introspect<T> {
     fn size() -> usize;
     fn layout(ref layout: Array<u8>);
     fn ty() -> Ty;

diff --git a/crates/dojo-core/src/database/introspect_test.cairo b/crates/dojo-core/src/database/introspect_test.cairo
@@ -0,0 +1,17 @@
+use dojo::database::introspect::Introspect;
+
+#[derive(Drop, Introspect)]
+struct Base {
+    value: u32,
+}
+
+#[derive(Drop, Introspect)]
+struct Generic<T> {
+    value: T,
+}
+
+#[test]
+#[available_gas(2000000)]
+fn test_generic_introspect() {
+    let generic = Generic { value: Base { value: 123 } };
+}
diff --git a/crates/dojo-core/src/model.cairo b/crates/dojo-core/src/model.cairo
@@ -10,5 +10,5 @@ trait Model<T> {
 trait IModel<T> {
     fn name(self: @T) -> felt252;
     fn layout(self: @T) -> Span<felt252>;
-    fn schema(self: @T) -> Span<dojo::database::schema::Member>;
+    fn schema(self: @T) -> Span<dojo::database::introspect::Member>;
 }
diff --git a/crates/dojo-core/src/packing_test.cairo b/crates/dojo-core/src/packing_test.cairo
@@ -5,7 +5,7 @@ use integer::U256BitAnd;
 use option::OptionTrait;
 use debug::PrintTrait;
 use traits::{Into, TryInto};
-use dojo::database::schema::SchemaIntrospection;
+use dojo::database::introspect::Introspect;
 
 #[test]
 #[available_gas(9000000)]

diff --git a/crates/dojo-core/src/world_test.cairo b/crates/dojo-core/src/world_test.cairo
@@ -10,7 +10,7 @@ use starknet::syscalls::deploy_syscall;
 use dojo::benchmarks;
 use dojo::executor::executor;
 use dojo::world::{IWorldDispatcher, IWorldDispatcherTrait, world};
-use dojo::database::schema::SchemaIntrospection;
+use dojo::database::introspect::Introspect;
 use dojo::test_utils::{spawn_test_world, deploy_with_world_address};
 use dojo::benchmarks::{Character, end};
 
@@ -38,7 +38,7 @@ trait Ibar<TContractState> {
 
 #[starknet::contract]
 mod bar {
-    use super::{Foo, IWorldDispatcher, IWorldDispatcherTrait, SchemaIntrospection};
+    use super::{Foo, IWorldDispatcher, IWorldDispatcherTrait, Introspect};
     use super::benchmarks::{Character, Abilities, Stats, Weapon, Sword};
     use traits::Into;
     use starknet::{get_caller_address, ContractAddress};
@@ -60,7 +60,7 @@ mod bar {
 
         fn delete_foo(self: @ContractState) {
             let mut layout = array![];
-            SchemaIntrospection::<Foo>::layout(ref layout);
+            Introspect::<Foo>::layout(ref layout);
             self
                 .world
                 .read()

diff --git a/crates/dojo-lang/src/inline_macros/get.rs b/crates/dojo-lang/src/inline_macros/get.rs
@@ -106,8 +106,7 @@ impl InlineMacroExprPlugin for GetMacro {
 
             builder.add_str(&format!(
                 "\n            let mut __{model}_layout__ = array::ArrayTrait::new();
-                 dojo::database::schema::SchemaIntrospection::<{model}>::layout(ref \
-                 __{model}_layout__);
+                 dojo::database::introspect::Introspect::<{model}>::layout(ref __{model}_layout__);
                  let mut __{model}_layout_clone__ = __{model}_layout__.clone();
                  let mut __{model}_layout_span__ = array::ArrayTrait::span(@__{model}_layout__);
                  let mut __{model}_layout_clone_span__ = \

diff --git a/crates/dojo-lang/src/introspect.rs b/crates/dojo-lang/src/introspect.rs
@@ -2,7 +2,9 @@ use std::collections::HashMap;
 
 use cairo_lang_defs::patcher::RewriteNode;
 use cairo_lang_defs::plugin::PluginDiagnostic;
-use cairo_lang_syntax::node::ast::{Expr, ItemEnum, ItemStruct, OptionTypeClause};
+use cairo_lang_syntax::node::ast::{
+    Expr, GenericParam, ItemEnum, ItemStruct, OptionTypeClause, OptionWrappedGenericParamList,
+};
 use cairo_lang_syntax::node::db::SyntaxGroup;
 use cairo_lang_syntax::node::helpers::QueryAttrs;
 use cairo_lang_syntax::node::{Terminal, TypedSyntaxNode};
@@ -57,43 +59,41 @@ pub fn handle_introspect_struct(db: &dyn SyntaxGroup, struct_ast: ItemStruct) ->
             if primitive_sizes.get(&ty).is_some() {
                 // It's a primitive type
                 member_types.push(format!(
-                    "
-                    dojo::database::schema::serialize_member(@dojo::database::schema::Member {{
-                        name: '{name}',
-                        ty: dojo::database::schema::Ty::Primitive('{ty}'),
-                        attrs: array![{}].span()
-                    }})\n",
+                    "dojo::database::introspect::serialize_member(@\
+                     dojo::database::introspect::Member {{
+                name: '{name}',
+                ty: dojo::database::introspect::Ty::Primitive('{ty}'),
+                attrs: array![{}].span()
+            }})",
                     attrs.join(","),
                 ));
             } else {
                 // It's a custom struct/enum
                 member_types.push(format!(
-                    "
-                    dojo::database::schema::serialize_member(@dojo::database::schema::Member {{
-                        name: '{name}',
-                        ty: dojo::database::schema::SchemaIntrospection::<{ty}>::ty(),
-                        attrs: array![{}].span()
-                    }})\n",
+                    "dojo::database::introspect::serialize_member(@\
+                     dojo::database::introspect::Member {{
+                name: '{name}',
+                ty: dojo::database::introspect::Introspect::<{ty}>::ty(),
+                attrs: array![{}].span()
+            }})",
                     attrs.join(","),
                 ));
             }
 
             Member { name, ty, key }
         })
         .collect::<_>();
-    drop(primitive_sizes);
 
     let type_ty = format!(
-        "
-        dojo::database::schema::Ty::Struct(dojo::database::schema::Struct {{
+        "dojo::database::introspect::Ty::Struct(dojo::database::introspect::Struct {{
             name: '{name}',
             attrs: array![].span(),
             children: array![{}].span()
         }})",
-        member_types.join(",\n")
+        member_types.join(", ")
     );
 
-    handle_introspect_internal(db, name, vec![], 0, type_ty, members)
+    handle_introspect_internal(db, name, struct_ast.generic_params(db), vec![], 0, type_ty, members)
 }
 
 /// A handler for Dojo code derives Introspect for an enum
@@ -108,6 +108,7 @@ pub fn handle_introspect_enum(
     enum_ast: ItemEnum,
 ) -> RewriteNode {
     let name = enum_ast.name(db).text(db).into();
+
     let variant_type = enum_ast.variants(db).elements(db).first().unwrap().type_clause(db);
     let variant_type_text = variant_type.as_syntax_node().get_text(db);
     let variant_type_text = variant_type_text.trim();
@@ -121,8 +122,8 @@ pub fn handle_introspect_enum(
                 variant_type_arr.push((
                     // Not using Ty right now, but still keeping it for later.
                     format!(
-                        "dojo::database::schema::serialize_member_type(
-                            @dojo::database::schema::Ty::Primitive('{}')
+                        "dojo::database::introspect::serialize_member_type(
+                            @dojo::database::introspect::Ty::Primitive('{}')
                         )",
                         ty_name
                     ),
@@ -134,8 +135,8 @@ pub fn handle_introspect_enum(
             variant_type_arr.push((
                 // Not using Ty right now, but still keeping it for later.
                 format!(
-                    "dojo::database::schema::serialize_member_type(
-                        @dojo::database::schema::SchemaIntrospection::<{}>::ty()
+                    "dojo::database::introspect::serialize_member_type(
+                        @dojo::database::introspect::Introspect::<{}>::ty()
                     )",
                     ty_name
                 ),
@@ -170,32 +171,27 @@ pub fn handle_introspect_enum(
 
         // @TODO: Prepare type struct
         arms_ty.push(format!(
-            "
-            (
-                '{member_name}',
-                dojo::database::schema::serialize_member_type(
-                @dojo::database::schema::Ty::Tuple(array![{}].span()))
-            )",
+            "(
+                    '{member_name}',
+                    dojo::database::introspect::serialize_member_type(
+                    @dojo::database::introspect::Ty::Tuple(array![{}].span()))
+                )",
             if !variant_type_arr.is_empty() {
                 let ty_cairo: Vec<_> =
                     variant_type_arr.iter().map(|(ty_cairo, _)| ty_cairo.to_string()).collect();
-                // format!("'{}'", &ty_cairo.join("', '"))
-                ty_cairo.join(",\n")
+                ty_cairo.join(", ")
             } else {
                 "".to_string()
             }
         ));
     });
 
     let type_ty = format!(
-        "
-        dojo::database::schema::Ty::Enum(
-            dojo::database::schema::Enum {{
+        "dojo::database::introspect::Ty::Enum(
+            dojo::database::introspect::Enum {{
                 name: '{name}',
                 attrs: array![].span(),
-                children: array![
-                {}
-                ].span()
+                children: array![{}].span()
             }}
         )",
         arms_ty.join(",\n")
@@ -204,12 +200,21 @@ pub fn handle_introspect_enum(
     // Enums have 1 size and 8 bit layout by default
     let layout = vec![RewriteNode::Text("layout.append(8);\n".into())];
     let size_precompute = 1;
-    handle_introspect_internal(db, name, layout, size_precompute, type_ty, members)
+    handle_introspect_internal(
+        db,
+        name,
+        enum_ast.generic_params(db),
+        layout,
+        size_precompute,
+        type_ty,
+        members,
+    )
 }
 
 fn handle_introspect_internal(
-    _db: &dyn SyntaxGroup,
+    db: &dyn SyntaxGroup,
     name: String,
+    generics: OptionWrappedGenericParamList,
     mut layout: Vec<RewriteNode>,
     mut size_precompute: usize,
     type_ty: String,
@@ -218,6 +223,30 @@ fn handle_introspect_internal(
     let mut size = vec![];
     let primitive_sizes = primitive_type_introspection();
 
+    let generics = if let OptionWrappedGenericParamList::WrappedGenericParamList(params) = generics
+    {
+        params
+            .generic_params(db)
+            .elements(db)
+            .iter()
+            .filter_map(|el| {
+                if let GenericParam::Type(typ) = el {
+                    Some(typ.name(db).text(db).to_string())
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>()
+    } else {
+        vec![]
+    };
+
+    let generic_impls = generics
+        .iter()
+        .map(|g| format!("{g}, impl {g}Introspect: dojo::database::introspect::Introspect<{g}>"))
+        .collect::<Vec<_>>()
+        .join(", ");
+
     members.iter().for_each(|m| {
         let primitive_intro = primitive_sizes.get(&m.ty);
         let mut attrs = vec![];
@@ -237,12 +266,9 @@ fn handle_introspect_internal(
             if m.key {
                 attrs.push("'key'");
             } else {
-                size.push(format!(
-                    "dojo::database::schema::SchemaIntrospection::<{}>::size()",
-                    m.ty,
-                ));
+                size.push(format!("dojo::database::introspect::Introspect::<{}>::size()", m.ty,));
                 layout.push(RewriteNode::Text(format!(
-                    "dojo::database::schema::SchemaIntrospection::<{}>::layout(ref layout);\n",
+                    "dojo::database::introspect::Introspect::<{}>::layout(ref layout);\n",
                     m.ty
                 )));
             }
@@ -255,26 +281,28 @@ fn handle_introspect_internal(
 
     RewriteNode::interpolate_patched(
         "
-        impl $name$SchemaIntrospection of dojo::database::schema::SchemaIntrospection<$name$> {
-            
-            #[inline(always)]
-            fn size() -> usize {
-                $size$
-            }
+impl $name$Introspect<$generics$> of \
+         dojo::database::introspect::Introspect<$name$<$generics_types$>> {
+    #[inline(always)]
+    fn size() -> usize {
+        $size$
+    }
 
-            #[inline(always)]
-            fn layout(ref layout: Array<u8>) {
-                $layout$
-            }
+    #[inline(always)]
+    fn layout(ref layout: Array<u8>) {
+        $layout$
+    }
 
-            #[inline(always)]
-            fn ty() -> dojo::database::schema::Ty {
-                $ty$
-            }
-        }
+    #[inline(always)]
+    fn ty() -> dojo::database::introspect::Ty {
+        $ty$
+    }
+}
         ",
         &UnorderedHashMap::from([
             ("name".to_string(), RewriteNode::Text(name)),
+            ("generics".to_string(), RewriteNode::Text(generic_impls)),
+            ("generics_types".to_string(), RewriteNode::Text(generics.join(", "))),
             ("size".to_string(), RewriteNode::Text(size.join(" + "))),
             ("layout".to_string(), RewriteNode::new_modified(layout)),
             ("ty".to_string(), RewriteNode::Text(type_ty)),