Add JMP, JZ, and Label

6 files changed, 147 insertions, 5 deletions

diff --git a/src/insn.rs b/src/insn.rs
index 2bfaf18..be2e689 100644
--- a/src/insn.rs
+++ b/src/insn.rs
@@ -1,5 +1,7 @@
 mod add;
 mod dec;
+mod jmp;
+mod jz;
 mod mov;
 mod ret;
 mod test;
@@ -12,6 +14,14 @@ pub trait Dec<T> {
     fn dec(&mut self, op1: T);
 }
 
+pub trait Jmp<T> {
+    fn jmp(&mut self, op1: T);
+}
+
+pub trait Jz<T> {
+    fn jz(&mut self, op1: T);
+}
+
 pub trait Mov<T, U> {
     fn mov(&mut self, op1: T, op2: U);
 }
diff --git a/src/insn/jmp.rs b/src/insn/jmp.rs
new file mode 100644
index 0000000..71d1dbc
--- /dev/null
+++ b/src/insn/jmp.rs
@@ -0,0 +1,7 @@
+use crate::prelude::*;
+
+impl Jmp<&mut Label> for Asm {
+    fn jmp(&mut self, op1: &mut Label) {
+        self.encode_jmp_label(&[0xe9], op1);
+    }
+}
diff --git a/src/insn/jz.rs b/src/insn/jz.rs
new file mode 100644
index 0000000..6563ca2
--- /dev/null
+++ b/src/insn/jz.rs
@@ -0,0 +1,7 @@
+use crate::prelude::*;
+
+impl Jz<&mut Label> for Asm {
+    fn jz(&mut self, op1: &mut Label) {
+        self.encode_jmp_label(&[0x0f, 0x84], op1);
+    }
+}
diff --git a/src/label.rs b/src/label.rs
new file mode 100644
index 0000000..b1f1133
--- /dev/null
+++ b/src/label.rs
@@ -0,0 +1,72 @@
+use std::collections::HashSet;
+
+/// A label which is used as target for jump instructions.
+///
+/// ```rust
+/// use juicebox_asm::prelude::*;
+///
+/// let mut lbl = Label::new();
+/// let mut asm = Asm::new();
+///
+/// // Skip the mov instruction.
+/// asm.jmp(&mut lbl);
+/// asm.mov(Reg64::rax, Reg64::rax);
+/// asm.bind(&mut lbl);
+/// ```
+///
+/// # Panics
+///
+/// Panics if the label is dropped while not yet bound, or having unresolved relocations.
+/// This is mainly a safety-guard to detect wrong usage.
+pub struct Label {
+    /// Location of the label. Will be set after the label is bound, else None.
+    location: Option<usize>,
+
+    /// Offsets that must be patched with the label location.
+    offsets: HashSet<usize>,
+}
+
+impl Label {
+    /// Create a new `unbound` [Label].
+    pub fn new() -> Label {
+        Label {
+            location: None,
+            offsets: HashSet::new(),
+        }
+    }
+
+    /// Bind the label to the `location`.
+    pub(crate) fn bind(&mut self, loc: usize) {
+        // A label can only be bound once!
+        assert!(!self.is_bound());
+
+        self.location = Some(loc);
+    }
+
+    /// Record an offset that must be patched with the label location.
+    pub(crate) fn record_offset(&mut self, off: usize) {
+        self.offsets.insert(off);
+    }
+
+    pub(crate) fn location(&self) -> Option<usize> {
+        self.location
+    }
+
+    pub(crate) fn offsets_mut(&mut self) -> &mut HashSet<usize> {
+        &mut self.offsets
+    }
+
+    /// Check whether the label is bound to a location.
+    const fn is_bound(&self) -> bool {
+        self.location.is_some()
+    }
+}
+
+impl Drop for Label {
+    fn drop(&mut self) {
+        // Ensure the label was bound when it is dropped.
+        assert!(self.is_bound());
+        // Ensure all offsets have been patched when the label is dropped.
+        assert!(self.offsets.is_empty());
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 27164ad..3368e3d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,12 +2,13 @@ pub mod prelude;
 
 mod imm;
 mod insn;
+mod label;
 mod reg;
 
 use imm::Imm;
-pub use imm::{Imm16, Imm32, Imm64, Imm8};
+use label::Label;
 use reg::Reg;
-pub use reg::{Reg16, Reg32, Reg64, Reg8};
+use reg::{Reg16, Reg32, Reg64, Reg8};
 
 pub enum MemOp {
     Indirect(Reg64),
@@ -69,6 +70,35 @@ impl Asm {
         }
     }
 
+    /// Bind the [Label] to the current location.
+    pub fn bind(&mut self, label: &mut Label) {
+        // Bind the label to the current offset.
+        label.bind(self.buf.len());
+
+        // Resolve any pending relocations for the label.
+        self.resolve(label);
+    }
+
+    /// If the [Label] is bound, patch any pending relocation.
+    pub fn resolve(&mut self, label: &mut Label) {
+        if let Some(loc) = label.location() {
+            let loc = i32::try_from(loc).expect("Label location did not fit into i32.");
+
+            // Resolve any pending relocations for the label.
+            for off in label.offsets_mut().drain() {
+                // Displacement is relative to the next instruction following the jump.
+                // We record the offset to patch at the first byte of the disp32 therefore we need
+                // to account for that in the disp computation.
+                let disp32 = loc - i32::try_from(off).expect("Label offset did not fit into i32") - 4 /* account for the disp32 */;
+
+                // Patch the relocation with the disp32.
+                self.emit_at(off, &disp32.to_ne_bytes());
+            }
+        }
+    }
+
+    // -- Encode utilities.
+
     fn encode_rr<T: Reg>(&mut self, opc: u8, op1: T, op2: T)
     where
         Self: EncodeRR<T>,
@@ -185,6 +215,20 @@ impl Asm {
         //   op2 -> modrm.rm
         self.encode_mr(opc, op2, op1);
     }
+
+    fn encode_jmp_label(&mut self, opc: &[u8], op1: &mut Label) {
+        // Emit the opcode.
+        self.emit(opc);
+
+        // Record relocation offset starting at the first byte of the disp32.
+        op1.record_offset(self.buf.len());
+
+        // Emit a zeroed disp32, which serves as placeholder for the relocation.
+        self.emit(&[0u8; 4]);
+
+        // Resolve any pending relocations for the label.
+        self.resolve(op1);
+    }
 }
 
 // -- Encoder helper.
diff --git a/src/prelude.rs b/src/prelude.rs
index bca972f..29e7265 100644
--- a/src/prelude.rs
+++ b/src/prelude.rs
@@ -1,6 +1,8 @@
 pub use crate::Asm;
 pub use crate::MemOp;
-pub use crate::{Imm16, Imm32, Imm64, Imm8};
-pub use crate::{Reg16, Reg32, Reg64, Reg8};
 
-pub use crate::insn::{Add, Dec, Mov, Test};
+pub use crate::imm::{Imm16, Imm32, Imm64, Imm8};
+pub use crate::label::Label;
+pub use crate::reg::{Reg16, Reg32, Reg64, Reg8};
+
+pub use crate::insn::{Add, Dec, Jmp, Jz, Mov, Test};