support assigning to a mutable reference in cube macro (#67)

crates/cubecl-core/tests/frontend/assign.rs

@@ -28,6 +28,12 @@ pub fn assign_vectorized(y: UInt) -> UInt {
     x + y
 }
 
+#[cube]
+pub fn assign_deref(y: &mut UInt) -> UInt {
+    *y = UInt::new(1);
+    *y
+}
+
 mod tests {
     use super::*;
     use cubecl_core::{
@@ -81,6 +87,18 @@ mod tests {
         assert_eq!(scope.operations, inline_macro_ref_assign_vectorized());
     }
 
+    #[test]
+    fn cube_assign_deref_test() {
+        let mut context = CubeContext::root();
+
+        let y = context.create_local(Item::new(UInt::as_elem()));
+        assign_deref::__expand(&mut context, y.into());
+
+        let scope = context.into_scope();
+
+        assert_eq!(scope.operations, inline_macro_ref_assign_deref());
+    }
+
     fn inline_macro_ref_mut_assign() -> Vec<Operation> {
         let context = CubeContext::root();
 
@@ -156,4 +174,16 @@ mod tests {
 
         scope.operations
     }
+
+    fn inline_macro_ref_assign_deref() -> Vec<Operation> {
+        let context = CubeContext::root();
+        let mut scope = context.into_scope();
+        let y = scope.create_local(Item::new(Elem::UInt));
+
+        let one: Variable = 1u32.into();
+
+        cpa!(scope, y = one);
+
+        scope.operations
+    }
 }

crates/cubecl-macros/src/codegen_function/variable.rs

@@ -150,19 +150,7 @@ pub(crate) fn codegen_assign(
                 }
             }
         }
-        syn::Expr::Path(_) => {
-            let lhs = codegen_expr(&assign.left, loop_level, variable_tracker);
-            let rhs = codegen_expr(&assign.right, loop_level, variable_tracker);
-
-            quote::quote! {
-                {
-                    let _assign_lhs = #lhs;
-                    let _assign_rhs = #rhs;
-                    cubecl::frontend::assign::expand(context, _assign_rhs, _assign_lhs)
-                }
-            }
-        }
-        syn::Expr::Field(_) => {
+        syn::Expr::Unary(_) | syn::Expr::Field(_) | syn::Expr::Path(_) => {
             let lhs = codegen_expr(&assign.left, loop_level, variable_tracker);
             let rhs = codegen_expr(&assign.right, loop_level, variable_tracker);