From 696ef741a6b654a4fe9f46bc47b502049e919aa9 Mon Sep 17 00:00:00 2001
From: Lucas Steuernagel <38472950+LucasSte@users.noreply.github.com>
Date: Wed, 13 Mar 2024 10:35:55 -0300
Subject: [PATCH] [SOL] Update SBF call ABI (#91)

---
 compiler/rustc_target/src/abi/call/sbf.rs | 32 +++++++++++++++--------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/compiler/rustc_target/src/abi/call/sbf.rs b/compiler/rustc_target/src/abi/call/sbf.rs
index d762c7fae9734..9c35e9a54bbdc 100644
--- a/compiler/rustc_target/src/abi/call/sbf.rs
+++ b/compiler/rustc_target/src/abi/call/sbf.rs
@@ -1,23 +1,33 @@
 // see https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/BPF/BPFCallingConv.td
-use crate::abi::call::{ArgAbi, FnAbi};
+use crate::abi::call::{ArgAbi, FnAbi, Reg, Uniform};
 
 fn classify_ret<Ty>(ret: &mut ArgAbi<'_, Ty>) {
-    if ret.layout.is_aggregate() || ret.layout.size.bits() > 64 {
-        if ret.layout.size.bits() != 128 {
-            ret.make_indirect();
-        }
-    } else {
+    let size = ret.layout.size;
+    let bits = size.bits();
+    if !ret.layout.is_aggregate() && bits <= 64 {
         ret.extend_integer_width_to(64);
+        return;
+    }
+
+    if bits <= 128 {
+        ret.cast_to(Uniform {unit: Reg::i64(), total: size});
+    } else {
+        ret.make_indirect();
     }
 }
 
 fn classify_arg<Ty>(arg: &mut ArgAbi<'_, Ty>) {
-    if arg.layout.is_aggregate() || arg.layout.size.bits() > 64 {
-        if arg.layout.size.bits() != 128 {
-            arg.make_indirect();
-        }
-    } else {
+    let size = arg.layout.size;
+    let bits = size.bits();
+    if !arg.layout.is_aggregate() && bits <= 64 {
         arg.extend_integer_width_to(64);
+        return;
+    }
+
+    if bits <= 128 {
+        arg.cast_to(Uniform {unit: Reg::i64(), total: size});
+    } else {
+        arg.make_indirect();
     }
 }