x86_64 backend: tons of fixes + raytracer example works!

libs/ast/include/pslang/ast/type.hpp

@@ -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);
+
 }

libs/ast/source/type.cpp

@@ -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;
+	}
+
 }

libs/jit/include/pslang/jit/arch/linux_x86_64/instruction_builder.hpp

@@ -72,6 +72,9 @@ namespace pslang::jit::linux_x86_64
 		// Store the zero-extended 16-bit GPR @reg_src in 64-bit register @reg_dst
 		void movzx_16(reg reg_src, reg reg_dst);
 
+		// Store the zero-extended 32-bit GPR @reg_src in 64-bit register @reg_dst
+		void movzx_32(reg reg_src, reg reg_dst);
+
 		// Store the sign-extended 8-bit GPR @reg_src in 64-bit register @reg_dst
 		void movsx_8(reg reg_src, reg reg_dst);
 
@@ -198,6 +201,9 @@ namespace pslang::jit::linux_x86_64
 		// Bitwise and the 64-bit GPR @reg_src with the 64-bit GPR @reg_dst and store the result in @reg_dst
 		void and_(reg reg_src, reg reg_dst);
 
+		// Bitwise and the 64-bit GPR @reg_src with the 8-bit @value (sign-extended to 64 bits) and store the result in @reg_dst
+		void and_imm8(reg reg_src, std::int8_t value);
+
 		// Bitwise or the 64-bit GPR @reg_src with the 64-bit GPR @reg_dst and store the result in @reg_dst
 		void or_(reg reg_src, reg reg_dst);
 
@@ -210,6 +216,9 @@ namespace pslang::jit::linux_x86_64
 		// Right-shift with zero-filling the 64-bit GPR @reg_dst by the value of the low 6 bits of RCX and store the result in @reg_dst
 		void shr(reg reg_dst);
 
+		// Right-shift with zero-filling the 64-bit GPR @reg_dst by 1 bit and store the result in @reg_dst
+		void shr1(reg reg_dst);
+
 		// Right-shift with sign-extending the 64-bit GPR @reg_dst by the value of the low 6 bits of RCX and store the result in @reg_dst
 		void sar(reg reg_dst);
 
@@ -309,6 +318,11 @@ namespace pslang::jit::linux_x86_64
 		// NB: RIP holds the address of the _next_ instruction
 		void jump_if_nonzero(std::int32_t offset);
 
+		// Conditional jump to 32-bit signed @offset relative to RIP
+		// if SF flag is set
+		// NB: RIP holds the address of the _next_ instruction
+		void jump_if_sign(std::int32_t offset);
+
 		// Assuming that @opcode refers to the location of a JUMP instruction,
 		// replace its 32-bit jump offset with @offset
 		void jump_inject(std::uint8_t * opcode, std::int32_t offset);
@@ -336,6 +350,22 @@ namespace pslang::jit::linux_x86_64
 		// of the LEA instruction itself
 		void lea_rip_inject_prev(std::uint8_t * opcode, std::int32_t offset);
 
+		// Call to 32-bit signed @offset relative to RIP
+		// NB: RIP holds the address of the _next_ instruction
+		void call_imm(std::int32_t offset);
+
+		// Assuming that @opcode refers to the location of a CALL instruction,
+		// replace its 32-bit call offset with @offset
+		void call_imm_inject(std::uint8_t * opcode, std::int32_t offset);
+
+		// Assuming that @opcode refers to the location of a CALL instruction,
+		// replace its 32-bit call offset with @offset, compensating for the size
+		// of the CALL instruction itself
+		void call_imm_inject_prev(std::uint8_t * opcode, std::int32_t offset);
+
+		// Call to address specified in @reg_addr
+		void call_reg(reg reg_addr);
+
 	private:
 		template <typename ... Args>
 		void do_push(Args ... values);

libs/jit/source/arch/linux_x86_64/compiler.cpp

@@ -15,9 +15,9 @@ namespace pslang::jit::linux_x86_64
 	{
 
 		template <typename ... Regs>
-		reg find_free_reg(Regs ... regs)
+		reg find_free_reg(reg start, Regs ... regs)
 		{
-			for (int i = 0; i < 16; ++i)
+			for (int i = (int)start; i < 16; ++i)
 			{
 				if ((true && ... &&  (regs != reg(i))))
 					return reg(i);
@@ -34,6 +34,22 @@ namespace pslang::jit::linux_x86_64
 			reg::r9,
 		};
 
+		static std::array<reg, 2> integer_ret_reg = {
+			reg::rax,
+			reg::rdx,
+		};
+
+		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 +64,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
@@ -59,8 +77,43 @@ namespace pslang::jit::linux_x86_64
 			std::vector<resolve_data> jump_resolve;
 			std::vector<resolve_data> cjump_resolve;
 			std::vector<resolve_data> node_resolve;
+			std::vector<resolve_data> call_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;
@@ -96,7 +149,7 @@ namespace pslang::jit::linux_x86_64
 
 			void apply(ir::extern_symbol const & node, types::type_ptr const &)
 			{
-				std::int32_t offset = pcontext.storage.size();
+ 				std::int32_t offset = pcontext.storage.size();
 				lcontext.extern_symbols[node.name] = offset;
 				pcontext.foreign_resolve.push_back({node.name, offset});
 				push_bytes(pcontext.storage.storage, (void *)nullptr);
@@ -160,6 +213,7 @@ namespace pslang::jit::linux_x86_64
 			std::vector<std::int32_t> argument_position;
 			std::unordered_map<ir::node_ref, std::int32_t> stack_position;
 			std::int32_t stack_size = 0;
+			bool return_value_is_large_struct = false;
 
 			void apply(ir::node_ref, ir::label const &, types::type_ptr const &)
 			{}
@@ -283,6 +337,7 @@ namespace pslang::jit::linux_x86_64
 			void apply(ir::node_ref it, ir::binary_operation const & node, types::type_ptr const & type)
 			{
 				auto const arg1_type = node.arg1->inferred_type;
+				auto const arg2_type = node.arg2->inferred_type;
 				bool const is_fp = types::is_floating_point_type(*arg1_type);
 				bool const result_is_fp = types::is_floating_point_type(*type);
 				auto const arg_size = ast::type_size(*arg1_type);
@@ -392,7 +447,7 @@ namespace pslang::jit::linux_x86_64
 					break;
 				case ast::binary_operation_type::right_shift:
 					reg_extend(reg::rax, reg::rax, *arg1_type);
-					reg_extend(reg::rbx, reg::rcx, *arg1_type);
+					reg_extend(reg::rbx, reg::rcx, *arg2_type);
 					if (types::is_unsigned_integer_type(*arg1_type))
 						builder.shr(reg::rax);
 					else
@@ -573,24 +628,66 @@ namespace pslang::jit::linux_x86_64
 						}
 						else if (types::is_floating_point_type(*dst_type))
 						{
-							if (types::is_unsigned_integer_type(*src_type) && ast::type_size(*src_type) == 8)
-								throw std::runtime_error("64-bit unsigned to floating-point conversion is not implemented");
-
 							reg_extend(reg::rax, reg::rax, *src_type);
 
 							auto dst_size = ast::type_size(*dst_type);
-							if (dst_size <= 4)
-								builder.gpr_to_xmm_32(reg::rax, reg::xmm0);
+
+							if (types::is_unsigned_integer_type(*src_type) && ast::type_size(*src_type) == 8)
+							{
+								// No u64 -> f32/f64 conversion on x86_64
+								// Emulate it by splitting x = 2 * (x >> 1) + (x & 1)
+
+								reg_extend(reg::rax, reg::rax, *src_type);
+
+								builder.test(reg::rax, reg::rax);
+								auto first_jump = builder.code.size();
+								builder.jump_if_sign(0);
+
+								if (dst_size <= 4)
+									builder.gpr_to_xmm_32(reg::rax, reg::xmm0);
+								else
+									builder.gpr_to_xmm(reg::rax, reg::xmm0);
+
+								auto second_jump = builder.code.size();
+								builder.jump(0);
+
+								auto first_jump_target = builder.code.size();
+
+								builder.mov(reg::rax, reg::rcx);
+								builder.shr1(reg::rcx);
+								builder.and_imm8(reg::rax, 1);
+								builder.or_(reg::rax, reg::rcx);
+
+								if (dst_size <= 4)
+								{
+									builder.gpr_to_xmm_32(reg::rcx, reg::xmm0);
+									builder.add_xmm_32(reg::xmm0, reg::xmm0);
+								}
+								else
+								{
+									builder.gpr_to_xmm(reg::rcx, reg::xmm0);
+									builder.add_xmm(reg::xmm0, reg::xmm0);
+								}
+
+								auto second_jump_target = builder.code.size();
+
+								builder.cjump_inject_prev(builder.code.data() + first_jump, first_jump_target - first_jump);
+								builder.jump_inject_prev(builder.code.data() + second_jump, second_jump_target - second_jump);
+							}
 							else
-								builder.gpr_to_xmm(reg::rax, reg::xmm0);
+							{
+								if (dst_size <= 4)
+									builder.gpr_to_xmm_32(reg::rax, reg::xmm0);
+								else
+									builder.gpr_to_xmm(reg::rax, reg::xmm0);
+							}
 						}
 					}
 					else if (types::is_floating_point_type(*src_type))
 					{
 						if (types::is_integer_type(*dst_type))
 						{
-							if (types::is_unsigned_integer_type(*dst_type) && ast::type_size(*dst_type) == 8)
-								throw std::runtime_error("Floating-point to 64-bit unsigned conversion is not implemented");
+							// TODO: correct floating-point to 64-bit unsigned conversion
 
 							auto src_size = ast::type_size(*src_type);
 							if (src_size <= 4)
@@ -639,7 +736,7 @@ namespace pslang::jit::linux_x86_64
 
 			void apply(ir::node_ref it, ir::extern_symbol const & node, types::type_ptr const & type)
 			{
-				builder.lea_rip(lcontext.extern_symbols[node.name] - (std::int32_t)pcontext.storage.size(), reg::rax);
+				builder.mov_read_rip_prev(lcontext.extern_symbols[node.name] - (std::int32_t)pcontext.storage.size(), reg::rax);
 				store(it, reg::rax);
 			}
 
@@ -678,6 +775,7 @@ namespace pslang::jit::linux_x86_64
 			void apply(ir::node_ref, ir::jump_if_zero const & node, types::type_ptr const & type)
 			{
 				load(node.condition, reg::rax);
+				reg_extend(reg::rax, reg::rax, *node.condition->inferred_type);
 				builder.test(reg::rax, reg::rax);
 				lcontext.cjump_resolve.push_back({(std::int32_t)pcontext.storage.size(), node.target});
 				builder.jump_if_zero(0);
@@ -686,19 +784,165 @@ namespace pslang::jit::linux_x86_64
 			void apply(ir::node_ref, ir::jump_if_nonzero const & node, types::type_ptr const & type)
 			{
 				load(node.condition, reg::rax);
+				reg_extend(reg::rax, reg::rax, *node.condition->inferred_type);
 				builder.test(reg::rax, reg::rax);
 				lcontext.cjump_resolve.push_back({(std::int32_t)pcontext.storage.size(), node.target});
 				builder.jump_if_nonzero(0);
 			}
 
-			void apply(ir::node_ref, ir::call const & node, types::type_ptr const & type)
+			template <typename Node, typename DoCall>
+			void apply_call(ir::node_ref it, Node const & node, types::type_ptr const & type, DoCall && do_call)
 			{
-				throw std::runtime_error("Not implemented");
+				bool return_value_is_large_struct = false;
+				if (std::holds_alternative<types::struct_type>(*type) || std::holds_alternative<types::array_type>(*type))
+					return_value_is_large_struct = !classify_small_struct(lcontext, type);
+
+				if (this->return_value_is_large_struct)
+				{
+					builder.push(reg::rdi);
+					stack_size += 8;
+				}
+
+				// Compute space required for stack arguments
+				std::int32_t arguments_stack_size = 0;
+				for (auto const & argument : node.arguments)
+				{
+					auto size = ast::type_size(*argument->inferred_type);
+					auto alignment = ast::type_size(*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 (struct_type || array_type)
+					{
+						if (!classify_small_struct(lcontext, argument->inferred_type))
+						{
+							arguments_stack_size += (alignment - (arguments_stack_size % alignment)) % alignment;
+							arguments_stack_size += size;
+						}
+					}
+				}
+
+				// Ensure stack is 16-byte aligned after CALL pushes return address on stack
+				// I.e. enforce (stack_size + arguments_stack_size) % 16 == 8
+				arguments_stack_size += (24 - ((arguments_stack_size + stack_size) % 16)) % 16;
+
+				// TODO: handle the case when there weren't enough registers
+				std::uint8_t reg_index = return_value_is_large_struct ? 1 : 0;
+				std::uint8_t fp_reg = 0;
+				std::int32_t stack_offset = 0;
+				for (auto const & argument : node.arguments)
+				{
+					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 (struct_type || array_type)
+					{
+						if (auto small_struct = classify_small_struct(lcontext, argument->inferred_type))
+						{
+							auto address = node_address(argument);
+
+							for (int o : {0, 1})
+							{
+								auto & octet = small_struct->octets[o];
+								if (!octet) break;
+
+								if (octet->has_integer)
+									load(argument, integer_arg_reg[reg_index++], 8 * o);
+								else if (octet->has_floating_point)
+									load_xmm(argument, (reg)(fp_reg++), 8, 8 * o);
+							}
+						}
+						else
+						{
+							stack_offset += (alignment - (stack_offset % alignment)) % alignment;
+							copy_memory(node_address(argument), {.base = reg::rsp, .offset = stack_offset - arguments_stack_size}, size);
+							stack_offset += size;
+						}
+					}
+					else if (types::is_integer_like_type(*argument->inferred_type))
+						load(argument, integer_arg_reg[reg_index++]);
+					else if (types::is_floating_point_type(*argument->inferred_type))
+						load_xmm(argument, (reg)(fp_reg++), size);
+					else
+						throw std::runtime_error("Unsupported function argument type");
+				}
+
+				if (return_value_is_large_struct)
+				{
+					auto address = node_address(it);
+					// Function call node cannot have RIP-relative address
+					builder.lea(address.base.value(), address.offset, reg::rdi);
+				}
+
+				if (arguments_stack_size > 0)
+				{
+					builder.sub_imm(reg::rsp, arguments_stack_size);
+					stack_size += arguments_stack_size;
+				}
+
+				do_call();
+
+				if (arguments_stack_size > 0)
+				{
+					builder.add_imm(reg::rsp, arguments_stack_size);
+					stack_size -= arguments_stack_size;
+				}
+
+				std::size_t const 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)
+				{
+					if (auto small_struct = classify_small_struct(lcontext, type))
+					{
+						std::uint8_t reg_index = 0;
+						std::uint8_t fp_reg = 0;
+						for (int o : {0, 1})
+						{
+							auto & octet = small_struct->octets[o];
+							if (!octet) break;
+
+							if (octet->has_integer)
+								store(it, integer_ret_reg[reg_index++], 8 * o);
+							else if (octet->has_floating_point)
+								store_xmm(it, (reg)(fp_reg++), 8, 8 * o);
+						}
+					}
+					else
+					{
+						// Nothing to be done - return value should already be in place
+					}
+				}
+				else if (types::is_integer_like_type(*type))
+					store(it, reg::rax);
+				else if (types::is_floating_point_type(*type))
+					store_xmm(it, reg::xmm0, size);
+				else
+					throw std::runtime_error("Unsupported return value type");
+
+				if (this->return_value_is_large_struct)
+				{
+					builder.pop(reg::rdi);
+					stack_size -= 8;
+				}
 			}
 
-			void apply(ir::node_ref, ir::call_pointer const & node, types::type_ptr const & type)
+			void apply(ir::node_ref it, ir::call const & node, types::type_ptr const & type)
 			{
-				throw std::runtime_error("Not implemented");
+				apply_call(it, node, type, [&]{
+					lcontext.call_resolve.emplace_back(pcontext.storage.size(), node.target);
+					builder.call_imm(0);
+				});
+			}
+
+			void apply(ir::node_ref it, ir::call_pointer const & node, types::type_ptr const & type)
+			{
+				apply_call(it, node, type, [&]{
+					load(node.pointer, reg::rax);
+					builder.call_reg(reg::rax);
+				});
 			}
 
 			void apply(ir::node_ref, ir::return_value const & node, types::type_ptr const & type)
@@ -713,7 +957,27 @@ namespace pslang::jit::linux_x86_64
 					{}
 					else if (struct_type || array_type)
 					{
-						throw std::runtime_error("Not implemented");
+						if (auto small_struct = classify_small_struct(lcontext, type))
+						{
+							auto address = node_address(*node.value);
+
+							std::uint8_t reg_index = 0;
+							std::uint8_t fp_reg = 0;
+							for (int o : {0, 1})
+							{
+								auto & octet = small_struct->octets[o];
+								if (!octet) break;
+
+								if (octet->has_integer)
+									load(*node.value, integer_ret_reg[reg_index++], 8 * o);
+								else if (octet->has_floating_point)
+									load_xmm(*node.value, (reg)(fp_reg++), 8, 8 * o);
+							}
+						}
+						else
+						{
+							copy_memory(node_address(*node.value), {reg::rdi, 0}, size);
+						}
 					}
 					else if (types::is_integer_like_type(*type))
 						load(*node.value, reg::rax);
@@ -723,11 +987,11 @@ namespace pslang::jit::linux_x86_64
 						throw std::runtime_error("Unsupported return value type");
 				}
 				if (stack_size > 0)
-				{
 					builder.add_imm(reg::rsp, stack_size);
-					if (lcontext.use_frame_pointer)
-						builder.pop(reg::rbp);
-				}
+
+				builder.pop(reg::rbx);
+				if (lcontext.use_frame_pointer)
+					builder.pop(reg::rbp);
 				builder.ret();
 			}
 
@@ -735,11 +999,9 @@ namespace pslang::jit::linux_x86_64
 			{
 				auto result_type = function_definition->inferred_result_type;
 
-				if (lcontext.use_frame_pointer)
-				{
-					builder.push(reg::rbp);
-					builder.mov(reg::rsp, reg::rbp);
-				}
+				if (result_type)
+					if (std::holds_alternative<types::struct_type>(*result_type) || std::holds_alternative<types::array_type>(*result_type))
+						return_value_is_large_struct = !classify_small_struct(lcontext, result_type);
 
 				stack_size = 0;
 
@@ -776,23 +1038,55 @@ namespace pslang::jit::linux_x86_64
 				auto it = begin;
 
 				lcontext.nodes[it] = pcontext.storage.size();
+
+				if (lcontext.use_frame_pointer)
+				{
+					builder.push(reg::rbp);
+					builder.mov(reg::rsp, reg::rbp);
+				}
+
+				builder.push(reg::rbx);
+
 				if (stack_size > 0)
 					builder.sub_imm(reg::rsp, stack_size);
 
 				// TODO: handle the case when there weren't enough registers
 
-				std::uint8_t reg_index = 0;
+				// If return value is large struct, RDI stores the pointer to return value
+				std::uint8_t reg_index = return_value_is_large_struct ? 1 : 0;
 				std::uint8_t fp_reg = 0;
+				// 8 for the return address
+				// Another 8 for RBP if using frame pointers
+				// Another 8 for saved RBX
+				std::int32_t stack_offset = stack_size + (lcontext.use_frame_pointer ? 16 : 8) + 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) break;
+								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))
 					{
@@ -820,7 +1114,7 @@ namespace pslang::jit::linux_x86_64
 						continue;
 
 					// Uncomment to debug per-node instruction generation:
-					// builder.nop();
+					builder.nop();
 
 					lcontext.nodes[it] = pcontext.storage.size();
 					std::visit([&](auto const & instruction){ apply(it, instruction, it->inferred_type); }, it->instruction);
@@ -838,6 +1132,10 @@ namespace pslang::jit::linux_x86_64
 						builder.movzx_8(reg_src, reg_dst);
 					else if (size == 2)
 						builder.movzx_16(reg_src, reg_dst);
+					else if (size == 4)
+						builder.movzx_32(reg_src, reg_dst);
+					else if (reg_src != reg_dst)
+						builder.mov(reg_src, reg_dst);
 				}
 				else if (types::is_signed_integer_type(type))
 				{
@@ -847,6 +1145,8 @@ namespace pslang::jit::linux_x86_64
 						builder.movsx_16(reg_src, reg_dst);
 					else if (size == 4)
 						builder.movsx_32(reg_src, reg_dst);
+					else if (reg_src != reg_dst)
+						builder.mov(reg_src, reg_dst);
 				}
 				else if (types::is_bool_type(type))
 				{
@@ -862,25 +1162,25 @@ namespace pslang::jit::linux_x86_64
 					return {.base = reg::rsp, .offset = stack_size - stack_position.at(it)};
 			}
 
-			void load(ir::node_ref it, reg reg_dst)
+			void load(ir::node_ref it, reg reg_dst, std::int32_t offset = 0)
 			{
 				auto const address = node_address(it);
 				if (address.base)
-					builder.mov_read(*address.base, address.offset, reg_dst);
+					builder.mov_read(*address.base, address.offset + offset, reg_dst);
 				else
-					builder.mov_read_rip_prev(address.offset, reg_dst);
+					builder.mov_read_rip_prev(address.offset + offset, reg_dst);
 			}
 
-			void store(ir::node_ref it, reg reg_src)
+			void store(ir::node_ref it, reg reg_src, std::int32_t offset = 0)
 			{
 				auto const address = node_address(it);
 				if (address.base)
-					builder.mov_write(reg_src, *address.base, address.offset);
+					builder.mov_write(reg_src, *address.base, address.offset + offset);
 				else
-					builder.mov_write_rip_prev(reg_src, address.offset);
+					builder.mov_write_rip_prev(reg_src, address.offset + offset);
 			}
 
-			void load_xmm(ir::node_ref it, reg reg_dst, std::uint8_t size)
+			void load_xmm(ir::node_ref it, reg reg_dst, std::uint8_t size, std::int32_t offset = 0)
 			{
 				if (size != 4 && size != 8)
 					throw std::runtime_error("Bad type size for load_xmm");
@@ -889,15 +1189,15 @@ namespace pslang::jit::linux_x86_64
 				if (address.base)
 				{
 					if (size == 4)
-						builder.mov_read_xmm_32(*address.base, address.offset, reg_dst);
+						builder.mov_read_xmm_32(*address.base, address.offset + offset, reg_dst);
 					else
-						builder.mov_read_xmm(*address.base, address.offset, reg_dst);
+						builder.mov_read_xmm(*address.base, address.offset + offset, reg_dst);
 				}
 				else
 					throw std::runtime_error("RIP-relative XMM read is not supported");
 			}
 
-			void store_xmm(ir::node_ref it, reg reg_src, std::uint8_t size)
+			void store_xmm(ir::node_ref it, reg reg_src, std::uint8_t size, std::int32_t offset = 0)
 			{
 				if (size != 4 && size != 8)
 					throw std::runtime_error("Bad type size for store_xmm");
@@ -906,9 +1206,9 @@ namespace pslang::jit::linux_x86_64
 				if (address.base)
 				{
 					if (size == 4)
-						builder.mov_write_xmm_32(reg_src, *address.base, address.offset);
+						builder.mov_write_xmm_32(reg_src, *address.base, address.offset + offset);
 					else
-						builder.mov_write_xmm(reg_src, *address.base, address.offset);
+						builder.mov_write_xmm(reg_src, *address.base, address.offset + offset);
 				}
 				else
 					throw std::runtime_error("RIP-relative XMM write is not supported");
@@ -923,7 +1223,7 @@ namespace pslang::jit::linux_x86_64
 					reg_src = *src.base;
 				else
 				{
-					reg_src = find_free_reg(dst.base);
+					reg_src = find_free_reg(reg::r10, dst.base);
 					builder.lea_rip_prev(0, reg_src);
 				}
 
@@ -932,13 +1232,13 @@ namespace pslang::jit::linux_x86_64
 					reg_dst = *dst.base;
 				else
 				{
-					reg_dst = find_free_reg(reg_src);
+					reg_dst = find_free_reg(reg::r10, reg_src);
 					builder.lea_rip_prev(storage_size_at_start - (std::int32_t)pcontext.storage.size(), reg_dst);
 				}
 
 				std::int32_t offset = 0;
 
-				auto copy_reg = find_free_reg(reg_src, reg_dst);
+				auto copy_reg = find_free_reg(reg::r10, reg_src, reg_dst);
 
 				while (size > 0)
 				{
@@ -1032,6 +1332,9 @@ namespace pslang::jit::linux_x86_64
 		for (auto const & resolve : lcontext.node_resolve)
 			builder.lea_rip_inject_prev(pcontext.storage.storage.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset);
 
+		for (auto const & resolve : lcontext.call_resolve)
+			builder.call_imm_inject_prev(pcontext.storage.storage.data() + resolve.offset, lcontext.nodes.at(resolve.target) - resolve.offset);
+
 		auto code_end = pcontext.storage.align();
 		pcontext.storage.code.push_back({code_begin, code_end});
 	}

libs/jit/source/arch/linux_x86_64/instruction_builder.cpp

@@ -163,6 +163,11 @@ namespace pslang::jit::linux_x86_64
 		do_push(rex(1, reg_dst, {}, reg_src), 0x0f_ub, 0xb7_ub, modrm(MOD_REG, reg_dst, reg_src));
 	}
 
+	void instruction_builder::movzx_32(reg reg_src, reg reg_dst)
+	{
+		do_push(rex(0, reg_src, {}, reg_dst), 0x89_ub, modrm(MOD_REG, reg_src, reg_dst));
+	}
+
 	void instruction_builder::movsx_8(reg reg_src, reg reg_dst)
 	{
 		do_push(rex(1, reg_dst, {}, reg_src), 0x0f_ub, 0xbe_ub, modrm(MOD_REG, reg_dst, reg_src));
@@ -392,6 +397,11 @@ namespace pslang::jit::linux_x86_64
 		do_push(rex(1, reg_dst, {}, reg_src), 0x23_ub, modrm(MOD_REG, reg_dst, reg_src));
 	}
 
+	void instruction_builder::and_imm8(reg reg_src, std::int8_t value)
+	{
+		do_push(rex(1, {}, {}, reg_src), 0x83_ub, modrm(MOD_REG, 0b100, reg_src), value);
+	}
+
 	void instruction_builder::or_(reg reg_src, reg reg_dst)
 	{
 		do_push(rex(1, reg_dst, {}, reg_src), 0x0b_ub, modrm(MOD_REG, reg_dst, reg_src));
@@ -412,6 +422,11 @@ namespace pslang::jit::linux_x86_64
 		do_push(rex(1, {}, {}, reg_dst), 0xd3_ub, modrm(MOD_REG, 0b101, reg_dst));
 	}
 
+	void instruction_builder::shr1(reg reg_dst)
+	{
+		do_push(rex(1, {}, {}, reg_dst), 0xd1_ub, modrm(MOD_REG, 0b101, reg_dst));
+	}
+
 	void instruction_builder::sar(reg reg_dst)
 	{
 		do_push(rex(1, {}, {}, reg_dst), 0xd3_ub, modrm(MOD_REG, 0b111, reg_dst));
@@ -567,6 +582,11 @@ namespace pslang::jit::linux_x86_64
 		do_push(0x0f_ub, 0x85_ub, offset);
 	}
 
+	void instruction_builder::jump_if_sign(std::int32_t offset)
+	{
+		do_push(0x0f_ub, 0x88_ub, offset);
+	}
+
 	void instruction_builder::jump_inject(std::uint8_t * opcode, std::int32_t offset)
 	{
 		auto src = (std::uint8_t const *)(&offset);
@@ -600,6 +620,30 @@ namespace pslang::jit::linux_x86_64
 		lea_rip_inject(opcode, offset - 7);
 	}
 
+	void instruction_builder::call_imm(std::int32_t offset)
+	{
+		do_push(0xe8_ub, offset);
+	}
+
+	void instruction_builder::call_imm_inject(std::uint8_t * opcode, std::int32_t offset)
+	{
+		auto src = (std::uint8_t const *)(&offset);
+		std::copy(src, src + 4, opcode + 1);
+	}
+
+	void instruction_builder::call_imm_inject_prev(std::uint8_t * opcode, std::int32_t offset)
+	{
+		call_imm_inject(opcode, offset - 5);
+	}
+
+	void instruction_builder::call_reg(reg reg_addr)
+	{
+		if ((std::uint8_t(reg_addr) & 0x8) == 0)
+			do_push(0xff_ub, modrm(MOD_REG, 0b010, reg_addr));
+		else
+			do_push(rex(0, {}, {}, reg_addr), 0xff_ub, modrm(MOD_REG, 0b010, reg_addr));
+	}
+
 	template <typename ... Args>
 	void instruction_builder::do_push(Args ... values)
 	{