Aarch64 compiler wip: variables

apps/interpreter/source/main.cpp

@@ -170,7 +170,7 @@ int main(int argc, char ** argv)
 		{
 			// TODO: remove, testing-only code; should execute entry point instead
 			auto offset = module.code.symbol_table.at("test");
-			auto fptr = (bool(*)())(module.code.memory.data.get() + offset);
+			auto fptr = (int(*)())(module.code.memory.data.get() + offset);
 			auto x = fptr();
 			std::cout << "Result: " << std::boolalpha << x << std::endl;
 		}

examples/jit_test.psl

@@ -1,2 +1,5 @@
 func test() -> i32:
-	return -10*9
+	let x = 42 + 69
+	mut y = (x * 2) / 3
+	y = y + 13
+	return x - y

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

@@ -8,6 +8,7 @@
 #include <stdexcept>
 #include <type_traits>
 #include <unordered_map>
+#include <vector>
 
 namespace pslang::jit::aarch64
 {
@@ -73,7 +74,36 @@ namespace pslang::jit::aarch64
 
 			context & context;
 			instruction_builder builder{context.code};
-			std::uint32_t stack_free_bytes = 0; // must be a multiple of 8
+
+			// Bytes starting at stack pointer that are free for use
+			// (pushing registers / allocating variables)
+			// Must be a multiple of 8
+			std::uint32_t stack_free_bytes = 0;
+
+			// Different between initial stack pointer at function enter
+			// and current stack pointer value
+			// Must be a multiple of 16
+			std::uint32_t stack_offset = 0;
+
+			struct variable_data
+			{
+				// Difference between initial stack pointer at scope enter
+				// and the variable address
+				// Must be a multiple of 8
+				std::uint32_t frame_offset;
+			};
+
+			struct scope
+			{
+				std::unordered_map<std::string, variable_data> variables = {};
+
+				// Different between initial stack pointer at scope enter
+				// and current stack pointer value
+				// Must be a multiple of 16
+				std::uint32_t stack_offset = 0;
+			};
+
+			std::vector<scope> scopes;
 			
 			template <typename Node>
 			void apply(Node const &)
@@ -108,10 +138,23 @@ namespace pslang::jit::aarch64
 
 				if (sizeof(T) < 8)
 				{
-					if (std::is_signed_v<T>)
+					if constexpr (std::is_signed_v<T>)
+					{
+						if (node.value < 0)
 							builder.sbfm(0, 0, sizeof(T) * 8);
-					else
+					}
-						builder.ubfm(0, 0, sizeof(T) * 8);
+				}
+			}
+
+			void apply(ast::identifier const & node)
+			{
+				for (auto it = scopes.rbegin(); it != scopes.rend(); ++it)
+				{
+					if (auto jt = it->variables.find(node.name); jt != it->variables.end())
+					{
+						builder.ldr(0, 31, (stack_offset - jt->second.frame_offset) / 8);
+						break;
+					}
 				}
 			}
 
@@ -267,9 +310,36 @@ namespace pslang::jit::aarch64
 				}
 			}
 
+			void apply(ast::assignment const & node)
+			{
+				auto identifier = std::get_if<ast::identifier>(node.lhs.get());
+				if (!identifier)
+					throw std::runtime_error("Not implemented");
+
+				apply(*node.rhs);
+
+				for (auto it = scopes.rbegin(); it != scopes.rend(); ++it)
+				{
+					if (auto jt = it->variables.find(identifier->name); jt != it->variables.end())
+					{
+						builder.str(0, 31, (stack_offset - jt->second.frame_offset) / 8);
+						break;
+					}
+				}
+			}
+
+			void apply(ast::variable_declaration const & node)
+			{
+				apply(*node.initializer);
+				push(0);
+				scopes.back().variables[node.name] = {.frame_offset = stack_offset - stack_free_bytes};
+			}
+
 			void apply(ast::return_statement const & node)
 			{
 				apply(*node.value);
+				if (stack_offset > 0)
+					builder.add_imm(31, 31, stack_offset);
 				builder.ret();
 			}
 
@@ -280,6 +350,7 @@ namespace pslang::jit::aarch64
 
 			void do_apply(ast::function_definition const & node)
 			{
+				scopes.emplace_back();
 				// TODO: arguments
 				apply(*node.statements);
 			}
@@ -291,6 +362,8 @@ namespace pslang::jit::aarch64
 				{
 					builder.sub_imm(31, 31, 16);
 					stack_free_bytes += 16;
+					stack_offset += 16;
+					scopes.back().stack_offset += 16;
 				}
 				builder.str(reg, 31, (stack_free_bytes - 8) / 8);
 				stack_free_bytes -= 8;
@@ -304,6 +377,8 @@ namespace pslang::jit::aarch64
 				{
 					builder.add_imm(31, 31, 16);
 					stack_free_bytes -= 16;
+					stack_offset -= 16;
+					scopes.back().stack_offset -= 16;
 				}
 			}