CodeGenerator: Add MOV (immediate)

include/oaknut/impl/imm.hpp

@@ -114,7 +114,7 @@ public:
 private:
     template<typename Policy>
     friend class BasicCodeGenerator;
-    std::uint32_t m_encoded;
+    std::uint32_t m_encoded = 0;
 };
 
 namespace detail {

include/oaknut/oaknut.hpp

@@ -99,6 +99,64 @@ public:
 #include "oaknut/impl/arm64_mnemonics.inc.hpp"
 #include "oaknut/impl/fpsimd_mnemonics.inc.hpp"
 
+    void RET()
+    {
+        return RET(XReg{30});
+    }
+
+    void MOV(WReg wd, uint32_t imm)
+    {
+        if (wd.index() == 31)
+            return;
+        if (MovImm16::is_valid(imm))
+            return MOVZ(wd, imm);
+        if (MovImm16::is_valid(~static_cast<std::uint64_t>(imm)))
+            return MOVN(wd, imm);
+        if (detail::encode_bit_imm(imm))
+            return ORR(wd, WzrReg{}, imm);
+
+        MOVZ(wd, {static_cast<std::uint16_t>(imm >> 0), MovImm16Shift::SHL_0});
+        MOVK(wd, {static_cast<std::uint16_t>(imm >> 16), MovImm16Shift::SHL_16});
+    }
+
+    void MOV(XReg xd, uint64_t imm)
+    {
+        if (xd.index() == 31)
+            return;
+        if (imm >> 32 == 0)
+            return MOV(xd.toW(), static_cast<std::uint32_t>(imm));
+        if (MovImm16::is_valid(imm))
+            return MOVZ(xd, imm);
+        if (MovImm16::is_valid(~imm))
+            return MOVN(xd, imm);
+        if (detail::encode_bit_imm(imm))
+            return ORR(xd, ZrReg{}, imm);
+
+        bool movz_done = false;
+        int shift_count = 0;
+
+        if (detail::encode_bit_imm(static_cast<std::uint32_t>(imm))) {
+            ORR(xd.toW(), WzrReg{}, static_cast<std::uint32_t>(imm));
+            imm >>= 32;
+            movz_done = true;
+            shift_count = 2;
+        }
+
+        while (imm != 0) {
+            const uint16_t hw = static_cast<uint16_t>(imm);
+            if (hw != 0) {
+                if (movz_done) {
+                    MOVK(xd, {hw, static_cast<MovImm16Shift>(shift_count)});
+                } else {
+                    MOVZ(xd, {hw, static_cast<MovImm16Shift>(shift_count)});
+                    movz_done = true;
+                }
+            }
+            imm >>= 16;
+            shift_count++;
+        }
+    }
+
 private:
 #include "oaknut/impl/arm64_encode_helpers.inc.hpp"
 

tests/basic.cpp

@@ -8,19 +8,20 @@
 
 #include "oaknut/code_block.hpp"
 #include "oaknut/oaknut.hpp"
+#include "rand_int.hpp"
+
+using namespace oaknut;
+using namespace oaknut::util;
 
 TEST_CASE("Basic Test")
 {
-    using namespace oaknut;
-    using namespace oaknut::util;
-
     CodeBlock mem{4096};
     CodeGenerator code{mem.ptr()};
 
     mem.unprotect();
 
-    code.MOVZ(W0, 42);
-    code.RET(X30);
+    code.MOV(W0, 42);
+    code.RET();
 
     mem.protect();
     mem.invalidate_all();
@@ -31,9 +32,6 @@ TEST_CASE("Basic Test")
 
 TEST_CASE("Fibonacci")
 {
-    using namespace oaknut;
-    using namespace oaknut::util;
-
     CodeBlock mem{4096};
     CodeGenerator code{mem.ptr()};
 
@@ -50,7 +48,7 @@ TEST_CASE("Fibonacci")
     code.SUBS(W0, W0, 1);
     code.B(LT, zero);
     code.B(NE, recurse);
-    code.MOVZ(W0, 1);
+    code.MOV(W0, 1);
     code.B(end);
 
     code.l(zero);
@@ -67,7 +65,7 @@ TEST_CASE("Fibonacci")
     code.l(end);
     code.LDP(X20, X19, SP, 16);
     code.LDP(X29, X30, SP, POST_INDEXED, 32);
-    code.RET(X30);
+    code.RET();
 
     mem.protect();
     mem.invalidate_all();
@@ -77,3 +75,43 @@ TEST_CASE("Fibonacci")
     REQUIRE(fib(5) == 5);
     REQUIRE(fib(9) == 34);
 }
+
+TEST_CASE("Immediate generation (32-bit)")
+{
+    CodeBlock mem{4096};
+
+    for (int i = 0; i < 0x100000; i++) {
+        const std::uint32_t value = RandInt<std::uint32_t>(0, 0xffffffff);
+
+        CodeGenerator code{mem.ptr()};
+
+        auto f = code.ptr<std::uint64_t (*)()>();
+        mem.unprotect();
+        code.MOV(W0, value);
+        code.RET();
+        mem.protect();
+        mem.invalidate_all();
+
+        REQUIRE(f() == value);
+    }
+}
+
+TEST_CASE("Immediate generation (64-bit)")
+{
+    CodeBlock mem{4096};
+
+    for (int i = 0; i < 0x100000; i++) {
+        const std::uint64_t value = RandInt<std::uint64_t>(0, 0xffffffff'ffffffff);
+
+        CodeGenerator code{mem.ptr()};
+
+        auto f = code.ptr<std::uint64_t (*)()>();
+        mem.unprotect();
+        code.MOV(X0, value);
+        code.RET();
+        mem.protect();
+        mem.invalidate_all();
+
+        REQUIRE(f() == value);
+    }
+}

tests/rand_int.hpp

@@ -0,0 +1,20 @@
+// SPDX-FileCopyrightText: Copyright (c) 2022 merryhime <https://mary.rs>
+// SPDX-License-Identifier: MIT
+
+#pragma once
+
+#include <random>
+#include <type_traits>
+
+template<typename T>
+T RandInt(T min, T max)
+{
+    static_assert(std::is_integral_v<T>, "T must be an integral type.");
+    static_assert(!std::is_same_v<T, signed char> && !std::is_same_v<T, unsigned char>,
+                  "Using char with uniform_int_distribution is undefined behavior.");
+
+    static std::random_device rd;
+    static std::mt19937 mt(rd());
+    std::uniform_int_distribution<T> rand(min, max);
+    return rand(mt);
+}