Fixed generated functions in trait-fns.

commit-id:cd467efe
starkware-libs · Jan 14, 2025 · b45d5db · b45d5db
1 parent f800627
commit b45d5db
Show file tree

Hide file tree

Showing 2 changed files with 98 additions and 1 deletion.
diff --git a/crates/cairo-lang-lowering/src/lower/test_data/loop b/crates/cairo-lang-lowering/src/lower/test_data/loop
@@ -1527,3 +1527,82 @@ Statements:
   (v27: core::panics::PanicResult::<(test::A, core::felt252, ())>) <- PanicResult::Err(v26)
 End:
   Return(v6, v7, v27)
+
+//! > ==========================================================================
+
+//! > Test default implementation with loop.
+
+//! > test_runner_name
+test_generated_function
+
+//! > function
+fn foo() {
+    MyTrait::impl_in_trait();
+}
+
+//! > function_name
+foo
+
+//! > module_code
+trait MyTrait {
+    fn impl_in_impl(x: u8) -> bool;
+    fn impl_in_trait() -> u8 {
+        let mut i = 0;
+        loop {
+            if Self::impl_in_impl(i) {
+                break;
+            }
+            i += 1;
+        };
+        i
+    }
+}
+
+impl MyImpl of MyTrait {
+    fn impl_in_impl(x: u8) -> bool {
+        x == 30
+    }
+}
+
+//! > expected_diagnostics
+
+//! > semantic_diagnostics
+
+//! > lowering
+Main:
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::integer::u8) <- test::MyImpl::impl_in_trait()
+  (v1: ()) <- struct_construct()
+End:
+  Return(v1)
+
+
+Final lowering:
+Parameters: v0: core::RangeCheck, v1: core::gas::GasBuiltin
+blk0 (root):
+Statements:
+  (v2: core::integer::u8) <- 0
+  (v3: core::RangeCheck, v4: core::gas::GasBuiltin, v5: core::panics::PanicResult::<(core::integer::u8, ())>) <- test::MyImpl::impl_in_trait[expr9](v0, v1, v2)
+End:
+  Match(match_enum(v5) {
+    PanicResult::Ok(v6) => blk1,
+    PanicResult::Err(v7) => blk2,
+  })
+
+blk1:
+Statements:
+  (v8: ()) <- struct_construct()
+  (v9: ((),)) <- struct_construct(v8)
+  (v10: core::panics::PanicResult::<((),)>) <- PanicResult::Ok(v9)
+End:
+  Return(v3, v4, v10)
+
+blk2:
+Statements:
+  (v11: core::panics::PanicResult::<((),)>) <- PanicResult::Err(v7)
+End:
+  Return(v3, v4, v11)
+
+//! > lowering_diagnostics
diff --git a/crates/cairo-lang-semantic/src/substitution.rs b/crates/cairo-lang-semantic/src/substitution.rs
@@ -464,7 +464,7 @@ add_basic_rewrites!(
     <'a>,
     SubstitutionRewriter<'a>,
     DiagnosticAdded,
-    @exclude TypeId TypeLongId ImplId ImplLongId ConstValue GenericFunctionId
+    @exclude TypeId TypeLongId ImplId ImplLongId ConstValue GenericFunctionId GenericFunctionWithBodyId
 );
 
 impl SemanticRewriter<TypeId, DiagnosticAdded> for SubstitutionRewriter<'_> {
@@ -633,3 +633,21 @@ impl SemanticRewriter<GenericFunctionId, DiagnosticAdded> for SubstitutionRewrit
         value.default_rewrite(self)
     }
 }
+impl SemanticRewriter<GenericFunctionWithBodyId, DiagnosticAdded> for SubstitutionRewriter<'_> {
+    fn internal_rewrite(&mut self, value: &mut GenericFunctionWithBodyId) -> Maybe<RewriteResult> {
+        if let GenericFunctionWithBodyId::Trait(id) = value {
+            if let Some(self_impl) = &self.substitution.self_impl {
+                if let ImplLongId::Concrete(concrete_impl_id) = self_impl.lookup_intern(self.db) {
+                    if id.concrete_trait(self.db.upcast()) == self_impl.concrete_trait(self.db)? {
+                        *value = GenericFunctionWithBodyId::Impl(ImplGenericFunctionWithBodyId {
+                            concrete_impl_id,
+                            function_body: ImplFunctionBodyId::Trait(id.trait_function(self.db)),
+                        });
+                        return Ok(RewriteResult::Modified);
+                    }
+                }
+            }
+        }
+        value.default_rewrite(self)
+    }
+}