Implement computing struct layout (size, alignment, field offsets) in type checking

examples/jit_test.psl

@@ -15,6 +15,8 @@ foreign func putchar(c: i32) -> i32
 func print(c: u8):
 	putchar(c as i32)
 
+let x = 0
+
 func test1():
 	print('H')
 	print('e')
@@ -42,9 +44,20 @@ func test() -> i32:
 	else:
 		return b()
 
-print('O')
-print('K')
-print('\n')
+struct interval:
+	min: f32
+	max: f32
+
+struct box2f:
+	x: interval
+	y: interval
+
+struct weird:
+	a: i32
+	b: u8
+	c: f32 -> f32
+	d: f64
+	e: f16
 
 //func test1():
 //	let str = ['H', 'e', 'l', 'l', 'o', ',', ' ', 'w', 'o', 'r', 'l', 'd', '!', '\n']

libs/ast/include/pslang/ast/struct.hpp

@@ -11,11 +11,24 @@
 namespace pslang::ast
 {
 
+	struct field_layout
+	{
+		std::size_t offset = 0;
+	};
+
 	struct field_definition
 	{
 		std::string name;
 		ast::type_ptr type;
 		ast::location location;
+
+		field_layout layout = {};
+	};
+
+	struct struct_layout
+	{
+		std::size_t size = 0;
+		std::size_t alignment = 1;
 	};
 
 	struct struct_definition
@@ -24,6 +37,8 @@ namespace pslang::ast
 		std::vector<field_definition> fields;
 		ast::location prelude_location;
 		ast::location location;
+
+		struct_layout layout = {};
 	};
 
 	struct field_access

libs/ast/source/print.cpp

@@ -382,13 +382,13 @@ namespace pslang::ast
 				put_indent(out, options);
 				out << "field { name = \"" << node.name << "\", type = ";
 				print(out, *node.type);
-				out << " }\n";
+				out << ", offset = " << node.layout.offset << " }\n";
 			}
 
 			void apply(struct_definition const & node)
 			{
 				put_indent(out, options);
-				out << "struct { name = \"" << node.name << "\" }\n";
+				out << "struct { name = \"" << node.name << "\", size = " << node.layout.size << ", align = " << node.layout.alignment << " }\n";
 				for (auto const & field : node.fields)
 					child(field);
 			}

libs/ast/source/type_check.cpp

@@ -5,6 +5,7 @@
 #include <pslang/ast/error.hpp>
 #include <pslang/types/print.hpp>
 #include <pslang/types/type.hpp>
+#include <pslang/types/type_visitor.hpp>
 
 #include <unordered_map>
 #include <sstream>
@@ -35,6 +36,10 @@ namespace pslang::ast
 				types::type_ptr type;
 			};
 
+			ast::struct_definition * node;
+			bool layout_being_computed = false;
+			bool layout_ready = false;
+
 			std::vector<field_data> fields;
 		};
 
@@ -45,11 +50,85 @@ namespace pslang::ast
 			std::unordered_map<std::string, struct_data> structs;
 
 			bool is_function_scope = false;
-			bool is_global_scope = false;
 
 			types::type_ptr expected_return_type = nullptr;
 		};
 
+		void compute_layout(struct_data & data, std::vector<scope> & scopes);
+
+		struct size_and_alignment
+		{
+			std::size_t size;
+			std::size_t alignment;
+		};
+
+		struct field_layout_visitor
+			: types::const_visitor<field_layout_visitor>
+		{
+			std::vector<scope> & scopes;
+
+			using const_visitor::apply;
+
+			size_and_alignment apply(types::unit_type const &)
+			{
+				return {.size = 1, .alignment = 1};
+			}
+
+			size_and_alignment apply(types::primitive_type const & type)
+			{
+				auto size = types::builtin_type_size(type);
+				return {.size = size, .alignment = size};
+			}
+
+			size_and_alignment apply(types::array_type const & type)
+			{
+				auto base = apply(*type.element_type);
+				return {.size = base.size * type.size, .alignment = base.alignment};
+			}
+
+			size_and_alignment apply(types::function_type const & type)
+			{
+				return {.size = 8, .alignment = 8};
+			}
+
+			size_and_alignment apply(types::named_type const & type)
+			{
+				for (auto it = scopes.rbegin(); it != scopes.rend(); ++it)
+					if (auto jt = it->structs.find(type.name); jt != it->structs.end())
+					{
+						// TODO: better error message (including the resursive inclusion path)
+						if (!jt->second.layout_ready && jt->second.layout_being_computed)
+							throw validation_error("Recursive structs are not allowed", jt->second.node->location);
+
+						compute_layout(jt->second, scopes);
+
+						auto & layout = jt->second.node->layout;
+						return {.size = layout.size, .alignment = layout.alignment};
+					}
+
+				throw std::runtime_error("Unknown type \"" + type.name + "\"");
+			}
+		};
+
+		void compute_layout(struct_data & data, std::vector<scope> & scopes)
+		{
+			if (data.layout_ready)
+				return;
+
+			data.layout_being_computed = true;
+			auto & layout = data.node->layout;
+			for (std::size_t i = 0; i < data.fields.size(); ++i)
+			{
+				auto field_layout = field_layout_visitor{{}, scopes}.apply(*data.fields[i].type);
+				layout.alignment = std::max(layout.alignment, field_layout.alignment);
+				layout.size = ((layout.size + field_layout.alignment - 1) / field_layout.alignment) * field_layout.alignment;
+				data.node->fields[i].layout.offset = layout.size;
+				layout.size += field_layout.size;
+			}
+			layout.size = ((layout.size + layout.alignment - 1) / layout.alignment) * layout.alignment;
+			data.layout_ready = true;
+		}
+
 		struct resolve_types_visitor
 			: type_visitor<resolve_types_visitor>
 		{
@@ -179,12 +258,13 @@ namespace pslang::ast
 				resolve_types(scopes, *node.type);
 			}
 
-			void apply(struct_definition const & node)
+			void apply(struct_definition & node)
 			{
 				for (auto const & field : node.fields)
 					apply(field);
 
 				auto & data = scopes.back().structs[node.name];
+				data.node = &node;
 				for (auto const & field : node.fields)
 					data.fields.push_back({.name = field.name, .type = get_type(*field.type)});
 			}
@@ -787,6 +867,9 @@ namespace pslang::ast
 				populate_globals(scopes, node);
 				for (auto const & statement : node.statements)
 					apply(*statement);
+
+				for (auto & struct_data : scopes.back().structs)
+					compute_layout(struct_data.second, scopes);
 			}
 		};
 
@@ -795,7 +878,7 @@ namespace pslang::ast
 	void check_and_infer_types(statement_list_ptr & statements)
 	{
 		std::vector<scope> scopes;
-		scopes.emplace_back().is_global_scope = true;
+		scopes.emplace_back();
 		check_visitor visitor{{}, {}, scopes};
 		visitor.apply(*statements);
 	}