diff --git a/libs/jit/source/arch/aarch64/compiler.cpp b/libs/jit/source/arch/aarch64/compiler.cpp
index 13ca073..2593ae8 100644
--- a/libs/jit/source/arch/aarch64/compiler.cpp
+++ b/libs/jit/source/arch/aarch64/compiler.cpp
@@ -793,11 +793,29 @@ namespace pslang::jit::aarch64
 					else
 						throw std::runtime_error("Unsupported function argument type");
 				}
-				if (return_value_is_large_struct)
+
+				std::int32_t size = ast::type_size(*type);
+				auto struct_type = std::get_if<types::struct_type>(type.get());
+				auto array_type = std::get_if<types::array_type>(type.get());
+				bool return_value_is_large_struct = false;
+				std::int32_t base_offset = stack_size - stack_position.at(it);
+				if (struct_type || array_type)
+				{
+					auto hfa = get_hfa_data(lcontext, type);
+					if (!(hfa && hfa->count <= 4) && !(size <= 16))
+						return_value_is_large_struct = true;
+				}
+
+				// Save x8 in case the called function changes it
+				if (this->return_value_is_large_struct)
 				{
 					builder.sub_imm(31, 31, 16);
 					builder.str(8, 31, 0);
 				}
+				if (return_value_is_large_struct)
+				{
+					builder.add_imm(31, 8, base_offset);
+				}
 				if (!lcontext.use_frame_pointer)
 				{
 					builder.sub_imm(31, 31, 16);
@@ -809,23 +827,18 @@ namespace pslang::jit::aarch64
 					builder.ldr(30, 31, 0);
 					builder.add_imm(31, 31, 16);
 				}
-				if (return_value_is_large_struct)
+				if (this->return_value_is_large_struct)
 				{
 					builder.ldr(8, 31, 0);
 					builder.add_imm(31, 31, 16);
 				}
 
-				// TODO: array return value?
-				auto size = ast::type_size(*type);
-				auto struct_type = std::get_if<types::struct_type>(type.get());
-				auto array_type = std::get_if<types::array_type>(type.get());
 				if (size == 0)
 				{}
 				else if (struct_type || array_type)
 				{
 					// NB: fixed-size arrays are handled in the same way
 					// as structs of N identical fields
-					auto base_offset = stack_size - stack_position.at(it);
 					if (auto hfa = get_hfa_data(lcontext, type); hfa && hfa->count <= 4)
 					{
 						auto fp_mode = fp_mode_for(*hfa->element_type);
@@ -843,8 +856,7 @@ namespace pslang::jit::aarch64
 					}
 					else
 					{
-						// Large struct - returned by pointer in x8 register
-						copy_memory(8, 0, 31, base_offset, size, 0);
+						// Nothing to be done - return value should already be in place
 					}
 				}
 				else if (types::is_integer_like_type(*type))
@@ -953,11 +965,13 @@ namespace pslang::jit::aarch64
 					else if (ir::is_value_instruction(it->instruction))
 					{
 						auto size = ast::type_size(*it->inferred_type);
-						if (size == 0) continue;
-						// TODO: inefficient for small types, maybe only round up to type alignment?
-						// Need to make sure all read/write arm64 instructions used can handle offsets that
-						// are not a multiple of 8
-						stack_size += ((size + 7) / 8) * 8;
+						if (size > 0)
+						{
+							// TODO: inefficient for small types, maybe only round up to type alignment?
+							// Need to make sure all read/write arm64 instructions used can handle offsets that
+							// are not a multiple of 8
+							stack_size += ((size + 7) / 8) * 8;
+						}
 						stack_position[it] = stack_size;
 					}
 				}