fix: allow the use of self and world in an exclusive manner (#2063)

* fix: allow the use of self and world in an exclusive manner

* fix: fmt

* fix: cairo fmt

* fix: ensure a warning is emitted with generate trait and world

* fix: cairo fmt

crates/dojo-lang/src/contract.rs

@@ -4,8 +4,10 @@ use cairo_lang_defs::patcher::{PatchBuilder, RewriteNode};
 use cairo_lang_defs::plugin::{
     DynGeneratedFileAuxData, PluginDiagnostic, PluginGeneratedFile, PluginResult,
 };
+use cairo_lang_diagnostics::Severity;
 use cairo_lang_syntax::node::ast::MaybeModuleBody;
 use cairo_lang_syntax::node::db::SyntaxGroup;
+use cairo_lang_syntax::node::helpers::QueryAttrs;
 use cairo_lang_syntax::node::{ast, ids, Terminal, TypedStablePtr, TypedSyntaxNode};
 use cairo_lang_utils::unordered_hash_map::UnorderedHashMap;
 use dojo_types::system::Dependency;
@@ -314,7 +316,7 @@ impl DojoContract {
         param_list: ast::ParamList,
         fn_diagnostic_item: ids::SyntaxStablePtrId,
     ) -> (String, bool) {
-        self_param::check_parameter(db, &param_list, fn_diagnostic_item, &mut self.diagnostics);
+        let is_self_used = self_param::check_parameter(db, &param_list);
 
         let world_injection = world_param::parse_world_injection(
             db,
@@ -323,6 +325,14 @@ impl DojoContract {
             &mut self.diagnostics,
         );
 
+        if is_self_used && world_injection != WorldParamInjectionKind::None {
+            self.diagnostics.push(PluginDiagnostic {
+                stable_ptr: fn_diagnostic_item,
+                message: "You cannot use `self` and `world` parameters together.".to_string(),
+                severity: Severity::Error,
+            });
+        }
+
         let mut params = param_list
             .elements(db)
             .iter()
@@ -340,7 +350,12 @@ impl DojoContract {
             .collect::<Vec<_>>();
 
         match world_injection {
-            WorldParamInjectionKind::None | WorldParamInjectionKind::View => {
+            WorldParamInjectionKind::None => {
+                if !is_self_used {
+                    params.insert(0, "self: @ContractState".to_string());
+                }
+            }
+            WorldParamInjectionKind::View => {
                 params.insert(0, "self: @ContractState".to_string());
             }
             WorldParamInjectionKind::External => {
@@ -372,10 +387,13 @@ impl DojoContract {
     ///  * removing `world` if present as first parameter (self excluded),
     ///  * adding `let world = self.world_dispatcher.read();` statement at the beginning of the
     ///    function to restore the removed `world` parameter.
+    ///  * if `has_generate_trait` is true, the implementation containing the function has the
+    ///    #[generate_trait] attribute.
     pub fn rewrite_function(
         &mut self,
         db: &dyn SyntaxGroup,
         fn_ast: ast::FunctionWithBody,
+        has_generate_trait: bool,
     ) -> Vec<RewriteNode> {
         let mut rewritten_fn = RewriteNode::from_ast(&fn_ast);
 
@@ -385,6 +403,16 @@ impl DojoContract {
             fn_ast.stable_ptr().untyped(),
         );
 
+        if has_generate_trait && was_world_injected {
+            self.diagnostics.push(PluginDiagnostic {
+                stable_ptr: fn_ast.stable_ptr().untyped(),
+                message: "You cannot use `world` and `#[generate_trait]` together. Use `self` \
+                          instead."
+                    .to_string(),
+                severity: Severity::Error,
+            });
+        }
+
         // We always rewrite the params as the self parameter is added based on the
         // world mutability.
         let rewritten_params = rewritten_fn
@@ -409,14 +437,17 @@ impl DojoContract {
 
     /// Rewrites all the functions of a Impl block.
     fn rewrite_impl(&mut self, db: &dyn SyntaxGroup, impl_ast: ast::ItemImpl) -> Vec<RewriteNode> {
+        let generate_attrs = impl_ast.attributes(db).query_attr(db, "generate_trait");
+        let has_generate_trait = !generate_attrs.is_empty();
+
         if let ast::MaybeImplBody::Some(body) = impl_ast.body(db) {
             let body_nodes: Vec<_> = body
                 .items(db)
                 .elements(db)
                 .iter()
                 .flat_map(|el| {
                     if let ast::ImplItem::Function(fn_ast) = el {
-                        return self.rewrite_function(db, fn_ast.clone());
+                        return self.rewrite_function(db, fn_ast.clone(), has_generate_trait);
                     }
                     vec![RewriteNode::Copied(el.as_syntax_node())]
                 })

crates/dojo-lang/src/interface.rs

@@ -93,7 +93,7 @@ impl DojoInterface {
             .map(|e| e.as_syntax_node().get_text(db))
             .collect::<Vec<_>>();
 
-        self_param::check_parameter(db, &param_list, diagnostic_item, &mut self.diagnostics);
+        let is_self_used = self_param::check_parameter(db, &param_list);
 
         let world_injection = world_param::parse_world_injection(
             db,
@@ -102,9 +102,19 @@ impl DojoInterface {
             &mut self.diagnostics,
         );
 
+        if is_self_used && world_injection != WorldParamInjectionKind::None {
+            self.diagnostics.push(PluginDiagnostic {
+                stable_ptr: diagnostic_item,
+                message: "You cannot use `self` and `world` parameters together.".to_string(),
+                severity: Severity::Error,
+            });
+        }
+
         match world_injection {
             WorldParamInjectionKind::None => {
-                params.insert(0, "self: @TContractState".to_string());
+                if !is_self_used {
+                    params.insert(0, "self: @TContractState".to_string());
+                }
             }
             WorldParamInjectionKind::View => {
                 params.remove(0);

crates/dojo-lang/src/plugin.rs

@@ -126,8 +126,6 @@ impl BuiltinDojoPlugin {
     ) -> PluginResult {
         PluginResult {
             code: None,
-            // All diagnostics are for now error. Severity may be moved as argument
-            // if warnings are required in this file.
             diagnostics: vec![PluginDiagnostic { stable_ptr, message, severity: Severity::Error }],
             remove_original_item: false,
         }