azm: Simplify label usage

robinlinden · Dec 5, 2023 · 29805c1 · 29805c1
1 parent 97c48f3
commit 29805c1
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 19 deletions.
diff --git a/azm/amd64/assembler.h b/azm/amd64/assembler.h
@@ -5,11 +5,13 @@
 #ifndef AZM_AMD64_ASSEMBLER_H_
 #define AZM_AMD64_ASSEMBLER_H_
 
+#include <cassert>
 #include <cstddef>
 #include <cstdint>
 #include <iostream>
 #include <optional>
 #include <utility>
+#include <variant>
 #include <vector>
 
 namespace azm::amd64 {
@@ -40,34 +42,42 @@ constexpr std::optional<std::uint8_t> register_index(Reg32 reg) {
 }
 
 struct Label {
-    std::size_t offset{};
-};
+    struct Linked {
+        std::size_t offset{};
+    };
+    struct Unlinked {
+        std::vector<std::size_t> patch_offsets{};
+    };
+
+    static Label linked(std::size_t jmp_target_offset) { return {Linked{jmp_target_offset}}; }
+    static Label unlinked() { return {Unlinked{}}; }
 
-struct UnlinkedLabel {
-    std::vector<std::size_t> patch_offsets{};
+    std::variant<Linked, Unlinked> v;
 };
 
 // https://www.felixcloutier.com/x86/
 class Assembler {
 public:
     [[nodiscard]] std::vector<std::uint8_t> take_assembled() { return std::exchange(assembled_, {}); }
 
-    Label label() const { return Label{assembled_.size()}; }
-    UnlinkedLabel unlinked_label() const { return UnlinkedLabel{}; }
+    Label label() const { return Label::linked(assembled_.size()); }
+    Label unlinked_label() const { return Label::unlinked(); }
 
-    Label link(UnlinkedLabel const &label) {
+    void link(Label &label) {
+        assert(std::holds_alternative<Label::Unlinked>(label.v));
         static constexpr int kInstructionSize = 4;
         std::size_t const jmp_target_offset = assembled_.size();
 
-        for (std::size_t patch_offset : label.patch_offsets) {
+        auto const &unlinked = std::get<Label::Unlinked>(label.v);
+        for (std::size_t patch_offset : unlinked.patch_offsets) {
             auto const rel32 = static_cast<std::uint32_t>(jmp_target_offset - patch_offset - kInstructionSize);
             assembled_[patch_offset + 0] = rel32 & 0xff;
             assembled_[patch_offset + 1] = (rel32 >> 8) & 0xff;
             assembled_[patch_offset + 2] = (rel32 >> 16) & 0xff;
             assembled_[patch_offset + 3] = (rel32 >> 24) & 0xff;
         }
 
-        return Label{jmp_target_offset};
+        label = Label::linked(jmp_target_offset);
     }
 
     // Instructions
@@ -82,17 +92,19 @@ class Assembler {
         emit(imm32);
     }
 
-    void jmp(Label label) {
+    void jmp(Label &label) {
         // JMP rel32
-        emit(0xe9);
-        static constexpr int kInstructionSize = 4;
-        emit(Imm32{static_cast<std::uint32_t>(label.offset - assembled_.size() - kInstructionSize)});
-    }
+        if (std::holds_alternative<Label::Linked>(label.v)) {
+            auto const &linked = std::get<Label::Linked>(label.v);
+            static constexpr int kInstructionSize = 4;
+            emit(0xe9);
+            emit(Imm32{static_cast<std::uint32_t>(linked.offset - assembled_.size() - kInstructionSize)});
+            return;
+        }
 
-    void jmp(UnlinkedLabel &label) {
-        // JMP rel32
+        auto &unlinked = std::get<Label::Unlinked>(label.v);
         emit(0xe9);
-        label.patch_offsets.push_back(assembled_.size());
+        unlinked.patch_offsets.push_back(assembled_.size());
         emit(Imm32{0xdeadbeef});
     }
 

diff --git a/azm/amd64/assembler_test.cpp b/azm/amd64/assembler_test.cpp
@@ -80,8 +80,8 @@ int main() {
         assembler.jmp(slot1);
         assembler.ud2();
         assembler.jmp(slot1);
-        auto slot2 = assembler.link(slot1);
-        assembler.jmp(slot2);
+        assembler.link(slot1);
+        assembler.jmp(slot1);
 
         a.expect_eq(assembler.take_assembled(),
                 CodeVec{