Rust: Disambiguate some method calls based on argument types

rust/ql/lib/codeql/rust/internal/TypeInference.qll

@@ -1231,11 +1231,133 @@ private Function getTypeParameterMethod(TypeParameter tp, string name) {
   result = getMethodSuccessor(tp.(ImplTraitTypeTypeParameter).getImplTraitTypeRepr(), name)
 }
 
+bindingset[t1, t2]
+private predicate typeMentionEqual(TypeMention t1, TypeMention t2) {
+  forex(TypePath path, Type type | t1.resolveTypeAt(path) = type | t2.resolveTypeAt(path) = type)
+}
+
+pragma[nomagic]
+private predicate implSiblingCandidate(
+  Impl impl, TraitItemNode trait, Type rootType, TypeMention selfTy
+) {
+  trait = impl.(ImplItemNode).resolveTraitTy() and
+  // If `impl` has an expansion from a macro attribute, then it's been
+  // superseded by the output of the expansion (and usually the expansion
+  // contains the same `impl` block so considering both would give spurious
+  // siblings).
+  not exists(impl.getAttributeMacroExpansion()) and
+  // We use this for resolving methods, so exclude traits that do not have methods.
+  exists(Function f | f = trait.getASuccessor(_) and f.getParamList().hasSelfParam()) and
+  selfTy = impl.getSelfTy() and
+  rootType = selfTy.resolveType()
+}
+
+/**
+ * Holds if `impl1` and `impl2` are a sibling implementations of `trait`. We
+ * consider implementations to be siblings if they implement the same trait for
+ * the same type. In that case `Self` is the same type in both implementations,
+ * and method calls to the implementations cannot be resolved unambiguously
+ * based only on the receiver type.
+ */
+pragma[inline]
+private predicate implSiblings(TraitItemNode trait, Impl impl1, Impl impl2) {
+  exists(Type rootType, TypeMention selfTy1, TypeMention selfTy2 |
+    impl1 != impl2 and
+    implSiblingCandidate(impl1, trait, rootType, selfTy1) and
+    implSiblingCandidate(impl2, trait, rootType, selfTy2) and
+    // In principle the second conjunct below should be superflous, but we still
+    // have ill-formed type mentions for types that we don't understand. For
+    // those checking both directions restricts further. Note also that we check
+    // syntactic equality, whereas equality up to renaming would be more
+    // correct.
+    typeMentionEqual(selfTy1, selfTy2) and
+    typeMentionEqual(selfTy2, selfTy1)
+  )
+}
+
+/**
+ * Holds if `impl` is an implementation of `trait` and if another implementation
+ * exists for the same type.
+ */
+pragma[nomagic]
+private predicate implHasSibling(Impl impl, Trait trait) { implSiblings(trait, impl, _) }
+
+/**
+ * Holds if a type parameter of `trait` occurs in the method with the name
+ * `methodName` at the `pos`th parameter at `path`.
+ */
+bindingset[trait]
+pragma[inline_late]
+private predicate traitTypeParameterOccurrence(
+  TraitItemNode trait, string methodName, int pos, TypePath path
+) {
+  exists(Function f | f = trait.getASuccessor(methodName) |
+    f.getParam(pos).getTypeRepr().(TypeMention).resolveTypeAt(path) =
+      trait.(TraitTypeAbstraction).getATypeParameter()
+  )
+}
+
+bindingset[f, pos, path]
+pragma[inline_late]
+private predicate methodTypeAtPath(Function f, int pos, TypePath path, Type type) {
+  f.getParam(pos).getTypeRepr().(TypeMention).resolveTypeAt(path) = type
+}
+
+/**
+ * Holds if resolving the method `f` in `impl` with the name `methodName`
+ * requires inspecting the types of applied _arguments_ in order to determine
+ * whether it is the correct resolution.
+ */
+pragma[nomagic]
+private predicate methodResolutionDependsOnArgument(
+  Impl impl, string methodName, Function f, int pos, TypePath path, Type type
+) {
+  /*
+   * As seen in the example below, when an implementation has a sibling for a
+   * trait we find occurrences of a type parameter of the trait in a method
+   * signature in the trait. We then find the type given in the implementation
+   * at the same position, which is a position that might disambiguate the
+   * method from its siblings.
+   *
+   * ```rust
+   * trait MyTrait<T> {
+   *     fn method(&self, value: Foo<T>) -> Self;
+   * //                   ^^^^^^^^^^^^^ `pos` = 0
+   * //                              ^ `path` = "T"
+   * }
+   * impl MyAdd<i64> for i64 {
+   *     fn method(&self, value: Foo<i64>) -> Self { ... }
+   * //                              ^^^ `type` = i64
+   * }
+   * ```
+   *
+   * Note that we only check the root type symbol at the position. If the type
+   * at that position is a type constructor (for instance `Vec<..>`) then
+   * inspecting the entire type tree could be necessary to disambiguate the
+   * method. In that case we will still resolve several methods.
+   */
+
+  exists(TraitItemNode trait |
+    implHasSibling(impl, trait) and
+    traitTypeParameterOccurrence(trait, methodName, pos, path) and
+    methodTypeAtPath(getMethodSuccessor(impl, methodName), pos, path, type) and
+    f = getMethodSuccessor(impl, methodName)
+  )
+}
+
 /** Gets a method from an `impl` block that matches the method call `mc`. */
 private Function getMethodFromImpl(MethodCall mc) {
   exists(Impl impl |
     IsInstantiationOf<MethodCall, IsInstantiationOfInput>::isInstantiationOf(mc, impl, _) and
     result = getMethodSuccessor(impl, mc.getMethodName())
+  |
+    not methodResolutionDependsOnArgument(impl, _, _, _, _, _) and
+    result = getMethodSuccessor(impl, mc.getMethodName())
+    or
+    exists(int pos, TypePath path, Type type |
+      methodResolutionDependsOnArgument(impl, mc.getMethodName(), result, pos, path, type) and
+      inferType(mc.getArgument(pos), path) = type
+    )
   )
 }
 

rust/ql/test/library-tests/type-inference/main.rs

@@ -1099,6 +1099,7 @@ mod method_call_type_conversion {
         println!("{:?}", x5.0); // $ fieldof=S
 
         let x6 = &S(S2); // $ SPURIOUS: type=x6:&T.&T.S
+
         // explicit dereference
         println!("{:?}", (*x6).m1()); // $ method=m1 method=deref
 
@@ -1668,17 +1669,18 @@ mod async_ {
     }
 
     fn f2() -> impl Future<Output = S1> {
-        async {
-            S1
-        }
+        async { S1 }
     }
 
     struct S2;
 
     impl Future for S2 {
         type Output = S1;
 
-        fn poll(self: std::pin::Pin<&mut Self>, _cx: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
+        fn poll(
+            self: std::pin::Pin<&mut Self>,
+            _cx: &mut std::task::Context<'_>,
+        ) -> std::task::Poll<Self::Output> {
             std::task::Poll::Ready(S1)
         }
     }
@@ -1692,14 +1694,11 @@ mod async_ {
         f2().await.f(); // $ method=S1f
         f3().await.f(); // $ method=S1f
         S2.await.f(); // $ method=S1f
-        let b = async {
-            S1
-        };
+        let b = async { S1 };
         b.await.f(); // $ method=S1f
     }
 }
 
-
 mod impl_trait {
     struct S1;
     struct S2;
@@ -1816,6 +1815,44 @@ mod macros {
     }
 }
 
+mod method_determined_by_argument_type {
+    trait MyAdd<T> {
+        fn my_add(&self, value: T) -> Self;
+    }
+
+    impl MyAdd<i64> for i64 {
+        // MyAdd<i64>::my_add
+        fn my_add(&self, value: i64) -> Self {
+            value
+        }
+    }
+
+    impl MyAdd<&i64> for i64 {
+        // MyAdd<&i64>::my_add
+        fn my_add(&self, value: &i64) -> Self {
+            *value // $ method=deref
+        }
+    }
+
+    impl MyAdd<bool> for i64 {
+        // MyAdd<bool>::my_add
+        fn my_add(&self, value: bool) -> Self {
+            if value {
+                1
+            } else {
+                0
+            }
+        }
+    }
+
+    pub fn f() {
+        let x: i64 = 73;
+        x.my_add(5i64); // $ method=MyAdd<i64>::my_add
+        x.my_add(&5i64); // $ method=MyAdd<&i64>::my_add
+        x.my_add(true); // $ method=MyAdd<bool>::my_add
+    }
+}
+
 fn main() {
     field_access::f();
     method_impl::f();
@@ -1839,4 +1876,5 @@ fn main() {
     impl_trait::f();
     indexers::f();
     macros::f();
+    method_determined_by_argument_type::f();
 }