From cb32ec345938a489f31f7acd7bb9bb4fa22d73eb Mon Sep 17 00:00:00 2001 From: lisyarus Date: Tue, 6 Jan 2026 19:10:52 +0300 Subject: [PATCH] Refactor jit-compiling interface: introduce common global program context with shared code & symbol tables --- apps/interpreter/source/main.cpp | 18 ++-- .../pslang/jit/arch/aarch64/compiler.hpp | 2 +- .../include/pslang/jit/compiled_module.hpp | 29 ------ libs/jit/include/pslang/jit/executable.hpp | 9 +- libs/jit/include/pslang/jit/jit.hpp | 5 +- .../include/pslang/jit/program_context.hpp | 20 ++++ libs/jit/source/arch/aarch64/compiler.cpp | 93 ++++++++----------- libs/jit/source/executable.cpp | 23 +---- libs/jit/source/jit.cpp | 7 +- 9 files changed, 88 insertions(+), 118 deletions(-) delete mode 100644 libs/jit/include/pslang/jit/compiled_module.hpp create mode 100644 libs/jit/include/pslang/jit/program_context.hpp diff --git a/apps/interpreter/source/main.cpp b/apps/interpreter/source/main.cpp index afad3f0..e138bbd 100644 --- a/apps/interpreter/source/main.cpp +++ b/apps/interpreter/source/main.cpp @@ -161,16 +161,20 @@ int main(int argc, char ** argv) if (jit) { - auto abi = jit::host_abi(); - std::vector modules; - for (auto const & ast : parsed) - modules.push_back(jit::make_host_executable(jit::compile(ast, abi))); + jit::program_context pcontext + { + .abi = jit::host_abi(), + }; + + for (auto const & ast : parsed) + jit::compile(pcontext, ast); + + auto executable = jit::make_host_executable(pcontext.code); - for (auto const & module : modules) { // TODO: remove, testing-only code; should execute entry point instead - auto offset = module.code.symbol_table.at("test"); - auto fptr = (float(*)())(module.code.memory.data.get() + offset); + auto offset = pcontext.symbols.at("test"); + auto fptr = (float(*)())(executable.data.get() + offset); auto x = fptr(); std::cout << "Result: " << std::boolalpha << x << std::endl; } diff --git a/libs/jit/include/pslang/jit/arch/aarch64/compiler.hpp b/libs/jit/include/pslang/jit/arch/aarch64/compiler.hpp index 0c75f35..ce8946e 100644 --- a/libs/jit/include/pslang/jit/arch/aarch64/compiler.hpp +++ b/libs/jit/include/pslang/jit/arch/aarch64/compiler.hpp @@ -5,6 +5,6 @@ namespace pslang::jit::aarch64 { - compiled_module compile(ast::statement_list_ptr const & statements); + void compile(program_context & context, ast::statement_list_ptr const & statements); } \ No newline at end of file diff --git a/libs/jit/include/pslang/jit/compiled_module.hpp b/libs/jit/include/pslang/jit/compiled_module.hpp deleted file mode 100644 index db76ec2..0000000 --- a/libs/jit/include/pslang/jit/compiled_module.hpp +++ /dev/null @@ -1,29 +0,0 @@ -#pragma once - -#include -#include - -#include -#include - -namespace pslang::jit -{ - - struct compiled_module - { - struct segment - { - blob memory; - std::unordered_map symbol_table; - }; - - // TODO: How do we reference data segment from the code segment? - // Maybe use the same relocation/dynamic linking mechanism - // as for imported modules? - segment data; - segment code; - std::size_t entry_point; // in code segment - jit::abi abi; - }; - -} diff --git a/libs/jit/include/pslang/jit/executable.hpp b/libs/jit/include/pslang/jit/executable.hpp index 907a495..7e67d3f 100644 --- a/libs/jit/include/pslang/jit/executable.hpp +++ b/libs/jit/include/pslang/jit/executable.hpp @@ -1,12 +1,13 @@ #pragma once -#include +#include + +#include +#include namespace pslang::jit { - blob allocate(std::size_t size); - - compiled_module make_host_executable(compiled_module module); + blob make_host_executable(std::vector const & code); } diff --git a/libs/jit/include/pslang/jit/jit.hpp b/libs/jit/include/pslang/jit/jit.hpp index 405bc97..f1e4ea7 100644 --- a/libs/jit/include/pslang/jit/jit.hpp +++ b/libs/jit/include/pslang/jit/jit.hpp @@ -1,12 +1,11 @@ #pragma once #include -#include -#include +#include namespace pslang::jit { - compiled_module compile(ast::statement_list_ptr const & statements, abi abi); + void compile(program_context & context, ast::statement_list_ptr const & statements); } diff --git a/libs/jit/include/pslang/jit/program_context.hpp b/libs/jit/include/pslang/jit/program_context.hpp new file mode 100644 index 0000000..5635257 --- /dev/null +++ b/libs/jit/include/pslang/jit/program_context.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include + +#include +#include +#include +#include + +namespace pslang::jit +{ + + struct program_context + { + jit::abi abi; + std::vector code = {}; + std::unordered_map symbols = {}; + }; + +} diff --git a/libs/jit/source/arch/aarch64/compiler.cpp b/libs/jit/source/arch/aarch64/compiler.cpp index 9101129..b581307 100644 --- a/libs/jit/source/arch/aarch64/compiler.cpp +++ b/libs/jit/source/arch/aarch64/compiler.cpp @@ -16,11 +16,8 @@ namespace pslang::jit::aarch64 namespace { - struct context + struct local_context { - std::vector code; - std::unordered_map code_symbol_table; - std::unordered_map f16_constants; std::unordered_map f32_constants; std::unordered_map f64_constants; @@ -42,7 +39,8 @@ namespace pslang::jit::aarch64 using const_expression_visitor::apply; using const_statement_visitor::apply; - context & context; + program_context & pcontext; + local_context & lcontext; template requires(!std::is_floating_point_v) @@ -51,27 +49,27 @@ namespace pslang::jit::aarch64 void apply(ast::f16_literal const & node) { - if (!context.f16_constants.contains(node.value.repr)) + if (!lcontext.f16_constants.contains(node.value.repr)) { - context.f16_constants[node.value.repr] = context.code.size(); + lcontext.f16_constants[node.value.repr] = pcontext.code.size(); push_bytes(node.value.repr); } } void apply(ast::f32_literal const & node) { - if (!context.f32_constants.contains(node.value)) + if (!lcontext.f32_constants.contains(node.value)) { - context.f32_constants[node.value] = context.code.size(); + lcontext.f32_constants[node.value] = pcontext.code.size(); push_bytes(node.value); } } void apply(ast::f64_literal const & node) { - if (!context.f64_constants.contains(node.value)) + if (!lcontext.f64_constants.contains(node.value)) { - context.f64_constants[node.value] = context.code.size(); + lcontext.f64_constants[node.value] = pcontext.code.size(); push_bytes(node.value); } } @@ -182,7 +180,7 @@ namespace pslang::jit::aarch64 { auto begin = (std::uint8_t const *)(&value); auto end = begin + sizeof(value); - context.code.insert(context.code.end(), begin, end); + pcontext.code.insert(pcontext.code.end(), begin, end); } }; @@ -236,8 +234,9 @@ namespace pslang::jit::aarch64 using const_statement_visitor::apply; using const_expression_visitor::apply; - context & context; - instruction_builder builder{context.code}; + program_context & pcontext; + local_context & lcontext; + instruction_builder builder{pcontext.code}; // Difference between initial stack pointer at function enter // and current virtual stack pointer value. The actual stack pointer @@ -307,23 +306,23 @@ namespace pslang::jit::aarch64 void apply(ast::f16_literal const & node) { - auto offset = context.f16_constants.at(node.value.repr); - std::int32_t current = context.code.size(); + auto offset = lcontext.f16_constants.at(node.value.repr); + std::int32_t current = pcontext.code.size(); builder.ldr_fp_pc(0, 0, (offset - current) / 4); builder.fcvt(0, 0b10, 0, 0b01); } void apply(ast::f32_literal const & node) { - auto offset = context.f32_constants.at(node.value); - std::int32_t current = context.code.size(); + auto offset = lcontext.f32_constants.at(node.value); + std::int32_t current = pcontext.code.size(); builder.ldr_fp_pc(0, 0, (offset - current) / 4); } void apply(ast::f64_literal const & node) { - auto offset = context.f64_constants.at(node.value); - std::int32_t current = context.code.size(); + auto offset = lcontext.f64_constants.at(node.value); + std::int32_t current = pcontext.code.size(); builder.ldr_fp_pc(0, 1, (offset - current) / 4); } @@ -631,7 +630,7 @@ namespace pslang::jit::aarch64 if (block.condition) { apply(*block.condition); - branch_skip = context.code.size(); + branch_skip = pcontext.code.size(); builder.cbz(0, 0); } @@ -642,30 +641,30 @@ namespace pslang::jit::aarch64 if (i + 1 < node.blocks.size()) { - branch_to_end.push_back(context.code.size()); + branch_to_end.push_back(pcontext.code.size()); builder.b(0); } if (branch_skip) { - auto branch_offset = context.code.size() - *branch_skip; - builder.cb_inject(context.code.data() + *branch_skip, branch_offset / 4); + auto branch_offset = pcontext.code.size() - *branch_skip; + builder.cb_inject(pcontext.code.data() + *branch_skip, branch_offset / 4); } } - auto end = context.code.size(); + auto end = pcontext.code.size(); for (auto instruction : branch_to_end) { auto delta = end - instruction; - builder.b_inject(context.code.data() + instruction, delta / 4); + builder.b_inject(pcontext.code.data() + instruction, delta / 4); } } void apply(ast::while_block const & node) { - std::int32_t start = context.code.size(); + std::int32_t start = pcontext.code.size(); apply(*node.condition); - std::int32_t skip = context.code.size(); + std::int32_t skip = pcontext.code.size(); builder.cbz(0, 0); scopes.emplace_back(); @@ -673,12 +672,12 @@ namespace pslang::jit::aarch64 scope_cleanup(); scopes.pop_back(); - std::int32_t loop = context.code.size(); + std::int32_t loop = pcontext.code.size(); builder.b(0); - std::int32_t end = context.code.size(); + std::int32_t end = pcontext.code.size(); - builder.cb_inject(context.code.data() + skip, (end - skip) / 4); - builder.b_inject(context.code.data() + loop, (start - loop) / 4); + builder.cb_inject(pcontext.code.data() + skip, (end - skip) / 4); + builder.b_inject(pcontext.code.data() + loop, (start - loop) / 4); } void apply(ast::return_statement const & node) @@ -793,8 +792,9 @@ namespace pslang::jit::aarch64 { using const_statement_visitor::apply; - context & context; - instruction_builder builder{context.code}; + program_context & pcontext; + local_context & lcontext; + instruction_builder builder{pcontext.code}; template void apply(Statement const &) @@ -804,34 +804,21 @@ namespace pslang::jit::aarch64 void apply(ast::function_definition const & node) { - context.code_symbol_table[node.name] = context.code.size(); - compile_function_visitor visitor{{}, {}, context}; + pcontext.symbols[node.name] = pcontext.code.size(); + compile_function_visitor visitor{{}, {}, pcontext, lcontext}; visitor.do_apply(node); } }; } - compiled_module compile(ast::statement_list_ptr const & statements) + void compile(program_context & pcontext, ast::statement_list_ptr const & statements) { - context context; + local_context lcontext; - populate_constants_visitor{{}, {}, context}.apply(*statements); + populate_constants_visitor{{}, {}, pcontext, lcontext}.apply(*statements); - compile_visitor{{}, context}.apply(*statements); - - auto code = allocate(context.code.size()); - std::copy(context.code.data(), context.code.data() + context.code.size(), code.data.get()); - - return compiled_module { - .data = {}, - .code = { - .memory = std::move(code), - .symbol_table = std::move(context.code_symbol_table), - }, - .entry_point = 0, - .abi = abi::armv8, - }; + compile_visitor{{}, pcontext, lcontext}.apply(*statements); } } \ No newline at end of file diff --git a/libs/jit/source/executable.cpp b/libs/jit/source/executable.cpp index 86fc0f0..3e245d0 100644 --- a/libs/jit/source/executable.cpp +++ b/libs/jit/source/executable.cpp @@ -15,28 +15,15 @@ namespace pslang::jit { - blob allocate(std::size_t size) + blob make_host_executable(std::vector const & code) { #if defined(__linux__) || defined(__APPLE__) +auto size = code.size(); auto ptr = (std::uint8_t *)mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, 0, 0); - auto shared_ptr = std::shared_ptr(ptr, [size](void * ptr){ munmap(ptr, size); }); - return blob(shared_ptr, size); -#else - throw std::runtime_error("Allocate not supported for this platform"); -#endif - } - - compiled_module make_host_executable(compiled_module module) - { -#if defined(__linux__) || defined(__APPLE__) - if (module.abi != abi::itanium && module.abi != abi::armv8) - throw std::runtime_error("Abi mismatch"); - - // Assume the module code memory was obtained via mmap - if (mprotect(module.code.memory.data.get(), module.code.memory.size, PROT_READ | PROT_EXEC) != 0) + std::copy(code.begin(), code.end(), ptr); + if (mprotect(ptr, size, PROT_READ | PROT_EXEC) != 0) throw std::system_error(errno, std::generic_category()); - - return module; + return blob(std::shared_ptr(ptr, [size](void * ptr){ munmap(ptr, size); }), size); #else throw std::runtime_error("Host-executable modules are not supported for this platform"); #endif diff --git a/libs/jit/source/jit.cpp b/libs/jit/source/jit.cpp index 322405d..57ba72e 100644 --- a/libs/jit/source/jit.cpp +++ b/libs/jit/source/jit.cpp @@ -6,16 +6,17 @@ namespace pslang::jit { - compiled_module compile(ast::statement_list_ptr const & statements, jit::abi abi) + void compile(program_context & context, ast::statement_list_ptr const & statements) { - switch (abi) + switch (context.abi) { case abi::itanium: throw std::runtime_error("Itanium ABI JIT not implemented"); case abi::msvc: throw std::runtime_error("MSVC ABI JIT not implemented"); case abi::armv8: - return aarch64::compile(statements); + aarch64::compile(context, statements); + break; } }