Support stubbing trait method implementations (#4587)

Enable stubbing of trait method implementations using fully-qualified
syntax (`<Type as Trait>::method`), including generic traits.
Previously, Kani's stub resolution only searched inherent `impl` blocks
and trait method implementations were silently ignored.

Resolves #1997.

## Problem

Users could not stub trait method implementations. Writing
`#[kani::stub(<MyType as MyTrait>::method, mock_method)]` failed because
the resolution algorithm did not search through trait implementations
for the concrete type. This was a significant limitation since many Rust
APIs are trait-based.

## Solution

### Resolution (`resolve.rs`)

`resolve_in_trait_impl` now correctly resolves trait method
implementations by:
1. Extracting generic arguments from the trait segment (e.g., `<u32>`
from `Convert<u32>`)
2. Building the full `GenericArgs` with the Self type + trait generic
parameters
3. Resolving the `Instance` for the concrete type's implementation

This handles both simple traits (`<Vec<u8> as MyTrait>::method`) and
generic traits (`<MyType as Convert<u32>>::convert`).

### Documentation (`stubbing.md`)

Updated the user-facing docs to document trait method stubbing with
examples covering:
- Basic trait method stubbing syntax
- Generic trait methods
- Supported patterns (dyn dispatch, supertraits, overridden defaults)
- Known limitation: non-overridden default methods

## Testing

**5 passing tests:**
- `stub_trait_method.rs` — basic `<Type as Trait>::method` stubbing
- `stub_generic_trait_method.rs` — generic trait `<Type as
Trait<T>>::method`
- `stub_trait_dyn_dispatch.rs` — dynamic dispatch through `&dyn Trait`
and `Box<dyn Trait>`
- `stub_trait_supertrait.rs` — methods defined in supertraits
- `stub_trait_default_overridden.rs` — default trait methods that are
overridden in the impl

**1 fixme test:**
- `fixme_stub_trait_default_method.rs` — non-overridden default methods
cause a type mismatch during stub validation (the default body uses
`Self` as a placeholder type)

---

By submitting this pull request, I confirm that my contribution is made
under the terms of the Apache 2.0 and MIT licenses.

---------

Signed-off-by: Felipe R. Monteiro <felisous@amazon.com>

Author: feliperodri

docs/src/reference/experimental/stubbing.md

@@ -170,13 +170,47 @@ In the following, we describe the known limitations of the stubbing feature.
 
 ### Trait method stubbing
 
-Stubbing trait method implementations (e.g., `<MyType as MyTrait>::method`) is **not supported**.
-Only free functions and inherent methods can be stubbed. Rust's attribute path syntax does not
-support the fully-qualified `<Type as Trait>::method` form, and the resolution algorithm does not
-search through trait implementations.
+Kani supports stubbing trait method implementations using fully-qualified syntax:
 
-If you need to stub a trait method, consider stubbing the calling function instead, or using
-conditional compilation (`#[cfg(kani)]`) to provide an alternative implementation.
+```rust
+trait MyTrait {
+    fn method(&self) -> u32;
+}
+
+struct MyType;
+
+impl MyTrait for MyType {
+    fn method(&self) -> u32 { 0 }
+}
+
+fn stub_method(_x: &MyType) -> u32 { 42 }
+
+#[kani::proof]
+#[kani::stub(<MyType as MyTrait>::method, stub_method)]
+fn check_trait_stub() {
+    let x = MyType;
+    assert_eq!(x.method(), 42);
+}
+```
+
+This also works with generic traits:
+
+```rust
+#[kani::stub(<MyType as Convert<u32>>::convert, stub_convert)]
+```
+
+When two traits define methods with the same name, the fully-qualified syntax
+disambiguates which implementation to stub.
+
+The following trait patterns are supported:
+
+- **Supertrait methods:** Stubbing a method defined in a supertrait (e.g., `<MyType as Base>::method`) works independently of subtrait methods.
+- **Overridden default methods:** If a type overrides a trait's default method, the override can be stubbed normally.
+- **Dynamic dispatch:** Stubs apply even when the method is called through a trait object (`&dyn Trait` or `Box<dyn Trait>`).
+
+**Known limitation:** Stubbing a trait's default method that is *not* overridden by the implementing type is not currently supported ([#4588](https://github.com/model-checking/kani/issues/4588)). The default method body uses `Self` as a placeholder type, which causes a type mismatch during stub validation. This applies only to default methods that are inherited as-is (i.e., the `impl` block does not provide its own definition).
+
+**Known limitation:** Traits with const generic parameters (e.g., `<Type as Buf<16>>::write`) or associated type constraints (e.g., `<Type as Iterator<Item = u32>>::next`) are not currently supported and will produce a resolution error.
 
 ### Usage restrictions
 
@@ -191,7 +225,6 @@ Support for stubbing is currently **limited to functions and methods**. All othe
 The following are examples of items that could be good candidates for stubbing, but aren't supported:
 - Types
 - Macros
-- Traits
 - Intrinsics
 
 We acknowledge that support for method stubbing isn't as ergonomic as it could be.

kani-compiler/src/kani_middle/resolve.rs

@@ -73,7 +73,24 @@ pub fn resolve_fn_path<'tcx>(
             let ty = type_resolution::resolve_ty(tcx, current_module, syn_ty)?;
             let trait_fn_id = resolve_path(tcx, current_module, &path.path)?;
             validate_kind!(tcx, trait_fn_id, "function / method", DefKind::Fn | DefKind::AssocFn)?;
-            resolve_in_trait_impl(tcx, ty, trait_fn_id)
+            // Extract generic args from the trait segment (e.g., <u32> from Convert<u32>).
+            // For `<Type as crate::mod::Trait<u32>>::method`, segments after `as` are
+            // [crate, mod, Trait<u32>, method]. The trait segment with generic args is
+            // always second-to-last (last is the method name).
+            let trait_generic_args = path
+                .path
+                .segments
+                .iter()
+                .rev()
+                .nth(1) // The trait segment is second-to-last (last is the method name)
+                .and_then(|seg| {
+                    if let syn::PathArguments::AngleBracketed(args) = &seg.arguments {
+                        Some(args)
+                    } else {
+                        None
+                    }
+                });
+            resolve_in_trait_impl(tcx, Some(current_module), ty, trait_fn_id, trait_generic_args)
         }
         // Qualified path for a primitive type, such as `<[u8]::sort>`.
         Some(QSelf { ty: syn_ty, .. }) if type_resolution::is_type_primitive(syn_ty) => {
@@ -399,8 +416,17 @@ fn resolve_in_any_trait<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty, name: &str) -> Option<F
         .into_iter()
         .filter_map(|trait_def| {
             resolve_in_trait_def_stable(tcx, trait_def, name).ok().and_then(|item| {
-                resolve_in_trait_impl(tcx, ty, rustc_internal::internal(tcx, item.def_id.def_id()))
-                    .ok()
+                // No current_module or trait_args needed: this path resolves
+                // unqualified method names (e.g., `x.foo()`), not user-specified
+                // `<Type as Trait<T>>::method` syntax.
+                resolve_in_trait_impl(
+                    tcx,
+                    None,
+                    ty,
+                    rustc_internal::internal(tcx, item.def_id.def_id()),
+                    None,
+                )
+                .ok()
             })
         })
         .collect();
@@ -411,16 +437,49 @@ fn resolve_in_any_trait<'tcx>(tcx: TyCtxt<'tcx>, ty: Ty, name: &str) -> Option<F
 /// Resolves a trait method implementation by checking if there exists an Instance of the trait method for `ty`.
 /// This is distinct from `resolve_in_trait_def`: that function checks if the associated function is defined for the trait,
 /// while this function checks if it's implemented for `ty`.
-fn resolve_in_trait_impl(
-    tcx: TyCtxt<'_>,
+fn resolve_in_trait_impl<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    current_module: Option<LocalDefId>,
     ty: Ty,
     trait_fn_id: DefId,
-) -> Result<FnResolution, ResolveError<'_>> {
-    debug!(?ty, "resolve_in_trait_impl");
+    trait_args: Option<&syn::AngleBracketedGenericArguments>,
+) -> Result<FnResolution, ResolveError<'tcx>> {
+    debug!(?ty, ?trait_args, "resolve_in_trait_impl");
     // Given the `FnDef` of the *definition* of the trait method, see if there exists an Instance
     // that implements that method for `ty`.
     let trait_fn_fn_def = stable_fn_def(tcx, trait_fn_id).unwrap();
-    let desired_generic_args = GenericArgs(vec![GenericArgKind::Type(ty)]);
+
+    // Build generic args: Self type + any trait generic parameters (e.g., T in Convert<T>).
+    let mut args = vec![GenericArgKind::Type(ty)];
+    if let Some(syn_args) = trait_args {
+        let Some(module) = current_module else {
+            return Err(ResolveError::UnsupportedPath {
+                kind: "stub resolution requires generic trait arguments to have a known module context",
+            });
+        };
+        for arg in &syn_args.args {
+            match arg {
+                syn::GenericArgument::Type(syn_ty) => {
+                    let resolved_ty = type_resolution::resolve_ty(tcx, module, syn_ty)?;
+                    args.push(GenericArgKind::Type(resolved_ty));
+                }
+                syn::GenericArgument::Const(_) => {
+                    return Err(ResolveError::UnsupportedPath {
+                        kind: "const generic arguments in trait stubs",
+                    });
+                }
+                syn::GenericArgument::Lifetime(_) => {
+                    // Lifetimes are erased — safe to skip.
+                }
+                _ => {
+                    return Err(ResolveError::UnsupportedPath {
+                        kind: "associated type/const constraints in trait stubs",
+                    });
+                }
+            }
+        }
+    }
+    let desired_generic_args = GenericArgs(args);
     let exists = Instance::resolve(trait_fn_fn_def, &desired_generic_args);
 
     // If such an Instance exists, return *its* FnDef (i.e., the FnDef inside the impl block for this `ty`)

tests/kani/Stubbing/fixme_stub_trait_default_method.rs

@@ -0,0 +1,32 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//
+// kani-flags: -Z stubbing
+//
+//! Test stubbing a trait's default method implementation.
+//!
+//! FIXME: Default trait methods use `Self` in their body types, which
+//! doesn't match the concrete type in the stub. The validation compares
+//! the trait definition's body (with `&Self`) against the stub (with
+//! `&MyType`), causing a false type mismatch.
+//! Tracked in: https://github.com/model-checking/kani/issues/4588
+
+trait HasDefault {
+    fn default_method(&self) -> u32 {
+        0
+    }
+}
+
+struct MyType;
+impl HasDefault for MyType {}
+
+fn stub_default(_x: &MyType) -> u32 {
+    42
+}
+
+#[kani::proof]
+#[kani::stub(<MyType as HasDefault>::default_method, stub_default)]
+fn check_stub_default_trait_method() {
+    let x = MyType;
+    assert_eq!(x.default_method(), 42);
+}

…unction-contracts/generic_slice_index.rs …ani/Stubbing/stub_generic_slice_index.rstests/ui/function-contracts/generic_slice_index.rs renamed to tests/kani/Stubbing/stub_generic_slice_index.rs tests/ui/function-contracts/generic_slice_index.rs renamed to tests/kani/Stubbing/stub_generic_slice_index.rs

@@ -1,12 +1,9 @@
 // Copyright Kani Contributors
 // SPDX-License-Identifier: Apache-2.0 OR MIT
 // kani-flags: -Zfunction-contracts
-//! Test function contracts on generic trait implementations based on SliceIndex,
-//! c.f. https://github.com/model-checking/kani/issues/4084
-//! This `proof_for_contract` should work,
-//! but we do not yet support stubbing/contracts on trait fns with generic arguments
-//! c.f. https://github.com/model-checking/kani/issues/1997#issuecomment-3134614734.
-//! So for now, test that we emit a nice error message.
+//! Test function contracts on generic trait implementations based on SliceIndex.
+//! Regression test for https://github.com/model-checking/kani/issues/1997
+//! and https://github.com/model-checking/kani/issues/4084.
 
 const INVALID_INDEX: usize = 10;
 

tests/kani/Stubbing/stub_generic_trait_method.rs

@@ -0,0 +1,38 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//
+// kani-flags: -Z stubbing
+//
+//! Test stubbing generic trait method implementations.
+//! Regression test for https://github.com/model-checking/kani/issues/1997
+
+trait Convert<T> {
+    fn convert(&self) -> T;
+}
+
+struct MyType;
+
+impl Convert<u32> for MyType {
+    fn convert(&self) -> u32 {
+        100
+    }
+}
+
+impl Convert<bool> for MyType {
+    fn convert(&self) -> bool {
+        false
+    }
+}
+
+fn stub_convert_u32(_x: &MyType) -> u32 {
+    42
+}
+
+#[kani::proof]
+#[kani::stub(<MyType as Convert<u32>>::convert, stub_convert_u32)]
+fn check_generic_trait_stub() {
+    let m = MyType;
+    assert_eq!(<MyType as Convert<u32>>::convert(&m), 42);
+    // The bool impl should NOT be affected
+    assert_eq!(<MyType as Convert<bool>>::convert(&m), false);
+}

tests/kani/Stubbing/stub_trait_default_overridden.rs

@@ -0,0 +1,32 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//
+// kani-flags: -Z stubbing
+//
+//! Test stubbing a trait default method that IS overridden in the impl.
+//! The stub should replace the overridden implementation, not the default.
+
+trait HasDefault {
+    fn method(&self) -> u32 {
+        0
+    }
+}
+
+struct MyType;
+
+impl HasDefault for MyType {
+    fn method(&self) -> u32 {
+        100
+    }
+}
+
+fn stub_method(_x: &MyType) -> u32 {
+    42
+}
+
+#[kani::proof]
+#[kani::stub(<MyType as HasDefault>::method, stub_method)]
+fn check_stub_overridden_default() {
+    let x = MyType;
+    assert_eq!(x.method(), 42);
+}

tests/kani/Stubbing/stub_trait_dyn_dispatch.rs

@@ -0,0 +1,38 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//
+// kani-flags: -Z stubbing
+//
+//! Test stubbing a trait method called through dynamic dispatch (trait objects).
+
+trait Animal {
+    fn sound(&self) -> u32;
+}
+
+struct Dog;
+
+impl Animal for Dog {
+    fn sound(&self) -> u32 {
+        100
+    }
+}
+
+fn stub_sound(_x: &Dog) -> u32 {
+    42
+}
+
+fn call_via_dyn(a: &dyn Animal) -> u32 {
+    a.sound()
+}
+
+#[kani::proof]
+#[kani::stub(<Dog as Animal>::sound, stub_sound)]
+fn check_stub_dyn_dispatch() {
+    let dog = Dog;
+    // Direct call — should be stubbed
+    assert_eq!(dog.sound(), 42);
+    // Call through trait object — stub applies because Kani replaces the
+    // function body globally, which is used regardless of dispatch mechanism.
+    let result = call_via_dyn(&dog);
+    assert_eq!(result, 42);
+}

tests/kani/Stubbing/stub_trait_method.rs

@@ -0,0 +1,70 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+//
+// kani-flags: -Z stubbing
+//
+//! Test stubbing trait method implementations using fully-qualified syntax.
+//! Regression test for https://github.com/model-checking/kani/issues/1997
+
+trait A {
+    fn foo(&self) -> u32;
+    fn bar(&self) -> u32;
+}
+
+trait B {
+    fn bar(&self) -> u32;
+}
+
+struct X {}
+
+impl X {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl A for X {
+    fn foo(&self) -> u32 {
+        100
+    }
+    fn bar(&self) -> u32 {
+        200
+    }
+}
+
+impl B for X {
+    fn bar(&self) -> u32 {
+        300
+    }
+}
+
+fn stub_1(_x: &X) -> u32 {
+    1
+}
+fn stub_2(_x: &X) -> u32 {
+    2
+}
+fn stub_3(_x: &X) -> u32 {
+    3
+}
+
+#[kani::proof]
+#[kani::stub(<X as A>::foo, stub_1)]
+fn check_stub_trait_a_foo() {
+    let x = X::new();
+    assert_eq!(x.foo(), 1);
+}
+
+#[kani::proof]
+#[kani::stub(<X as A>::bar, stub_2)]
+fn check_stub_trait_a_bar() {
+    let x = X::new();
+    assert_eq!(A::bar(&x), 2);
+}
+
+#[kani::proof]
+#[kani::stub(<X as B>::bar, stub_3)]
+fn check_stub_trait_b_bar() {
+    let x = X::new();
+    assert_eq!(<X as B>::bar(&x), 3);
+}