fix: "Missing trait impl" error in trait dispatch (#3440)

Co-authored-by: kevaundray <kevtheappdev@gmail.com>
Co-authored-by: Maxim Vezenov <mvezenov@gmail.com>

compiler/noirc_frontend/src/hir/resolution/resolver.rs

@@ -29,7 +29,7 @@ use crate::hir::def_map::{LocalModuleId, ModuleDefId, TryFromModuleDefId, MAIN_F
 use crate::hir_def::stmt::{HirAssignStatement, HirForStatement, HirLValue, HirPattern};
 use crate::node_interner::{
     DefinitionId, DefinitionKind, ExprId, FuncId, NodeInterner, StmtId, StructId, TraitId,
-    TraitImplId,
+    TraitImplId, TraitImplKind,
 };
 use crate::{
     hir::{def_map::CrateDefMap, resolution::path_resolver::PathResolver},
@@ -1207,8 +1207,12 @@ impl<'a> Resolver<'a> {
                 Literal::Unit => HirLiteral::Unit,
             }),
             ExpressionKind::Variable(path) => {
-                if let Some(expr) = self.resolve_trait_generic_path(&path) {
-                    expr
+                if let Some((hir_expr, object_type)) = self.resolve_trait_generic_path(&path) {
+                    let expr_id = self.interner.push_expr(hir_expr);
+                    self.interner.push_expr_location(expr_id, expr.span, self.file);
+                    self.interner
+                        .select_impl_for_ident(expr_id, TraitImplKind::Assumed { object_type });
+                    return expr_id;
                 } else {
                     // If the Path is being used as an Expression, then it is referring to a global from a separate module
                     // Otherwise, then it is referring to an Identifier
@@ -1370,6 +1374,8 @@ impl<'a> Resolver<'a> {
             ExpressionKind::Parenthesized(sub_expr) => return self.resolve_expression(*sub_expr),
         };
 
+        // If these lines are ever changed, make sure to change the early return
+        // in the ExpressionKind::Variable case as well
         let expr_id = self.interner.push_expr(hir_expr);
         self.interner.push_expr_location(expr_id, expr.span, self.file);
         expr_id
@@ -1576,7 +1582,10 @@ impl<'a> Resolver<'a> {
     }
 
     // this resolves Self::some_static_method, inside an impl block (where we don't have a concrete self_type)
-    fn resolve_trait_static_method_by_self(&mut self, path: &Path) -> Option<HirExpression> {
+    fn resolve_trait_static_method_by_self(
+        &mut self,
+        path: &Path,
+    ) -> Option<(HirExpression, Type)> {
         if let Some(trait_id) = self.trait_id {
             if path.kind == PathKind::Plain && path.segments.len() == 2 {
                 let name = &path.segments[0].0.contents;
@@ -1590,7 +1599,7 @@ impl<'a> Resolver<'a> {
                             the_trait.self_type_typevar,
                             crate::TypeVariableKind::Normal,
                         );
-                        return Some(HirExpression::TraitMethodReference(self_type, method));
+                        return Some((HirExpression::TraitMethodReference(method), self_type));
                     }
                 }
             }
@@ -1599,7 +1608,10 @@ impl<'a> Resolver<'a> {
     }
 
     // this resolves a static trait method T::trait_method by iterating over the where clause
-    fn resolve_trait_method_by_named_generic(&mut self, path: &Path) -> Option<HirExpression> {
+    fn resolve_trait_method_by_named_generic(
+        &mut self,
+        path: &Path,
+    ) -> Option<(HirExpression, Type)> {
         if path.segments.len() != 2 {
             return None;
         }
@@ -1621,15 +1633,15 @@ impl<'a> Resolver<'a> {
                         the_trait.find_method(path.segments.last().unwrap().clone())
                     {
                         let self_type = self.resolve_type(typ.clone());
-                        return Some(HirExpression::TraitMethodReference(self_type, method));
+                        return Some((HirExpression::TraitMethodReference(method), self_type));
                     }
                 }
             }
         }
         None
     }
 
-    fn resolve_trait_generic_path(&mut self, path: &Path) -> Option<HirExpression> {
+    fn resolve_trait_generic_path(&mut self, path: &Path) -> Option<(HirExpression, Type)> {
         self.resolve_trait_static_method_by_self(path)
             .or_else(|| self.resolve_trait_method_by_named_generic(path))
     }

compiler/noirc_frontend/src/hir/type_check/errors.rs

@@ -113,6 +113,8 @@ pub enum TypeCheckError {
     },
     #[error("No matching impl found")]
     NoMatchingImplFound { constraints: Vec<(Type, String)>, span: Span },
+    #[error("Constraint for `{typ}: {trait_name}` is not needed, another matching impl is already in scope")]
+    UnneededTraitConstraint { trait_name: String, typ: Type, span: Span },
 }
 
 impl TypeCheckError {
@@ -269,6 +271,10 @@ impl From<TypeCheckError> for Diagnostic {
 
                 diagnostic
             }
+            TypeCheckError::UnneededTraitConstraint { trait_name, typ, span } => {
+                let msg = format!("Constraint for `{typ}: {trait_name}` is not needed, another matching impl is already in scope");
+                Diagnostic::simple_warning(msg, "Unnecessary trait constraint in where clause".into(), span)
+            }
         }
     }
 }

compiler/noirc_frontend/src/hir/type_check/expr.rs

@@ -10,7 +10,7 @@ use crate::{
         },
         types::Type,
     },
-    node_interner::{DefinitionKind, ExprId, FuncId, TraitMethodId},
+    node_interner::{DefinitionKind, ExprId, FuncId, TraitId, TraitMethodId},
     BinaryOpKind, Signedness, TypeBinding, TypeVariableKind, UnaryOp,
 };
 
@@ -132,7 +132,9 @@ impl<'interner> TypeChecker<'interner> {
             HirExpression::Index(index_expr) => self.check_index_expression(expr_id, index_expr),
             HirExpression::Call(call_expr) => {
                 self.check_if_deprecated(&call_expr.func);
+
                 let function = self.check_expression(&call_expr.func);
+
                 let args = vecmap(&call_expr.arguments, |arg| {
                     let typ = self.check_expression(arg);
                     (typ, *arg, self.interner.expr_span(arg))
@@ -160,21 +162,29 @@ impl<'interner> TypeChecker<'interner> {
                         // so that the backend doesn't need to worry about methods
                         let location = method_call.location;
 
-                        let mut func_id = None;
-                        if let HirMethodReference::FuncId(id) = method_ref {
-                            func_id = Some(id);
-
-                            // Automatically add `&mut` if the method expects a mutable reference and
-                            // the object is not already one.
-                            if id != FuncId::dummy_id() {
-                                let func_meta = self.interner.function_meta(&id);
-                                self.try_add_mutable_reference_to_object(
-                                    &mut method_call,
-                                    &func_meta.typ,
-                                    &mut args,
-                                );
+                        let trait_id = match &method_ref {
+                            HirMethodReference::FuncId(func_id) => {
+                                // Automatically add `&mut` if the method expects a mutable reference and
+                                // the object is not already one.
+                                if *func_id != FuncId::dummy_id() {
+                                    let func_meta = self.interner.function_meta(func_id);
+                                    self.try_add_mutable_reference_to_object(
+                                        &mut method_call,
+                                        &func_meta.typ,
+                                        &mut args,
+                                    );
+                                }
+
+                                let meta = self.interner.function_meta(func_id);
+                                meta.trait_impl.map(|impl_id| {
+                                    self.interner
+                                        .get_trait_implementation(impl_id)
+                                        .borrow()
+                                        .trait_id
+                                })
                             }
-                        }
+                            HirMethodReference::TraitMethodId(method) => Some(method.trait_id),
+                        };
 
                         let (function_id, function_call) = method_call.into_function_call(
                             method_ref.clone(),
@@ -185,29 +195,8 @@ impl<'interner> TypeChecker<'interner> {
                         let span = self.interner.expr_span(expr_id);
                         let ret = self.check_method_call(&function_id, method_ref, args, span);
 
-                        if let Some(func_id) = func_id {
-                            let meta = self.interner.function_meta(&func_id);
-
-                            if let Some(impl_id) = meta.trait_impl {
-                                let trait_impl = self.interner.get_trait_implementation(impl_id);
-
-                                let result = self.interner.lookup_trait_implementation(
-                                    &object_type,
-                                    trait_impl.borrow().trait_id,
-                                );
-
-                                if let Err(erroring_constraints) = result {
-                                    let constraints = vecmap(erroring_constraints, |constraint| {
-                                        let r#trait = self.interner.get_trait(constraint.trait_id);
-                                        (constraint.typ, r#trait.name.to_string())
-                                    });
-
-                                    self.errors.push(TypeCheckError::NoMatchingImplFound {
-                                        constraints,
-                                        span,
-                                    });
-                                }
-                            }
+                        if let Some(trait_id) = trait_id {
+                            self.verify_trait_constraint(&object_type, trait_id, function_id, span);
                         }
 
                         self.interner.replace_expr(expr_id, function_call);
@@ -285,7 +274,7 @@ impl<'interner> TypeChecker<'interner> {
 
                 Type::Function(params, Box::new(lambda.return_type), Box::new(env_type))
             }
-            HirExpression::TraitMethodReference(_, method) => {
+            HirExpression::TraitMethodReference(method) => {
                 let the_trait = self.interner.get_trait(method.trait_id);
                 let method = &the_trait.methods[method.method_index];
 
@@ -305,6 +294,26 @@ impl<'interner> TypeChecker<'interner> {
         typ
     }
 
+    fn verify_trait_constraint(
+        &mut self,
+        object_type: &Type,
+        trait_id: TraitId,
+        function_ident_id: ExprId,
+        span: Span,
+    ) {
+        match self.interner.lookup_trait_implementation(object_type, trait_id) {
+            Ok(impl_kind) => self.interner.select_impl_for_ident(function_ident_id, impl_kind),
+            Err(erroring_constraints) => {
+                let constraints = vecmap(erroring_constraints, |constraint| {
+                    let r#trait = self.interner.get_trait(constraint.trait_id);
+                    (constraint.typ, r#trait.name.to_string())
+                });
+
+                self.errors.push(TypeCheckError::NoMatchingImplFound { constraints, span });
+            }
+        }
+    }
+
     /// Check if the given method type requires a mutable reference to the object type, and check
     /// if the given object type is already a mutable reference. If not, add one.
     /// This is used to automatically transform a method call: `foo.bar()` into a function
@@ -512,13 +521,11 @@ impl<'interner> TypeChecker<'interner> {
 
                 let func_meta = self.interner.function_meta(&func_id);
                 let param_len = func_meta.parameters.len();
-
                 (func_meta.typ, param_len)
             }
-            HirMethodReference::TraitMethodId(_, method) => {
+            HirMethodReference::TraitMethodId(method) => {
                 let the_trait = self.interner.get_trait(method.trait_id);
                 let method = &the_trait.methods[method.method_index];
-
                 (method.get_type(), method.arguments.len())
             }
         };
@@ -537,7 +544,6 @@ impl<'interner> TypeChecker<'interner> {
 
         self.interner.store_instantiation_bindings(*function_ident_id, instantiation_bindings);
         self.interner.push_expr_type(function_ident_id, function_type.clone());
-
         self.bind_function_type(function_type, arguments, span)
     }
 
@@ -889,10 +895,7 @@ impl<'interner> TypeChecker<'interner> {
                             if method.name.0.contents == method_name {
                                 let trait_method =
                                     TraitMethodId { trait_id: constraint.trait_id, method_index };
-                                return Some(HirMethodReference::TraitMethodId(
-                                    object_type.clone(),
-                                    trait_method,
-                                ));
+                                return Some(HirMethodReference::TraitMethodId(trait_method));
                             }
                         }
                     }

compiler/noirc_frontend/src/hir/type_check/mod.rs

@@ -43,23 +43,46 @@ pub fn type_check_func(interner: &mut NodeInterner, func_id: FuncId) -> Vec<Type
     let mut type_checker = TypeChecker::new(interner);
     type_checker.current_function = Some(func_id);
 
+    let meta = type_checker.interner.function_meta(&func_id);
+    let mut errors = Vec::new();
+
+    // Temporarily add any impls in this function's `where` clause to scope
+    for constraint in &meta.trait_constraints {
+        let object = constraint.typ.clone();
+        let trait_id = constraint.trait_id;
+
+        if !type_checker.interner.add_assumed_trait_implementation(object, trait_id) {
+            let trait_name = type_checker.interner.get_trait(trait_id).name.to_string();
+            let typ = constraint.typ.clone();
+            let span = meta.name.location.span;
+            errors.push(TypeCheckError::UnneededTraitConstraint { trait_name, typ, span });
+        }
+    }
+
     // Bind each parameter to its annotated type.
     // This is locally obvious, but it must be bound here so that the
     // Definition object of the parameter in the NodeInterner is given the correct type.
     for param in meta.parameters.into_iter() {
         type_checker.bind_pattern(&param.0, param.1);
     }
 
-    let (function_last_type, delayed_type_check_functions, mut errors) =
+    let (function_last_type, delayed_type_check_functions, mut body_errors) =
         type_checker.check_function_body(function_body_id);
 
+    errors.append(&mut body_errors);
+
     // Go through any delayed type checking errors to see if they are resolved, or error otherwise.
     for type_check_fn in delayed_type_check_functions {
         if let Err(error) = type_check_fn() {
             errors.push(error);
         }
     }
 
+    // Now remove all the `where` clause constraints we added
+    for constraint in &meta.trait_constraints {
+        interner.remove_assumed_trait_implementations_for_trait(constraint.trait_id);
+    }
+
     // Check declared return type and actual return type
     if !can_ignore_ret {
         let (expr_span, empty_function) = function_info(interner, function_body_id);

compiler/noirc_frontend/src/hir_def/expr.rs

@@ -29,7 +29,7 @@ pub enum HirExpression {
     If(HirIfExpression),
     Tuple(Vec<ExprId>),
     Lambda(HirLambda),
-    TraitMethodReference(Type, TraitMethodId),
+    TraitMethodReference(TraitMethodId),
     Error,
 }
 
@@ -156,7 +156,7 @@ pub enum HirMethodReference {
     /// Or a method can come from a Trait impl block, in which case
     /// the actual function called will depend on the instantiated type,
     /// which can be only known during monomorphization.
-    TraitMethodId(Type, TraitMethodId),
+    TraitMethodId(TraitMethodId),
 }
 
 impl HirMethodCallExpression {
@@ -174,8 +174,8 @@ impl HirMethodCallExpression {
                 let id = interner.function_definition_id(func_id);
                 HirExpression::Ident(HirIdent { location, id })
             }
-            HirMethodReference::TraitMethodId(typ, method_id) => {
-                HirExpression::TraitMethodReference(typ, method_id)
+            HirMethodReference::TraitMethodId(method_id) => {
+                HirExpression::TraitMethodReference(method_id)
             }
         };
         let func = interner.push_expr(expr);

compiler/noirc_frontend/src/hir_def/traits.rs

@@ -151,5 +151,6 @@ impl TraitFunction {
             Box::new(self.return_type.clone()),
             Box::new(Type::Unit),
         )
+        .generalize()
     }
 }