Fix bug: Kani unwinds loops with contract in generic function (with -Z loop-contracts) (model-checking#4232)

This PR fixes the bug that even with -Z loop-contracts, Kani still
unwind the loop if the loop is inside a generic function.

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

Author: thanhnguyen-aws

kani-compiler/src/kani_middle/transform/loop_contracts.rs

@@ -661,7 +661,19 @@ impl LoopContractPass {
                 let TyKind::RigidTy(RigidTy::Closure(_, genarg)) = arg_ty.kind() else {
                     return false;
                 };
-                let GenericArgKind::Type(arg_ty) = genarg.0[2] else { return false };
+                // We look for the args' types of the kani_register_loop_contract function
+                // They are always stored in a tuple, which is next to the FnPtr generic args of kani_registered_loop_contract fn
+                // All the generic args before the FnPtr are from the outer function
+                let mut fnptrpos = 0;
+                for (i, arg) in genarg.0.iter().enumerate() {
+                    if let GenericArgKind::Type(arg_ty) = arg
+                        && let TyKind::RigidTy(RigidTy::FnPtr(_)) = arg_ty.kind()
+                    {
+                        fnptrpos = i;
+                        break;
+                    }
+                }
+                let GenericArgKind::Type(arg_ty) = genarg.0[fnptrpos + 1] else { return false };
                 let TyKind::RigidTy(RigidTy::Tuple(args)) = arg_ty.kind() else { return false };
                 // Check if the invariant involves any local variable
                 if !args.is_empty() {

tests/expected/loop-contract/loop_in_generic_function.expected

@@ -0,0 +1,9 @@
+sum::<u8>.loop_invariant_base.1\
+	 - Status: SUCCESS\
+	 - Description: "Check invariant before entry for loop sum::<u8>.0"
+
+main.assertion.1\
+	 - Status: SUCCESS\
+	 - Description: "assertion failed: j <= 2560"
+
+VERIFICATION:- SUCCESSFUL

tests/expected/loop-contract/loop_in_generic_function.rs

@@ -0,0 +1,30 @@
+// Copyright Kani Contributors
+// SPDX-License-Identifier: Apache-2.0 OR MIT
+
+// kani-flags: -Z loop-contracts
+
+//! Check loop in generic function
+
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+
+fn sum<T: Clone>(a: &[T]) -> u32
+where
+    u32: std::convert::From<T>,
+{
+    let mut j: u32 = 0;
+    let mut i: usize = 0;
+    #[kani::loop_invariant(i<=10 && j <= (u8::MAX as u32) * (i as u32))]
+    while i < 10 {
+        j = j + std::convert::Into::<u32>::into(a[i].clone());
+        i = i + 1;
+    }
+    j
+}
+
+#[kani::proof]
+fn main() {
+    let a: [u8; 10] = kani::any();
+    let j = sum(a.as_slice());
+    assert!(j <= 2560);
+}