From b79f3689c2c2bb26bb9f85bf21f27fd76ee8e3a3 Mon Sep 17 00:00:00 2001
From: Ori Ziv <oriz@starkware.co>
Date: Wed, 8 Jan 2025 01:49:54 +0200
Subject: [PATCH] Added better support for concrete trait by output type.

---
 .../cairo-lang-semantic/src/expr/inference.rs | 43 +++++++++++--------
 .../src/expr/test_data/closure                |  2 +-
 .../src/expr/test_data/inference              | 23 ++++++++++
 3 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/crates/cairo-lang-semantic/src/expr/inference.rs b/crates/cairo-lang-semantic/src/expr/inference.rs
index 17d0df63c06..d4460c427d2 100644
--- a/crates/cairo-lang-semantic/src/expr/inference.rs
+++ b/crates/cairo-lang-semantic/src/expr/inference.rs
@@ -913,29 +913,34 @@ impl<'db> Inference<'db> {
         let concrete_trait_id = self.rewrite(concrete_trait_id).no_err();
         enrich_lookup_context(self.db, concrete_trait_id, &mut lookup_context);
 
-        // Don't try to resolve impls if the first generic param is a variable.
-        let generic_args = concrete_trait_id.generic_args(self.db);
-        match generic_args.first() {
-            Some(GenericArgumentId::Type(ty)) => {
-                if let TypeLongId::Var(_) = ty.lookup_intern(self.db) {
-                    // Don't try to infer such impls.
-                    return Ok(SolutionSet::Ambiguous(Ambiguity::WillNotInfer(concrete_trait_id)));
+        // If has any type assignments, we can try to resolve it, even if first generic arg is a
+        // var.
+        if !impl_var_trait_item_mappings
+            .types
+            .iter()
+            .any(|(_, ty)| ty.is_var_free(self.db) && !ty.is_missing(self.db))
+        {
+            // Don't try to resolve impls if the first generic param is a variable.
+            let dont_infer = Ok(SolutionSet::Ambiguous(Ambiguity::WillNotInfer(concrete_trait_id)));
+            match concrete_trait_id.generic_args(self.db).first() {
+                Some(GenericArgumentId::Type(ty))
+                    if matches!(ty.lookup_intern(self.db), TypeLongId::Var(_)) =>
+                {
+                    return dont_infer;
                 }
-            }
-            Some(GenericArgumentId::Impl(imp)) => {
-                // Don't try to infer such impls.
-                if let ImplLongId::ImplVar(_) = imp.lookup_intern(self.db) {
-                    return Ok(SolutionSet::Ambiguous(Ambiguity::WillNotInfer(concrete_trait_id)));
+                Some(GenericArgumentId::Impl(imp))
+                    if matches!(imp.lookup_intern(self.db), ImplLongId::ImplVar(_)) =>
+                {
+                    return dont_infer;
                 }
-            }
-            Some(GenericArgumentId::Constant(const_value)) => {
-                if let ConstValue::Var(_, _) = const_value.lookup_intern(self.db) {
-                    // Don't try to infer such impls.
-                    return Ok(SolutionSet::Ambiguous(Ambiguity::WillNotInfer(concrete_trait_id)));
+                Some(GenericArgumentId::Constant(const_value))
+                    if matches!(const_value.lookup_intern(self.db), ConstValue::Var(_, _)) =>
+                {
+                    return dont_infer;
                 }
+                _ => {}
             }
-            _ => {}
-        };
+        }
         let (canonical_trait, canonicalizer) = CanonicalTrait::canonicalize(
             self.db,
             self.inference_id,
diff --git a/crates/cairo-lang-semantic/src/expr/test_data/closure b/crates/cairo-lang-semantic/src/expr/test_data/closure
index ef7d8096174..de2c696ebaa 100644
--- a/crates/cairo-lang-semantic/src/expr/test_data/closure
+++ b/crates/cairo-lang-semantic/src/expr/test_data/closure
@@ -390,7 +390,7 @@ fn baz<T, +core::ops::FnOnce<T, (u32,)>>(c: T) -> core::ops::FnOnce::<T, (u32,)>
 }
 
 //! > expected_diagnostics
-error: Trait has no implementation in context: core::ops::function::Fn::<core::felt252, ()>.
+error: Inference cycle detected
  --> lib.cairo:7:19
     let _x: u32 = bar();
                   ^^^^^
diff --git a/crates/cairo-lang-semantic/src/expr/test_data/inference b/crates/cairo-lang-semantic/src/expr/test_data/inference
index d633ca3a125..180a6c19595 100644
--- a/crates/cairo-lang-semantic/src/expr/test_data/inference
+++ b/crates/cairo-lang-semantic/src/expr/test_data/inference
@@ -755,3 +755,26 @@ error: Type annotations needed. Failed to infer ?0.
  --> lib.cairo:1:13
 fn foo() -> Option {
             ^^^^^^
+
+//! > ==========================================================================
+
+//! > Impl resolution by type.
+
+//! > test_runner_name
+test_function_diagnostics(expect_diagnostics: false)
+
+//! > function
+fn foo() {
+    let _u32: u32 = 2.pow(5);
+    const _U32: u32 = 2.pow(5);
+    let _felt252: felt252 = 2.pow(5);
+    const _FELT252: felt252 = 2.pow(5);
+}
+
+//! > function_name
+foo
+
+//! > module_code
+use core::num::traits::Pow;
+
+//! > expected_diagnostics