Support structs and arrays as function arguments on x86_64

This commit is contained in:
Nikita Lisitsa 2026-07-30 18:05:17 +03:00
parent ceb5c8014c
commit 1af72d91de
3 changed files with 149 additions and 1 deletions

View file

@ -62,4 +62,12 @@ namespace pslang::ast
std::size_t type_size(types::type const & type);
std::size_t type_alignment(types::type const & type);
struct flat_field_layout
{
types::type_ptr type;
std::size_t offset;
};
std::vector<flat_field_layout> flat_layout(types::type const & type);
}

View file

@ -118,6 +118,68 @@ namespace pslang::ast
}
};
struct flat_layout_visitor
: types::const_visitor<flat_layout_visitor>
{
using const_visitor::apply;
std::vector<flat_field_layout> & result;
flat_layout_visitor(std::vector<flat_field_layout> & result)
: result(result)
{}
bool apply(types::unit_type const & type)
{
return false;
}
bool apply(types::primitive_type const & type)
{
return false;
}
bool apply(types::array_type const & type)
{
auto const this_offset = current_offset_;
auto const element_size = ast::type_size(*type.element_type);
for (std::uint64_t i = 0; i < type.size; ++i)
{
current_offset_ = this_offset + i * element_size;
if (apply(*type.element_type)) continue;
result.push_back({.type = type.element_type, .offset = current_offset_});
}
return true;
}
bool apply(types::function_type const &)
{
return false;
}
bool apply(types::pointer_type const &)
{
return false;
}
bool apply(types::struct_type const & type)
{
auto const this_offset = current_offset_;
for (auto const & field : type.node->fields)
{
current_offset_ = this_offset + field.layout.offset;
if (apply(*field.inferred_type)) continue;
result.push_back({.type = field.inferred_type, .offset = current_offset_});
}
return true;
}
private:
std::size_t current_offset_ = 0;
};
}
types::type_ptr get_type(type const & type)
@ -135,4 +197,11 @@ namespace pslang::ast
return alignment_visitor{}.apply(type);
}
std::vector<flat_field_layout> flat_layout(types::type const & type)
{
std::vector<flat_field_layout> result;
flat_layout_visitor{result}.apply(type);
return result;
}
}

View file

@ -34,6 +34,17 @@ namespace pslang::jit::linux_x86_64
reg::r9,
};
struct small_struct_data
{
struct octet
{
bool has_integer = false;
bool has_floating_point = false;
};
std::optional<octet> octets[2];
};
struct value_address
{
// None means RIP-based addressing
@ -48,6 +59,8 @@ namespace pslang::jit::linux_x86_64
std::unordered_map<std::string, std::int32_t> extern_symbols;
std::unordered_map<types::type_ptr, small_struct_data> small_structs;
std::unordered_map<ir::node_ref, std::int32_t> nodes;
struct resolve_data
@ -61,6 +74,40 @@ namespace pslang::jit::linux_x86_64
std::vector<resolve_data> node_resolve;
};
std::optional<small_struct_data> classify_small_struct(local_context & lcontext, types::type_ptr type)
{
// TODO: unit type?
if (ast::type_size(*type) > 16)
return std::nullopt;
if (auto it = lcontext.small_structs.find(type); it != lcontext.small_structs.end())
return it->second;
if (std::holds_alternative<types::struct_type>(*type) || std::holds_alternative<types::array_type>(*type))
{
small_struct_data result;
auto layout = ast::flat_layout(*type);
for (auto const & field : layout)
{
std::size_t octet_index = field.offset / 8;
auto & octet = result.octets[octet_index].emplace();
if (types::is_floating_point_type(*field.type))
octet.has_floating_point = true;
else
octet.has_integer = true;
}
return (lcontext.small_structs[type] = result);
}
if (types::is_floating_point_type(*type))
return small_struct_data{.octets = {small_struct_data::octet{.has_floating_point = true}, std::nullopt}};
return small_struct_data{.octets = {small_struct_data::octet{.has_integer = true}, std::nullopt}};
}
struct populate_globals_visitor
{
std::vector<std::uint8_t> & storage;
@ -770,6 +817,9 @@ namespace pslang::jit::linux_x86_64
}
}
// Ensure (RSP % 16) == 8 for future function calls
stack_size = ((stack_size + 7) / 16) * 16 + 8;
if (!std::holds_alternative<ir::label>(begin->instruction))
throw std::runtime_error("First IR node of a function must be a label");
@ -783,16 +833,37 @@ namespace pslang::jit::linux_x86_64
std::uint8_t reg_index = 0;
std::uint8_t fp_reg = 0;
// 8 for the return address
// Another 8 for RBP if using frame pointers
std::int32_t stack_offset = stack_size + (lcontext.use_frame_pointer ? 16 : 8);
for (std::size_t i = 0; i < function_definition->arguments.size(); ++i)
{
auto const & argument = function_definition->arguments[i];
auto size = ast::type_size(*argument.inferred_type);
auto alignment = ast::type_alignment(*argument.inferred_type);
auto struct_type = std::get_if<types::struct_type>(argument.inferred_type.get());
auto array_type = std::get_if<types::array_type>(argument.inferred_type.get());
if (size == 0) continue;
if (struct_type || array_type)
{
throw std::runtime_error("Not implemented");
if (auto small_struct = classify_small_struct(lcontext, argument.inferred_type))
{
for (std::size_t o : {0, 1})
{
auto & octet = small_struct->octets[o];
if (!octet) continue;
if (octet->has_integer)
builder.mov_write(integer_arg_reg[reg_index++], reg::rsp, stack_size - argument_position[i] + 8 * o);
else if (octet->has_floating_point)
builder.mov_write_xmm((reg)(fp_reg++), reg::rsp, stack_size - argument_position[i] + 8 * o);
}
}
else
{
stack_offset += (alignment - (stack_offset % alignment)) % alignment;
copy_memory({.base = reg::rsp, .offset = stack_offset}, {.base = reg::rsp, .offset = stack_size - argument_position[i]}, size);
stack_offset += size;
}
}
else if (types::is_integer_like_type(*argument.inferred_type))
{