riscv64: Extend integer args passed host-to-guest

The x86_64 ABI does not specify the content of the upper bits of integer
arguments narrower than 64 bits. RV64 requires sign-appropriate
extension to 32 bits (sign-extension for signed types, zero-extension
for unsigned), followed by sign-extension to 64 bits.

Bug: 330396865
Test: CtsJniTestCases and berberis_run_host_tests
Change-Id: I50b35e0ddd9ad8d7bba16982a906fc9d7a529bb0

5 files changed, 256 insertions, 31 deletions

diff --git a/code_gen_lib/code_gen_lib_riscv64_test.cc b/code_gen_lib/code_gen_lib_riscv64_test.cc
index db81612c..789cf5fe 100644
--- a/code_gen_lib/code_gen_lib_riscv64_test.cc
+++ b/code_gen_lib/code_gen_lib_riscv64_test.cc
@@ -147,24 +147,163 @@ TEST(CodeGenLib, GenWrapGuestFunction) {
   ASSERT_TRUE(g_called);
 }
 
+void Run10UInt8(GuestAddr pc, GuestArgumentBuffer* buf) {
+  ASSERT_EQ(ToGuestAddr(&g_insn), pc);
+  ASSERT_NE(buf, nullptr);
+  ASSERT_EQ(buf->argc, 8);
+  ASSERT_EQ(buf->stack_argc, 16);
+  ASSERT_EQ(buf->resc, 1);
+  ASSERT_EQ(buf->argv[0], 0U);
+  ASSERT_EQ(buf->argv[1], 0xffU);
+  ASSERT_EQ(buf->argv[2], 2U);
+  ASSERT_EQ(buf->argv[3], 3U);
+  ASSERT_EQ(buf->argv[4], 4U);
+  ASSERT_EQ(buf->argv[5], 5U);
+  ASSERT_EQ(buf->argv[6], 6U);
+  ASSERT_EQ(buf->argv[7], 0xf9U);
+  ASSERT_EQ(buf->stack_argv[0], 0xf8U);
+  ASSERT_EQ(buf->stack_argv[1], 9U);
+  buf->argv[0] = 0xf6;
+}
+
+TEST(CodeGenLib, GenWrapGuestFunction_Run10UInt8) {
+  MachineCode machine_code;
+
+  GenWrapGuestFunction(
+      &machine_code, ToGuestAddr(&g_insn), "zzzzzzzzzzz", AsHostCode(Run10UInt8), "Run10UInt8");
+
+  ScopedExecRegion exec(&machine_code);
+
+  using Func = uint8_t(
+      uint8_t, uint8_t, uint8_t, uint8_t, uint8_t, uint8_t, uint8_t, uint8_t, uint8_t, uint8_t);
+  uint8_t res = exec.get<Func>()(0, 0xff, 2, 3, 4, 5, 6, 0xf9, 0xf8, 9);
+  ASSERT_EQ(res, 0xF6u);
+}
+
+void Run10Int8(GuestAddr pc, GuestArgumentBuffer* buf) {
+  ASSERT_EQ(ToGuestAddr(&g_insn), pc);
+  ASSERT_NE(buf, nullptr);
+  ASSERT_EQ(buf->argc, 8);
+  ASSERT_EQ(buf->stack_argc, 16);
+  ASSERT_EQ(buf->resc, 1);
+  ASSERT_EQ(buf->argv[0], 0U);
+  ASSERT_EQ(buf->argv[1], 0xffff'ffff'ffff'ffffULL);
+  ASSERT_EQ(buf->argv[2], 2U);
+  ASSERT_EQ(buf->argv[3], 3U);
+  ASSERT_EQ(buf->argv[4], 4U);
+  ASSERT_EQ(buf->argv[5], 5U);
+  ASSERT_EQ(buf->argv[6], 6U);
+  ASSERT_EQ(buf->argv[7], 0xffff'ffff'ffff'fff9ULL);
+  ASSERT_EQ(buf->stack_argv[0], 0xffff'ffff'ffff'fff8ULL);
+  ASSERT_EQ(buf->stack_argv[1], 9U);
+  buf->argv[0] = 0xffff'ffff'ffff'fff6;
+}
+
+TEST(CodeGenLib, GenWrapGuestFunction_Run10Int8) {
+  MachineCode machine_code;
+
+  GenWrapGuestFunction(
+      &machine_code, ToGuestAddr(&g_insn), "bbbbbbbbbbb", AsHostCode(Run10Int8), "Run10Int8");
+
+  ScopedExecRegion exec(&machine_code);
+
+  using Func =
+      int8_t(int8_t, int8_t, int8_t, int8_t, int8_t, int8_t, int8_t, int8_t, int8_t, int8_t);
+  int8_t res = exec.get<Func>()(0, -1, 2, 3, 4, 5, 6, -7, -8, 9);
+  ASSERT_EQ(res, -10);
+}
+
+void Run10UInt16(GuestAddr pc, GuestArgumentBuffer* buf) {
+  ASSERT_EQ(ToGuestAddr(&g_insn), pc);
+  ASSERT_NE(buf, nullptr);
+  ASSERT_EQ(buf->argc, 8);
+  ASSERT_EQ(buf->stack_argc, 16);
+  ASSERT_EQ(buf->resc, 1);
+  ASSERT_EQ(buf->argv[0], 0U);
+  ASSERT_EQ(buf->argv[1], 0xffffU);
+  ASSERT_EQ(buf->argv[2], 2U);
+  ASSERT_EQ(buf->argv[3], 3U);
+  ASSERT_EQ(buf->argv[4], 4U);
+  ASSERT_EQ(buf->argv[5], 5U);
+  ASSERT_EQ(buf->argv[6], 6U);
+  ASSERT_EQ(buf->argv[7], 0xfff9U);
+  ASSERT_EQ(buf->stack_argv[0], 0xfff8U);
+  ASSERT_EQ(buf->stack_argv[1], 9U);
+  buf->argv[0] = 0xfff6;
+}
+
+TEST(CodeGenLib, GenWrapGuestFunction_Run10UInt16) {
+  MachineCode machine_code;
+
+  GenWrapGuestFunction(
+      &machine_code, ToGuestAddr(&g_insn), "ccccccccccc", AsHostCode(Run10UInt16), "Run10UInt16");
+
+  ScopedExecRegion exec(&machine_code);
+
+  using Func = uint16_t(uint16_t,
+                        uint16_t,
+                        uint16_t,
+                        uint16_t,
+                        uint16_t,
+                        uint16_t,
+                        uint16_t,
+                        uint16_t,
+                        uint16_t,
+                        uint16_t);
+  uint16_t res = exec.get<Func>()(0, 0xffff, 2, 3, 4, 5, 6, 0xfff9, 0xfff8, 9);
+  ASSERT_EQ(res, 0xfff6U);
+}
+
+void Run10Int16(GuestAddr pc, GuestArgumentBuffer* buf) {
+  ASSERT_EQ(ToGuestAddr(&g_insn), pc);
+  ASSERT_NE(buf, nullptr);
+  ASSERT_EQ(buf->argc, 8);
+  ASSERT_EQ(buf->stack_argc, 16);
+  ASSERT_EQ(buf->resc, 1);
+  ASSERT_EQ(buf->argv[0], 0U);
+  ASSERT_EQ(buf->argv[1], 0xffff'ffff'ffff'ffffULL);
+  ASSERT_EQ(buf->argv[2], 2U);
+  ASSERT_EQ(buf->argv[3], 3U);
+  ASSERT_EQ(buf->argv[4], 4U);
+  ASSERT_EQ(buf->argv[5], 5U);
+  ASSERT_EQ(buf->argv[6], 6U);
+  ASSERT_EQ(buf->argv[7], 0xffff'ffff'ffff'fff9ULL);
+  ASSERT_EQ(buf->stack_argv[0], 0xffff'ffff'ffff'fff8ULL);
+  ASSERT_EQ(buf->stack_argv[1], 9U);
+  buf->argv[0] = 0xffff'ffff'ffff'fff6;
+}
+
+TEST(CodeGenLib, GenWrapGuestFunction_Run10Int16) {
+  MachineCode machine_code;
+
+  GenWrapGuestFunction(
+      &machine_code, ToGuestAddr(&g_insn), "sssssssssss", AsHostCode(Run10Int16), "Run10Int16");
+
+  ScopedExecRegion exec(&machine_code);
+
+  using Func = int16_t(
+      int16_t, int16_t, int16_t, int16_t, int16_t, int16_t, int16_t, int16_t, int16_t, int16_t);
+  int16_t res = exec.get<Func>()(0, -1, 2, 3, 4, 5, 6, -7, -8, 9);
+  ASSERT_EQ(res, -10);
+}
+
 void Run10Int(GuestAddr pc, GuestArgumentBuffer* buf) {
-  ASSERT_EQ(pc, ToGuestAddr(&g_insn));
-  ASSERT_NE(nullptr, buf);
-  ASSERT_EQ(8, buf->argc);
-  ASSERT_EQ(16, buf->stack_argc);
-  ASSERT_EQ(1, buf->resc);
-  // For 32-bit parameters, only least-significant bits are defined!
-  ASSERT_EQ(0u, static_cast<uint32_t>(buf->argv[0]));
-  ASSERT_EQ(1u, static_cast<uint32_t>(buf->argv[1]));
-  ASSERT_EQ(2u, static_cast<uint32_t>(buf->argv[2]));
-  ASSERT_EQ(3u, static_cast<uint32_t>(buf->argv[3]));
-  ASSERT_EQ(4u, static_cast<uint32_t>(buf->argv[4]));
-  ASSERT_EQ(5u, static_cast<uint32_t>(buf->argv[5]));
-  ASSERT_EQ(6u, static_cast<uint32_t>(buf->argv[6]));
-  ASSERT_EQ(7u, static_cast<uint32_t>(buf->argv[7]));
-  ASSERT_EQ(8u, static_cast<uint32_t>(buf->stack_argv[0]));
-  ASSERT_EQ(9u, static_cast<uint32_t>(buf->stack_argv[1]));
-  buf->argv[0] = 45;
+  ASSERT_EQ(ToGuestAddr(&g_insn), pc);
+  ASSERT_NE(buf, nullptr);
+  ASSERT_EQ(buf->argc, 8);
+  ASSERT_EQ(buf->stack_argc, 16);
+  ASSERT_EQ(buf->resc, 1);
+  ASSERT_EQ(buf->argv[0], 0U);
+  ASSERT_EQ(buf->argv[1], 0xffff'ffff'ffff'ffffULL);
+  ASSERT_EQ(buf->argv[2], 2U);
+  ASSERT_EQ(buf->argv[3], 3U);
+  ASSERT_EQ(buf->argv[4], 4U);
+  ASSERT_EQ(buf->argv[5], 5U);
+  ASSERT_EQ(buf->argv[6], 6U);
+  ASSERT_EQ(buf->argv[7], 0xffff'ffff'ffff'fff9ULL);
+  ASSERT_EQ(buf->stack_argv[0], 0xffff'ffff'ffff'fff8ULL);
+  ASSERT_EQ(buf->stack_argv[1], 9U);
+  buf->argv[0] = 0xffff'ffff'ffff'fff6;
 }
 
 TEST(CodeGenLib, GenWrapGuestFunction_Run10Int) {
@@ -176,8 +315,8 @@ TEST(CodeGenLib, GenWrapGuestFunction_Run10Int) {
   ScopedExecRegion exec(&machine_code);
 
   using Func = int(int, int, int, int, int, int, int, int, int, int);
-  int res = exec.get<Func>()(0, 1, 2, 3, 4, 5, 6, 7, 8, 9);
-  ASSERT_EQ(45, res);
+  int res = exec.get<Func>()(0, -1, 2, 3, 4, 5, 6, -7, -8, 9);
+  ASSERT_EQ(res, -10);
 }
 
 void Run10Fp(GuestAddr pc, GuestArgumentBuffer* buf) {
diff --git a/code_gen_lib/gen_wrapper_riscv64_to_x86_64.cc b/code_gen_lib/gen_wrapper_riscv64_to_x86_64.cc
index 5ec2be55..65c2811b 100644
--- a/code_gen_lib/gen_wrapper_riscv64_to_x86_64.cc
+++ b/code_gen_lib/gen_wrapper_riscv64_to_x86_64.cc
@@ -28,6 +28,31 @@ namespace berberis {
 
 using x86_64::Assembler;
 
+namespace {
+
+void ExtendIntArg(MacroAssembler<Assembler>& as,
+                  char type,
+                  Assembler::Register dst,
+                  Assembler::Register src) {
+  if (type == 'z') {
+    as.Movzxbq(dst, src);
+  } else if (type == 'b') {
+    as.Movsxbq(dst, src);
+  } else if (type == 's') {
+    as.Movsxwq(dst, src);
+  } else if (type == 'c') {
+    as.Movzxwq(dst, src);
+  } else if (type == 'i') {
+    as.Movsxlq(dst, src);
+  } else if (type == 'p' || type == 'l') {
+    // Do nothing. The parameter is already 64 bits.
+  } else {
+    FATAL("non-integer signature char '%c'", type);
+  }
+}
+
+}  // namespace
+
 void GenWrapGuestFunction(MachineCode* mc,
                           GuestAddr pc,
                           const char* signature,
@@ -80,7 +105,8 @@ void GenWrapGuestFunction(MachineCode* mc,
   int stack_argc = 0;
   int host_stack_argc = 0;
   for (size_t i = 1; signature[i] != '\0'; ++i) {
-    if (signature[i] == 'i' || signature[i] == 'p' || signature[i] == 'l') {
+    if (signature[i] == 'z' || signature[i] == 'b' || signature[i] == 's' || signature[i] == 'c' ||
+        signature[i] == 'i' || signature[i] == 'p' || signature[i] == 'l') {
       static constexpr Assembler::Register kParamRegs[] = {
           Assembler::rdi,
           Assembler::rsi,
@@ -90,16 +116,19 @@ void GenWrapGuestFunction(MachineCode* mc,
           Assembler::r9,
       };
       if (argc < static_cast<int>(std::size(kParamRegs))) {
+        ExtendIntArg(as, signature[i], kParamRegs[argc], kParamRegs[argc]);
         as.Movq({.base = Assembler::rsp, .disp = kArgvOffset + argc * 8}, kParamRegs[argc]);
       } else if (argc < 8) {
         as.Movq(Assembler::rax,
                 {.base = Assembler::rsp, .disp = params_offset + host_stack_argc * 8});
         ++host_stack_argc;
+        ExtendIntArg(as, signature[i], Assembler::rax, Assembler::rax);
         as.Movq({.base = Assembler::rsp, .disp = kArgvOffset + argc * 8}, Assembler::rax);
       } else {
         as.Movq(Assembler::rax,
                 {.base = Assembler::rsp, .disp = params_offset + host_stack_argc * 8});
         ++host_stack_argc;
+        ExtendIntArg(as, signature[i], Assembler::rax, Assembler::rax);
         as.Movq({.base = Assembler::rsp, .disp = kStackArgvOffset + stack_argc * 8},
                 Assembler::rax);
         ++stack_argc;
@@ -151,7 +180,8 @@ void GenWrapGuestFunction(MachineCode* mc,
   as.Movl({.base = Assembler::rsp, .disp = kStackArgcOffset}, stack_argc * 8);
 
   // Set resc.
-  if (signature[0] == 'i' || signature[0] == 'p' || signature[0] == 'l') {
+  if (signature[0] == 'z' || signature[0] == 'b' || signature[0] == 's' || signature[0] == 'c' ||
+      signature[0] == 'i' || signature[0] == 'p' || signature[0] == 'l') {
     as.Movl({.base = Assembler::rsp, .disp = kRescOffset}, 1);
     as.Movl({.base = Assembler::rsp, .disp = kFpRescOffset}, 0);
   } else if (signature[0] == 'f' || signature[0] == 'd') {
@@ -169,7 +199,10 @@ void GenWrapGuestFunction(MachineCode* mc,
   as.Call(guest_runner);
 
   // Get the result.
-  if (signature[0] == 'i' || signature[0] == 'p' || signature[0] == 'l') {
+  if (signature[0] == 'z' || signature[0] == 'b' || signature[0] == 's' || signature[0] == 'c' ||
+      signature[0] == 'i' || signature[0] == 'p' || signature[0] == 'l') {
+    // It is not necessary to unbox integer return values. The callee will truncate rax to only
+    // retrieve the bits appropriate for the return type.
     as.Movq(Assembler::rax, {.base = Assembler::rsp, .disp = kArgvOffset});
   } else if (signature[0] == 'f') {
     // Only take the lower 32 bits of the result register because floats are 1-extended (NaN boxed)
diff --git a/guest_abi/guest_function_wrapper_signature_test.cc b/guest_abi/guest_function_wrapper_signature_test.cc
index 6ac05860..a136779f 100644
--- a/guest_abi/guest_function_wrapper_signature_test.cc
+++ b/guest_abi/guest_function_wrapper_signature_test.cc
@@ -22,8 +22,10 @@ namespace berberis {
 
 namespace {
 
-static_assert('i' == kGuestFunctionWrapperSignatureChar<bool>);
-static_assert('i' == kGuestFunctionWrapperSignatureChar<char>);
+static_assert('z' == kGuestFunctionWrapperSignatureChar<bool>);
+static_assert('b' == kGuestFunctionWrapperSignatureChar<char>);
+static_assert('s' == kGuestFunctionWrapperSignatureChar<int16_t>);
+static_assert('c' == kGuestFunctionWrapperSignatureChar<uint16_t>);
 static_assert('i' == kGuestFunctionWrapperSignatureChar<int>);
 static_assert('l' == kGuestFunctionWrapperSignatureChar<long long>);
 static_assert('p' == kGuestFunctionWrapperSignatureChar<void*>);
diff --git a/guest_abi/include/berberis/guest_abi/guest_function_wrapper_signature.h b/guest_abi/include/berberis/guest_abi/guest_function_wrapper_signature.h
index f879e37f..382e83ae 100644
--- a/guest_abi/include/berberis/guest_abi/guest_function_wrapper_signature.h
+++ b/guest_abi/include/berberis/guest_abi/guest_function_wrapper_signature.h
@@ -27,8 +27,12 @@ namespace berberis {
 //
 // Supported types:
 //   'v': void (as return type)
-//   'i': integer and enum types <= 32bit
-//   'l': integer and enum types == 64bit
+//   'z': unsigned integer and enum types == 8bit (jboolean equivalent)
+//   'b': signed integer and enum types == 8bit (jbyte equivalent)
+//   's': signed integer and enum types == 16bit (jshort equivalent)
+//   'c': unsigned integer types == 16bit (jchar equivalent)
+//   'i': integer and enum types == 32bit (jint equivalent)
+//   'l': integer and enum types == 64bit (jfloat equivalent)
 //   'p': pointers (to objects and functions but not to members)
 //   'f': float (floating point 32 bits)
 //   'd': double (floating point 64 bits)
@@ -47,7 +51,42 @@ class kGuestFunctionWrapperSignatureCharHelper {
   }
 
   template <typename Type,
-            std::enable_if_t<(std::is_integral_v<Type> || std::is_enum_v<Type>)&&sizeof(Type) <=
+            std::enable_if_t<(std::is_integral_v<Type> ||
+                              std::is_enum_v<Type>)&&std::is_unsigned_v<Type> &&
+                                 sizeof(Type) == sizeof(uint8_t),
+                             int> = 0>
+  static constexpr char Value() {
+    return 'z';
+  }
+
+  template <
+      typename Type,
+      std::enable_if_t<(std::is_integral_v<Type> || std::is_enum_v<Type>)&&std::is_signed_v<Type> &&
+                           sizeof(Type) == sizeof(int8_t),
+                       int> = 0>
+  static constexpr char Value() {
+    return 'b';
+  }
+
+  template <
+      typename Type,
+      std::enable_if_t<(std::is_integral_v<Type> || std::is_enum_v<Type>)&&std::is_signed_v<Type> &&
+                           sizeof(Type) == sizeof(int16_t),
+                       int> = 0>
+  static constexpr char Value() {
+    return 's';
+  }
+
+  template <typename Type,
+            std::enable_if_t<std::is_integral_v<Type> && std::is_unsigned_v<Type> &&
+                                 sizeof(Type) == sizeof(uint16_t),
+                             int> = 0>
+  static constexpr char Value() {
+    return 'c';
+  }
+
+  template <typename Type,
+            std::enable_if_t<(std::is_integral_v<Type> || std::is_enum_v<Type>)&&sizeof(Type) ==
                                  sizeof(int32_t),
                              int> = 0>
   static constexpr char Value() {
diff --git a/jni/jni_trampolines.cc b/jni/jni_trampolines.cc
index cea6acba..c2ca4d6e 100644
--- a/jni/jni_trampolines.cc
+++ b/jni/jni_trampolines.cc
@@ -43,9 +43,13 @@ char ConvertDalvikTypeCharToWrapperTypeChar(char c) {
     case 'V':  // void
       return 'v';
     case 'Z':  // boolean
+      return 'z';
     case 'B':  // byte
+      return 'b';
     case 'S':  // short
+      return 's';
     case 'C':  // char
+      return 'c';
     case 'I':  // int
       return 'i';
     case 'L':  // class object - pointer
@@ -132,10 +136,18 @@ std::vector<jvalue> ConvertVAList(JNIEnv* env, jmethodID methodID, GuestVAListPa
     jvalue& arg = result[i];
     char c = short_signature[i];
     switch (c) {
-      case 'Z':  // boolean (u8) - passed as int
-      case 'B':  // byte (i8) - passed as int
-      case 'S':  // short (i16) - passed as int
-      case 'C':  // char (u16) - passed as int
+      case 'Z':  // boolean (u8)
+        arg.z = params.GetParam<uint8_t>();
+        break;
+      case 'B':  // byte (i8)
+        arg.b = params.GetParam<int8_t>();
+        break;
+      case 'S':  // short (i16)
+        arg.s = params.GetParam<int16_t>();
+        break;
+      case 'C':  // char (u16)
+        arg.c = params.GetParam<uint16_t>();
+        break;
       case 'I':  // int (i32)
         arg.i = params.GetParam<int32_t>();
         break;