Fix returning large structs from functions in Aarch64

This commit is contained in:
Nikita Lisitsa 2026-03-30 16:53:02 +03:00
parent 6d69f846f4
commit 52a2a00d1f

View file

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