From 8681d08e7b3951405904492337c63adfed3bad68 Mon Sep 17 00:00:00 2001
From: Ori Ziv <oriz@starkware.co>
Date: Wed, 8 Jan 2025 16:57:48 +0200
Subject: [PATCH] Made associated items not replaced if they aren't of the
 `Self` trait.

---
 .../src/items/tests/trait_default_fn          | 24 ++++++++
 .../cairo-lang-semantic/src/substitution.rs   | 58 ++++++++++---------
 2 files changed, 55 insertions(+), 27 deletions(-)

diff --git a/crates/cairo-lang-semantic/src/items/tests/trait_default_fn b/crates/cairo-lang-semantic/src/items/tests/trait_default_fn
index 30cedcfb5fd..3c8406d1c0d 100644
--- a/crates/cairo-lang-semantic/src/items/tests/trait_default_fn
+++ b/crates/cairo-lang-semantic/src/items/tests/trait_default_fn
@@ -83,3 +83,27 @@ error: Parameter of trait function `MyTrait::foo` can't be defined as mutable.
  --> lib.cairo:2:12
     fn foo(mut x: u32) -> u32 {
            ^^^
+
+//! > ==========================================================================
+
+//! > Test default implementation with associated types.
+
+//! > test_runner_name
+test_function_diagnostics(expect_diagnostics: false)
+
+//! > function
+fn foo() {}
+
+//! > function_name
+foo
+
+//! > module_code
+pub trait MyTrait {
+    type Associated;
+    fn impl_required() -> Option<Self::Associated>;
+    fn default_impl() -> bool {
+        Self::impl_required().is_some()
+    }
+}
+
+//! > expected_diagnostics
diff --git a/crates/cairo-lang-semantic/src/substitution.rs b/crates/cairo-lang-semantic/src/substitution.rs
index e42b1d488aa..0ae2177e339 100644
--- a/crates/cairo-lang-semantic/src/substitution.rs
+++ b/crates/cairo-lang-semantic/src/substitution.rs
@@ -509,14 +509,16 @@ impl SemanticRewriter<TypeLongId, DiagnosticAdded> for SubstitutionRewriter<'_>
             }
             TypeLongId::TraitType(trait_type_id) => {
                 if let Some(self_impl) = &self.substitution.self_impl {
-                    assert_eq!(
-                        trait_type_id.trait_id(self.db.upcast()),
-                        self_impl.concrete_trait(self.db)?.trait_id(self.db)
-                    );
-                    let impl_type_id = ImplTypeId::new(*self_impl, *trait_type_id, self.db);
-                    *value =
-                        self.db.impl_type_concrete_implized(impl_type_id)?.lookup_intern(self.db);
-                    return Ok(RewriteResult::Modified);
+                    if trait_type_id.trait_id(self.db.upcast())
+                        == self_impl.concrete_trait(self.db)?.trait_id(self.db)
+                    {
+                        let impl_type_id = ImplTypeId::new(*self_impl, *trait_type_id, self.db);
+                        *value = self
+                            .db
+                            .impl_type_concrete_implized(impl_type_id)?
+                            .lookup_intern(self.db);
+                        return Ok(RewriteResult::Modified);
+                    }
                 }
             }
             _ => {}
@@ -550,18 +552,18 @@ impl SemanticRewriter<ConstValue, DiagnosticAdded> for SubstitutionRewriter<'_>
             }
             ConstValue::TraitConstant(trait_constant_id) => {
                 if let Some(self_impl) = &self.substitution.self_impl {
-                    assert_eq!(
-                        trait_constant_id.trait_id(self.db.upcast()),
-                        self_impl.concrete_trait(self.db)?.trait_id(self.db)
-                    );
-                    let impl_const_id =
-                        ImplConstantId::new(*self_impl, *trait_constant_id, self.db);
-                    *value = self
-                        .db
-                        .impl_constant_concrete_implized_value(impl_const_id)?
-                        .lookup_intern(self.db);
+                    if trait_constant_id.trait_id(self.db.upcast())
+                        == self_impl.concrete_trait(self.db)?.trait_id(self.db)
+                    {
+                        let impl_const_id =
+                            ImplConstantId::new(*self_impl, *trait_constant_id, self.db);
+                        *value = self
+                            .db
+                            .impl_constant_concrete_implized_value(impl_const_id)?
+                            .lookup_intern(self.db);
 
-                    return Ok(RewriteResult::Modified);
+                        return Ok(RewriteResult::Modified);
+                    }
                 }
             }
             _ => {}
@@ -595,15 +597,17 @@ impl SemanticRewriter<ImplLongId, DiagnosticAdded> for SubstitutionRewriter<'_>
             }
             ImplLongId::TraitImpl(trait_impl_id) => {
                 if let Some(self_impl) = &self.substitution.self_impl {
-                    assert_eq!(
-                        trait_impl_id.trait_id(self.db.upcast()),
-                        self_impl.concrete_trait(self.db)?.trait_id(self.db)
-                    );
-                    let impl_impl_id = ImplImplId::new(*self_impl, *trait_impl_id, self.db);
-                    *value =
-                        self.db.impl_impl_concrete_implized(impl_impl_id)?.lookup_intern(self.db);
+                    if trait_impl_id.trait_id(self.db.upcast())
+                        == self_impl.concrete_trait(self.db)?.trait_id(self.db)
+                    {
+                        let impl_impl_id = ImplImplId::new(*self_impl, *trait_impl_id, self.db);
+                        *value = self
+                            .db
+                            .impl_impl_concrete_implized(impl_impl_id)?
+                            .lookup_intern(self.db);
 
-                    return Ok(RewriteResult::Modified);
+                        return Ok(RewriteResult::Modified);
+                    }
                 }
             }
             _ => {}