From caf9cbbdc82a93925f4b3c2eeb44823aa53ac226 Mon Sep 17 00:00:00 2001
From: SachinVin <26602104+SachinVin@users.noreply.github.com>
Date: Tue, 22 Nov 2022 02:58:00 +0530
Subject: [PATCH] oaknut: Implement ADRL and MOVP2R pseudo instructions (#2)

* offset.hpp: constify PageOffset

* okanut: Implement ADRL and MOVP2R pseudo instructions
---
 include/oaknut/impl/offset.hpp |  4 +--
 include/oaknut/oaknut.hpp      | 20 ++++++++++++++-
 tests/basic.cpp                | 46 ++++++++++++++++++++++++++++++++++
 3 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/include/oaknut/impl/offset.hpp b/include/oaknut/impl/offset.hpp
index 470de8f..db65fe3 100644
--- a/include/oaknut/impl/offset.hpp
+++ b/include/oaknut/impl/offset.hpp
@@ -66,7 +66,7 @@ private:
 
 template<std::size_t bitsize, std::size_t shift_amount>
 struct PageOffset {
-    PageOffset(void* ptr)
+    PageOffset(const void* ptr)
         : m_payload(ptr)
     {}
 
@@ -86,7 +86,7 @@ struct PageOffset {
 private:
     template<typename Policy>
     friend class BasicCodeGenerator;
-    std::variant<Label*, void*> m_payload;
+    std::variant<Label*, const void*> m_payload;
 };
 
 template<std::size_t bitsize, std::size_t alignment>
diff --git a/include/oaknut/oaknut.hpp b/include/oaknut/oaknut.hpp
index e6ab27e..e512afc 100644
--- a/include/oaknut/oaknut.hpp
+++ b/include/oaknut/oaknut.hpp
@@ -110,6 +110,12 @@ public:
         return RET(XReg{30});
     }
 
+    void ADRL(XReg xd, const void* addr)
+    {
+        ADRP(xd, addr);
+        ADD(xd, xd, reinterpret_cast<uint64_t>(addr) & 0xFFF);
+    }
+
     void MOV(WReg wd, uint32_t imm)
     {
         if (wd.index() == 31)
@@ -163,6 +169,18 @@ public:
         }
     }
 
+    // Convenience function for moving pointers to registers 
+    void MOVP2R(XReg xd, const void* addr) {
+        int64_t diff = reinterpret_cast<uint64_t>(addr) - Policy::current_address();
+        if (diff >= -0xF'FFFF && diff <= 0xF'FFFF) {
+            ADR(xd, addr);
+        } else if (diff >= -int64_t{0xFFFF'FFFF} && diff <= int64_t{0xFFFF'FFFF}) {
+            ADRL(xd, addr);
+        } else {
+            MOV(xd, reinterpret_cast<uint64_t>(addr));
+        }
+    }
+
     void align(std::size_t alignment)
     {
         if (alignment < 4 || (alignment & (alignment - 1)) != 0)
@@ -242,7 +260,7 @@ private:
                                   label->m_wbs.emplace_back(Label::Writeback{Policy::current_address(), ~splat, static_cast<Label::EmitFunctionType>(encode_fn)});
                                   return 0u;
                               },
-                              [&](void* p) {
+                              [&](const void* p) {
                                   return encode_fn(Policy::current_address(), reinterpret_cast<std::uintptr_t>(p));
                               },
                           },
diff --git a/tests/basic.cpp b/tests/basic.cpp
index 9bcd28f..e1d3c5e 100644
--- a/tests/basic.cpp
+++ b/tests/basic.cpp
@@ -3,6 +3,7 @@
 
 #include <cstdint>
 #include <cstdio>
+#include <limits>
 
 #include <catch2/catch_test_macros.hpp>
 
@@ -159,3 +160,48 @@ TEST_CASE("ADRP")
         REQUIRE(f() == expect);
     }
 }
+
+TEST_CASE("ADRL")
+{
+    CodeBlock mem{4096};
+
+    for (int i = 0; i < 0x200000; i++) {
+        const std::int64_t diff = RandInt<std::int64_t>(-4294967296, 4294967295);
+        const std::intptr_t value = reinterpret_cast<std::intptr_t>(mem.ptr()) + diff;
+
+        CodeGenerator code{mem.ptr()};
+
+        auto f = code.ptr<std::uint64_t (*)()>();
+        mem.unprotect();
+        code.ADRL(X0, reinterpret_cast<void*>(value));
+        code.RET();
+        mem.protect();
+        mem.invalidate_all();
+
+        INFO(i);
+        REQUIRE(f() == static_cast<std::uint64_t>(value));
+    }
+}
+
+TEST_CASE("MOVP2R")
+{
+    CodeBlock mem{4096};
+
+    for (int i = 0; i < 0x200'0000; i++) 
+    {
+        const std::int64_t diff = RandInt<std::int64_t>(std::numeric_limits<std::int64_t>::min(), 
+                                                        std::numeric_limits<std::int64_t>::max());
+        const std::intptr_t value = reinterpret_cast<std::intptr_t>(mem.ptr()) + diff;
+
+        CodeGenerator code{mem.ptr()};
+
+        auto f = code.ptr<std::uint64_t (*)()>();
+        mem.unprotect();
+        code.MOVP2R(X0, reinterpret_cast<void*>(value));
+        code.RET();
+        mem.protect();
+        mem.invalidate_all();
+
+        REQUIRE(f() == static_cast<std::uint64_t>(value));
+    }
+}