mir-importer: require closure receiver for rust-call lowering

crates/mir-importer/src/translator/terminator/mod.rs

@@ -903,15 +903,13 @@ fn translate_call(
         }
     }
 
-    // Handle closure trait method calls (FnOnce::call_once, FnMut::call_mut, Fn::call)
-    // These calls pass arguments as a tuple, but the closure body expects unpacked args.
-    // We need to unpack the tuple before calling the closure.
-    //
-    // MIR shows: <{closure} as FnMut<(u32,)>>::call_mut(self_ref, tuple_args)
-    // But the closure body expects: fn(self_ref, unpacked_arg1, unpacked_arg2, ...)
+    // Handle genuine closure trait method calls. The receiver test matters:
+    // wrapper ADTs can carry a closure in their generic substitutions while
+    // still expecting the rust-call tuple as one argument.
     if let Some(ref name) = pattern_name
         && (name.contains("call_once") || name.contains("call_mut") || name.ends_with("::call"))
-        && substs_contains("Closure")
+        && !args.is_empty()
+        && receiver_is_closure(&args[0], body)
     {
         return translate_closure_call(
             ctx,
@@ -1451,6 +1449,46 @@ fn translate_closure_call(
     }
 }
 
+/// True only when the rust-call receiver is itself a closure.
+///
+/// We type the operand's base local (or constant) and ignore any
+/// `place.projection`, and we peel at most one reference. Both are safe for
+/// the rust-call receiver specifically: rustc lowers a closure-trait call
+/// (`Fn::call` / `FnMut::call_mut` / `FnOnce::call_once`) so that the receiver
+/// argument is the closure passed by value, or a single `&`/`&mut` borrow of
+/// it, materialized into its own temporary local, never an in-place projection
+/// of a larger aggregate and never behind multiple references. So a closure
+/// reached through a field (`(self.f)(x)`) still arrives here as a base local
+/// of type `&{closure}`, and one `Ref` peel plus a base-local type check covers
+/// every genuine closure-call shape. A wrapper ADT that merely carries a
+/// closure in its generic substitutions (the case this guard exists to reject)
+/// has a non-closure receiver type and is correctly left on the ordinary call
+/// path with its rust-call tuple intact.
+fn receiver_is_closure(receiver: &mir::Operand, body: &mir::Body) -> bool {
+    let ty = match receiver {
+        mir::Operand::Copy(place) | mir::Operand::Move(place) => {
+            let local: usize = place.local;
+            let local_decls: Vec<_> = body.local_decls().collect();
+            match local_decls.get(local).map(|(_, decl)| decl.ty) {
+                Some(ty) => ty,
+                None => return false,
+            }
+        }
+        mir::Operand::Constant(const_op) => const_op.const_.ty(),
+        _ => return false,
+    };
+
+    let inner = match ty.kind() {
+        rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Ref(_, inner, _)) => inner,
+        _ => ty,
+    };
+
+    matches!(
+        inner.kind(),
+        rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Closure(_, _))
+    )
+}
+
 /// Extracts the closure body's mangled name from a closure operand.
 ///
 /// In unified compilation with `std`, when we see a call like:

crates/rustc-codegen-cuda/examples/device_closures/src/main.rs

@@ -136,6 +136,21 @@ mod kernels {
         f(x)
     }
 
+    struct ClosureWrapper<F> {
+        f: F,
+        offset: u32,
+    }
+
+    impl<F> ClosureWrapper<F>
+    where
+        F: Fn(u32) -> u32,
+    {
+        #[inline(never)]
+        fn call(&self, args: (u32,)) -> u32 {
+            (self.f)(args.0) + self.offset
+        }
+    }
+
     // Note: apply_twice with Fn trait is commented out due to control flow issues.
     // The `Fn` trait requires the closure to be called multiple times, which
     // can create complex control flow that the backend doesn't fully support yet.
@@ -201,6 +216,59 @@ mod kernels {
             *out_elem = apply_closure(scale, val);
         }
     }
+
+    /// Wrapper method named `call` whose generic substitutions contain a
+    /// closure. This is not a closure trait shim: the receiver is the wrapper,
+    /// so the tuple argument must remain one ordinary method argument.
+    #[kernel]
+    pub fn test_wrapper_method_named_call(factor: u32, mut out: DisjointSlice<u32>) {
+        let idx = thread::index_1d();
+        let idx_raw = idx.get();
+        if let Some(out_elem) = out.get_mut(idx) {
+            let scale = |x: u32| x * factor;
+            let wrapper = ClosureWrapper {
+                f: scale,
+                offset: 1,
+            };
+            *out_elem = wrapper.call((idx_raw as u32,));
+        }
+    }
+
+    /// Genuine closure call whose receiver is reached through a struct field
+    /// (`(holder.f)(x)`). The tightened `receiver_is_closure` guard must still
+    /// unpack the rust-call tuple here: rustc materializes the `&{closure}`
+    /// receiver into its own temporary local, so the base-local type check
+    /// sees a closure even though the closure lives in a field.
+    #[kernel]
+    pub fn test_closure_field_projection(factor: u32, mut out: DisjointSlice<u32>) {
+        struct Holder<F> {
+            f: F,
+        }
+        let idx = thread::index_1d();
+        let idx_raw = idx.get();
+        if let Some(out_elem) = out.get_mut(idx) {
+            let holder = Holder {
+                f: |x: u32| x * factor,
+            };
+            *out_elem = (holder.f)(idx_raw as u32);
+        }
+    }
+
+    /// Genuine closure call through two reference levels (`(**rr)(x)`). The
+    /// guard peels a single reference; this still classifies as a closure
+    /// because rustc resolves the receiver to one `&{closure}` borrow at the
+    /// call site regardless of the source-level double indirection.
+    #[kernel]
+    pub fn test_closure_double_ref(factor: u32, mut out: DisjointSlice<u32>) {
+        let idx = thread::index_1d();
+        let idx_raw = idx.get();
+        if let Some(out_elem) = out.get_mut(idx) {
+            let closure = |x: u32| x * factor;
+            let r = &closure;
+            let rr = &r;
+            *out_elem = (**rr)(idx_raw as u32);
+        }
+    }
 }
 
 // =============================================================================
@@ -456,5 +524,88 @@ fn main() {
         println!("  ✓ PASSED\n");
     }
 
+    // =========================================================================
+    // Test 9: Wrapper method named call with closure generic
+    // =========================================================================
+    println!("Test 9: Wrapper::call with closure generic");
+    {
+        const N: usize = 8;
+        let factor: u32 = 6;
+        let mut out_dev = DeviceBuffer::<u32>::zeroed(&stream, N).unwrap();
+
+        println!("  factor = {}", factor);
+
+        module
+            .test_wrapper_method_named_call(
+                (stream).as_ref(),
+                LaunchConfig::for_num_elems(N as u32),
+                factor,
+                &mut out_dev,
+            )
+            .expect("Kernel launch failed");
+
+        let out: Vec<u32> = out_dev.to_host_vec(&stream).unwrap();
+        let expected: Vec<u32> = (0..N).map(|i| (i as u32) * factor + 1).collect();
+
+        println!("  Output:   {:?}", out);
+        println!("  Expected: {:?}", expected);
+        assert_eq!(out, expected);
+        println!("  ✓ PASSED\n");
+    }
+
+    // =========================================================================
+    // Test 10: genuine closure called through a struct field projection
+    // =========================================================================
+    println!("Test 10: closure via field projection");
+    {
+        const N: usize = 8;
+        let factor: u32 = 5;
+        let mut out_dev = DeviceBuffer::<u32>::zeroed(&stream, N).unwrap();
+
+        module
+            .test_closure_field_projection(
+                (stream).as_ref(),
+                LaunchConfig::for_num_elems(N as u32),
+                factor,
+                &mut out_dev,
+            )
+            .expect("Kernel launch failed");
+
+        let out: Vec<u32> = out_dev.to_host_vec(&stream).unwrap();
+        let expected: Vec<u32> = (0..N).map(|i| (i as u32) * factor).collect();
+
+        println!("  Output:   {:?}", out);
+        println!("  Expected: {:?}", expected);
+        assert_eq!(out, expected);
+        println!("  ✓ PASSED\n");
+    }
+
+    // =========================================================================
+    // Test 11: genuine closure called through a double reference
+    // =========================================================================
+    println!("Test 11: closure via double reference");
+    {
+        const N: usize = 8;
+        let factor: u32 = 7;
+        let mut out_dev = DeviceBuffer::<u32>::zeroed(&stream, N).unwrap();
+
+        module
+            .test_closure_double_ref(
+                (stream).as_ref(),
+                LaunchConfig::for_num_elems(N as u32),
+                factor,
+                &mut out_dev,
+            )
+            .expect("Kernel launch failed");
+
+        let out: Vec<u32> = out_dev.to_host_vec(&stream).unwrap();
+        let expected: Vec<u32> = (0..N).map(|i| (i as u32) * factor).collect();
+
+        println!("  Output:   {:?}", out);
+        println!("  Expected: {:?}", expected);
+        assert_eq!(out, expected);
+        println!("  ✓ PASSED\n");
+    }
+
     println!("=== All device closure tests passed! ===");
 }