[mono] Add support for SwiftIndirectResult in Mono JIT and Interpreter (#104111)

* Add support for SwiftIndirectResult at P/Invoke boundary

* Update struct lowering

* Introduce arm64 support

* Introduce amd64 support

* Enable runtime tests

* Free RAX if used for the call

* Update src/mono/mono/mini/mini-amd64.c

---------

Co-authored-by: Matous Kozak <[email protected]>

Author: kotlarmilos

src/mono/mono/metadata/marshal.c

@@ -132,6 +132,7 @@ static GENERATE_TRY_GET_CLASS_WITH_CACHE (unmanaged_callconv_attribute, "System.
 
 GENERATE_TRY_GET_CLASS_WITH_CACHE (swift_error, "System.Runtime.InteropServices.Swift", "SwiftError")
 GENERATE_TRY_GET_CLASS_WITH_CACHE (swift_self, "System.Runtime.InteropServices.Swift", "SwiftSelf")
+GENERATE_TRY_GET_CLASS_WITH_CACHE (swift_indirect_result, "System.Runtime.InteropServices.Swift", "SwiftIndirectResult")
 
 static gboolean type_is_blittable (MonoType *type);
 
@@ -3698,8 +3699,9 @@ mono_marshal_get_native_wrapper (MonoMethod *method, gboolean check_exceptions,
 		if (mono_method_signature_has_ext_callconv (csig, MONO_EXT_CALLCONV_SWIFTCALL)) {
 			MonoClass *swift_self = mono_class_try_get_swift_self_class ();
 			MonoClass *swift_error = mono_class_try_get_swift_error_class ();
+			MonoClass *swift_indirect_result = mono_class_try_get_swift_indirect_result_class ();
 			MonoClass *swift_error_ptr = mono_class_create_ptr (m_class_get_this_arg (swift_error));
-			int swift_error_args = 0, swift_self_args = 0;
+			int swift_error_args = 0, swift_self_args = 0, swift_indirect_result_args = 0;
 			for (int i = 0; i < method->signature->param_count; ++i) {
 				MonoClass *param_klass = mono_class_from_mono_type_internal (method->signature->params [i]);
 				if (param_klass) {
@@ -3711,16 +3713,18 @@ mono_marshal_get_native_wrapper (MonoMethod *method, gboolean check_exceptions,
 						swift_error_args++;
 					} else if (param_klass == swift_self) {
 						swift_self_args++;
+					} else if (param_klass == swift_indirect_result) {
+						swift_indirect_result_args++;
 					} else if (!type_is_blittable (method->signature->params [i]) || m_class_is_simd_type (param_klass)) {
-						swift_error_args = swift_self_args = 0;
+						swift_error_args = swift_self_args = swift_indirect_result_args = 0;
 						mono_error_set_generic_error (emitted_error, "System", "InvalidProgramException", "Passing non-blittable types to a P/Invoke with the Swift calling convention is unsupported.");
 						break;
 					}
 				}
 			}
 
-			if (swift_self_args > 1 || swift_error_args > 1) {
-				mono_error_set_generic_error (emitted_error, "System", "InvalidProgramException", "Method signature contains multiple SwiftSelf or SwiftError arguments.");
+			if (swift_self_args > 1 || swift_error_args > 1 || swift_indirect_result_args > 1) {
+				mono_error_set_generic_error (emitted_error, "System", "InvalidProgramException", "Method signature contains multiple SwiftSelf, SwiftError, SwiftIndirectResult arguments.");
 			}
 
 			if (!is_ok (emitted_error)) {

src/mono/mono/metadata/marshal.h

@@ -742,6 +742,7 @@ mono_marshal_get_mono_callbacks_for_ilgen (void);
 
 GENERATE_TRY_GET_CLASS_WITH_CACHE_DECL (swift_self)
 GENERATE_TRY_GET_CLASS_WITH_CACHE_DECL (swift_error)
+GENERATE_TRY_GET_CLASS_WITH_CACHE_DECL (swift_indirect_result)
 
 typedef struct {
 	gboolean by_reference;

src/mono/mono/mini/interp/transform.c

@@ -3407,6 +3407,7 @@ interp_emit_swiftcall_struct_lowering (TransformData *td, MonoMethodSignature *c
 	int align;
 	MonoClass *swift_self = mono_class_try_get_swift_self_class ();
 	MonoClass *swift_error = mono_class_try_get_swift_error_class ();
+	MonoClass *swift_indirect_result = mono_class_try_get_swift_indirect_result_class ();
 	/*
 	* Go through the lowered arguments, if the argument is a struct,
 	* we need to replace it with a sequence of lowered arguments.
@@ -3415,8 +3416,8 @@ interp_emit_swiftcall_struct_lowering (TransformData *td, MonoMethodSignature *c
 	for (int idx_param = 0; idx_param < csignature->param_count; ++idx_param) {
 		MonoType *ptype = csignature->params [idx_param];
 		MonoClass *klass = mono_class_from_mono_type_internal (ptype);
-		// SwiftSelf and SwiftError are special cases where we need to preserve the class information for the codegen to handle them correctly.
-		if (mono_type_is_struct (ptype) && !(klass == swift_self || klass == swift_error)) {
+		// SwiftSelf, SwiftError, and SwiftIndirectResult are special cases where we need to preserve the class information for the codegen to handle them correctly.
+		if (mono_type_is_struct (ptype) && !(klass == swift_self || klass == swift_error || klass == swift_indirect_result)) {
 			SwiftPhysicalLowering lowered_swift_struct = mono_marshal_get_swift_physical_lowering (ptype, FALSE);
 			if (!lowered_swift_struct.by_reference) {
 				for (uint32_t idx_lowered = 0; idx_lowered < lowered_swift_struct.num_lowered_elements; ++idx_lowered) {

src/mono/mono/mini/method-to-ir.c

@@ -7534,6 +7534,7 @@ mono_method_to_ir (MonoCompile *cfg, MonoMethod *method, MonoBasicBlock *start_b
 				uint32_t new_param_count = 0;
 				MonoClass *swift_self = mono_class_try_get_swift_self_class ();
 				MonoClass *swift_error = mono_class_try_get_swift_error_class ();
+				MonoClass *swift_indirect_result = mono_class_try_get_swift_indirect_result_class ();
 				/*
 				 * Go through the lowered arguments, if the argument is a struct, 
 				 * we need to replace it with a sequence of lowered arguments.
@@ -7543,8 +7544,8 @@ mono_method_to_ir (MonoCompile *cfg, MonoMethod *method, MonoBasicBlock *start_b
 					MonoType *ptype = fsig->params [idx_param];
 					MonoClass *klass = mono_class_from_mono_type_internal (ptype);
 
-					// SwiftSelf and SwiftError are special cases where we need to preserve the class information for the codegen to handle them correctly.
-					if (mono_type_is_struct (ptype) && !(klass == swift_self || klass == swift_error)) {
+					// SwiftSelf, SwiftError, and SwiftIndirectResult are special cases where we need to preserve the class information for the codegen to handle them correctly.
+					if (mono_type_is_struct (ptype) && !(klass == swift_self || klass == swift_error || klass == swift_indirect_result)) {
 						SwiftPhysicalLowering lowered_swift_struct = mono_marshal_get_swift_physical_lowering (ptype, FALSE);
 						if (!lowered_swift_struct.by_reference) {
 							// Create a new local variable to store the base address of the struct

src/mono/mono/mini/mini-amd64.c

@@ -876,6 +876,7 @@ get_call_info (MonoMemPool *mp, MonoMethodSignature *sig)
 	cinfo->nargs = n;
 	cinfo->gsharedvt = mini_is_gsharedvt_variable_signature (sig);
 	cinfo->swift_error_index = -1;
+	cinfo->swift_indirect_result_index = -1;
 
 	gr = 0;
 	fr = 0;
@@ -1014,17 +1015,23 @@ get_call_info (MonoMemPool *mp, MonoMethodSignature *sig)
 		if (mono_method_signature_has_ext_callconv (sig, MONO_EXT_CALLCONV_SWIFTCALL)) {
 			MonoClass *swift_self = mono_class_try_get_swift_self_class ();
 			MonoClass *swift_error = mono_class_try_get_swift_error_class ();
+			MonoClass *swift_indirect_result = mono_class_try_get_swift_indirect_result_class ();
 			MonoClass *swift_error_ptr = mono_class_create_ptr (m_class_get_this_arg (swift_error));
 			MonoClass *klass = mono_class_from_mono_type_internal (sig->params [i]);
-			if (klass == swift_self && sig->pinvoke) {
+			if (klass == swift_indirect_result)
+				cinfo->swift_indirect_result_index = i;
+
+			if ((klass == swift_self || klass == swift_indirect_result) && sig->pinvoke) {
 				guint32 size = mini_type_stack_size_full (m_class_get_byval_arg (klass), NULL, sig->pinvoke && !sig->marshalling_disabled);
 				g_assert (size == 8);
 
 				ainfo->storage = ArgValuetypeInReg;
 				ainfo->pair_storage [0] = ArgInIReg;
 				ainfo->pair_storage [1] = ArgNone;
 				ainfo->nregs = 1;
-				ainfo->pair_regs [0] = GINT32_TO_UINT8 (AMD64_R13);
+				// The indirect result is passed in RAX on AMD64. However, since we don't prevent reg allocator to use RAX in m2n wrappers, the RAX is used to pass the function pointer.
+				// We use the R10 register instead and before the native call, the value is moved from R10 to RAX.
+				ainfo->pair_regs [0] = (klass == swift_self) ? GINT32_TO_UINT8 (AMD64_R13) : GINT32_TO_UINT8 (AMD64_R10);
 				ainfo->pair_size [0] = size;
 				continue;
 			} else if (klass == swift_error || klass == swift_error_ptr) {
@@ -4634,6 +4641,25 @@ amd64_handle_varargs_call (MonoCompile *cfg, guint8 *code, MonoCallInst *call, g
 #endif
 }
 
+static guint8*
+amd64_handle_swift_indirect_result (MonoCompile *cfg, guint8 *code, MonoCallInst *call)
+{
+	// Ideally, this should be in mono_arch_emit_prolog, but RAX may be used for the call, and it is required to free RAX.
+	if (mono_method_signature_has_ext_callconv (cfg->method->signature, MONO_EXT_CALLCONV_SWIFTCALL) && 
+		cfg->arch.cinfo->swift_indirect_result_index > -1) {
+
+		MonoInst *ins = (MonoInst*)call;
+		if (ins->sreg1 == AMD64_RAX) {
+			amd64_mov_reg_reg (code, AMD64_R11, AMD64_RAX, 8);
+			ins->sreg1 = AMD64_R11;
+		}
+		amd64_mov_reg_reg (code, AMD64_RAX, AMD64_R10, 8);
+	}
+
+	return code;
+}
+
+
 void
 mono_arch_output_basic_block (MonoCompile *cfg, MonoBasicBlock *bb)
 {
@@ -5619,6 +5645,9 @@ mono_arch_output_basic_block (MonoCompile *cfg, MonoBasicBlock *bb)
 			// FIXME Just like NT the direct cases are are not ideal.
 			amd64_mov_membase_reg (code, AMD64_RSP, -8, AMD64_RAX, 8);
 			code = amd64_handle_varargs_call (cfg, code, call, FALSE);
+#ifdef MONO_ARCH_HAVE_SWIFTCALL
+			code = amd64_handle_swift_indirect_result (cfg, code, call);
+#endif
 			amd64_jump_membase (code, AMD64_RSP, -8);
 #endif
 			ins->flags |= MONO_INST_GC_CALLSITE;
@@ -5644,6 +5673,9 @@ mono_arch_output_basic_block (MonoCompile *cfg, MonoBasicBlock *bb)
 			call = (MonoCallInst*)ins;
 
 			code = amd64_handle_varargs_call (cfg, code, call, FALSE);
+#ifdef MONO_ARCH_HAVE_SWIFTCALL
+			code = amd64_handle_swift_indirect_result (cfg, code, call);
+#endif
 			code = emit_call (cfg, call, code, MONO_JIT_ICALL_ZeroIsReserved);
 			ins->flags |= MONO_INST_GC_CALLSITE;
 			ins->backend.pc_offset = GPTRDIFF_TO_INT (code - cfg->native_code);
@@ -5664,6 +5696,9 @@ mono_arch_output_basic_block (MonoCompile *cfg, MonoBasicBlock *bb)
 			}
 
 			code = amd64_handle_varargs_call (cfg, code, call, TRUE);
+#ifdef MONO_ARCH_HAVE_SWIFTCALL
+			code = amd64_handle_swift_indirect_result (cfg, code, call);
+#endif
 			amd64_call_reg (code, ins->sreg1);
 			ins->flags |= MONO_INST_GC_CALLSITE;
 			ins->backend.pc_offset = GPTRDIFF_TO_INT (code - cfg->native_code);
@@ -8206,6 +8241,14 @@ MONO_RESTORE_WARNING
 			amd64_mov_membase_reg (code, cfg->vret_addr->inst_basereg, cfg->vret_addr->inst_offset, cinfo->ret.reg, 8);
 	}
 
+#ifdef MONO_ARCH_HAVE_SWIFTCALL
+	if (mono_method_signature_has_ext_callconv (sig, MONO_EXT_CALLCONV_SWIFTCALL) && 
+	    cfg->method->wrapper_type == MONO_WRAPPER_NATIVE_TO_MANAGED && 
+	    cfg->arch.cinfo->swift_indirect_result_index > -1) {
+		amd64_mov_reg_reg (code, AMD64_R10, AMD64_RAX, 8);
+	}
+#endif
+
 	/* Keep this in sync with emit_load_volatile_arguments */
 	for (guint i = 0; i < sig->param_count + sig->hasthis; ++i) {
 		ArgInfo *ainfo = cinfo->args + i;

src/mono/mono/mini/mini-amd64.h

@@ -330,6 +330,7 @@ struct CallInfo {
 	guint32 reg_usage;
 	guint32 freg_usage;
 	gint32 swift_error_index;
+	gint32 swift_indirect_result_index;
 	gboolean need_stack_align;
 	gboolean gsharedvt;
 	/* The index of the vret arg in the argument list */

src/mono/mono/mini/mini-arm64.c

@@ -1862,16 +1862,17 @@ get_call_info (MonoMemPool *mp, MonoMethodSignature *sig)
 		if (mono_method_signature_has_ext_callconv (sig, MONO_EXT_CALLCONV_SWIFTCALL)) {
 			MonoClass *swift_self = mono_class_try_get_swift_self_class ();
 			MonoClass *swift_error = mono_class_try_get_swift_error_class ();
+			MonoClass *swift_indirect_result = mono_class_try_get_swift_indirect_result_class ();
 			MonoClass *swift_error_ptr = mono_class_create_ptr (m_class_get_this_arg (swift_error));
 			MonoClass *klass = mono_class_from_mono_type_internal (sig->params [pindex]);
-			if (klass == swift_self && sig->pinvoke) {
+			if ((klass == swift_self || klass == swift_indirect_result) && sig->pinvoke) {
 				guint32 align;
 				MonoType *ptype = mini_get_underlying_type (sig->params [pindex]);
 				int size = mini_type_stack_size_full (ptype, &align, cinfo->pinvoke);
 				g_assert (size == 8);
 
 				ainfo->storage = ArgVtypeInIRegs;
-				ainfo->reg = ARMREG_R20;
+				ainfo->reg = (klass == swift_self) ? ARMREG_R20 : ARMREG_R8;
 				ainfo->nregs = 1;
 				ainfo->size = size;
 				continue;

src/mono/mono/mini/tramp-amd64.c

@@ -1162,6 +1162,8 @@ mono_arch_get_interp_to_native_trampoline (MonoTrampInfo **info)
 		amd64_mov_membase_reg (code, AMD64_RBP, off_ctxregs - i * sizeof (target_mgreg_t), i + CTX_REGS_OFFSET, sizeof (target_mgreg_t));
 		amd64_mov_reg_membase (code, i + CTX_REGS_OFFSET, AMD64_R11, MONO_STRUCT_OFFSET (CallContext, gregs) + (i + CTX_REGS_OFFSET) * sizeof (target_mgreg_t), sizeof (target_mgreg_t));		
 	}
+	// Move R10 to RAX before the native call as R10 is used as a proxy for SwiftIndirectResult.
+	amd64_mov_reg_membase (code, AMD64_RAX, AMD64_R11, MONO_STRUCT_OFFSET (CallContext, gregs) + 10 * sizeof (target_mgreg_t), sizeof (target_mgreg_t));
 #endif
 
 	/* load target addr */
@@ -1267,6 +1269,9 @@ mono_arch_get_native_to_interp_trampoline (MonoTrampInfo **info)
 	/* set context registers to CallContext */
 	for (i = 0; i < CTX_REGS; i++)
 		amd64_mov_membase_reg (code, AMD64_RSP, ctx_offset + MONO_STRUCT_OFFSET (CallContext, gregs) + (i + CTX_REGS_OFFSET) * sizeof (target_mgreg_t), i + CTX_REGS_OFFSET, sizeof (target_mgreg_t));
+
+	// Move RAX to R10 before the managed call as R10 is used as a proxy for SwiftIndirectResult.
+	amd64_mov_membase_reg (code, AMD64_RSP, ctx_offset + MONO_STRUCT_OFFSET (CallContext, gregs) + 10 * sizeof (target_mgreg_t), AMD64_RAX, sizeof (target_mgreg_t));
 #endif
 
 	/* set the stack pointer to the value at call site */

src/tests/issues.targets

@@ -1917,9 +1917,6 @@
         <ExcludeList Include="$(XunitTestBinBase)/Interop/Swift/SwiftCallbackAbiStress/**">
             <Issue>https://github.com/dotnet/runtime/issues/93631: Swift reverse pinvokes are not implemented on Mono yet</Issue>
         </ExcludeList>
-        <ExcludeList Include="$(XunitTestBinBase)/Interop/Swift/SwiftIndirectResult/**">
-            <Issue>https://github.com/dotnet/runtime/issues/93631: SwiftIndirectResult is not implemented on Mono yet</Issue>
-        </ExcludeList>
     </ItemGroup>
 
     <!-- Known failures for mono runtime on Windows -->