1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+11
+12
+13
+14
+15
+16
+17
+18
+19
+20
+21
+22
+23
+24
+25
+26
+27
+28
+29
+30
+31
+32
+33
+34
+35
+36
+37
+38
+39
+40
+41
+42
+43
+44
+45
+46
+47
+48
+49
+50
+51
+52
+53
+54
+55
+56
+57
+58
+59
+60
+61
+62
+63
+64
+65
+66
+67
+68
+69
+70
+71
+72
+73
+74
+75
+76
+77
+78
+79
+80
+81
+82
+83
+84
+85
+86
+87
+88
+89
+90
+91
+92
+93
+94
+95
+96
+97
+98
+99
+100
+101
+102
+103
+104
+105
+106
+107
+108
+109
+110
+111
+112
+113
+114
+115
+116
+117
+118
+119
+120
+121
+122
+123
+124
+125
+126
+127
+128
+129
+130
+131
+132
+133
+134
+135
+136
+137
+138
+139
+140
+141
+142
+143
+144
+145
+146
+147
+148
+149
+150
+151
+152
+153
+154
+155
+156
+157
+158
+159
+160
+161
+162
+163
+164
+165
+166
+167
+168
+169
+170
+171
+172
+173
+174
+175
+176
+177
+178
+179
+180
+181
+182
+183
+184
+185
+186
+187
+188
+189
+190
+191
+192
+193
+194
+195
+196
+197
+198
+199
+200
+201
+202
+203
+204
+205
+206
+207
+208
+209
+210
+211
+212
+213
+214
+215
+216
+217
+218
+219
+220
+221
+222
+223
+224
+225
+226
+227
+228
+229
+230
+231
+232
+233
+234
+235
+236
+237
+238
+239
+240
+241
+242
+243
+244
+245
+246
+247
+248
+249
+250
+251
+252
+253
+254
+255
+256
+257
+258
+259
+260
+261
+262
+263
+264
+265
+266
+267
+268
+269
+270
+271
+272
+273
+274
+275
+276
+277
+278
+279
+280
+281
+282
+283
+284
+285
+286
+287
+288
+289
+290
+291
+292
+293
+294
+295
+296
+297
+298
+299
+300
+301
+302
+303
+304
+305
+306
+307
+308
+309
+310
+311
+312
+313
+314
+315
+316
+317
+318
+319
+320
+321
+322
+323
+324
+325
+326
+327
+328
+329
+330
+331
+332
+333
+334
+335
+336
+337
+338
+339
+340
+341
+342
+343
+344
+

//! The `x64` jit assembler.
+
+use crate::*;
+use imm::Imm;
+use reg::Reg;
+
+/// Encode the `REX` byte.
+const fn rex(w: bool, r: u8, x: u8, b: u8) -> u8 {
+    let w = if w { 1 } else { 0 };
+    let r = (r >> 3) & 1;
+    let x = (x >> 3) & 1;
+    let b = (b >> 3) & 1;
+    0b0100_0000 | ((w & 1) << 3) | (r << 2) | (x << 1) | b
+}
+
+/// Encode the `ModR/M` byte.
+const fn modrm(mod_: u8, reg: u8, rm: u8) -> u8 {
+    ((mod_ & 0b11) << 6) | ((reg & 0b111) << 3) | (rm & 0b111)
+}
+
+/// `x64` jit assembler.
+pub struct Asm {
+    buf: Vec<u8>,
+}
+
+impl Asm {
+    /// Create a new `x64` jit assembler.
+    pub fn new() -> Asm {
+        // Some random default capacity.
+        let buf = Vec::with_capacity(1024);
+        Asm { buf }
+    }
+
+    /// Consume the assembler and get the emitted code.
+    pub fn into_code(self) -> Vec<u8> {
+        self.buf
+    }
+
+    /// Emit a slice of bytes.
+    pub(crate) fn emit(&mut self, bytes: &[u8]) {
+        self.buf.extend_from_slice(bytes);
+    }
+
+    /// Emit a slice of optional bytes.
+    fn emit_optional(&mut self, bytes: &[Option<u8>]) {
+        for byte in bytes.iter().filter_map(|&b| b) {
+            self.buf.push(byte);
+        }
+    }
+
+    /// Emit a slice of bytes at `pos`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if [pos..pos+len] indexes out of bound of the underlying code buffer.
+    fn emit_at(&mut self, pos: usize, bytes: &[u8]) {
+        if let Some(buf) = self.buf.get_mut(pos..pos + bytes.len()) {
+            buf.copy_from_slice(bytes);
+        } else {
+            unimplemented!();
+        }
+    }
+
+    /// 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.
+    fn resolve(&mut self, label: &mut Label) {
+        if let Some(loc) = label.location() {
+            // For now we only support disp32 as 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.
+
+    /// Encode an register-register instruction.
+    pub(crate) fn encode_rr<T: Reg>(&mut self, opc: u8, op1: T, op2: T)
+    where
+        Self: EncodeRR<T>,
+    {
+        // MR operand encoding.
+        //   op1 -> modrm.rm
+        //   op2 -> modrm.reg
+        let modrm = modrm(
+            0b11,      /* mod */
+            op2.idx(), /* reg */
+            op1.idx(), /* rm */
+        );
+
+        let prefix = <Self as EncodeRR<T>>::legacy_prefix();
+        let rex = <Self as EncodeRR<T>>::rex(op1, op2);
+
+        self.emit_optional(&[prefix, rex]);
+        self.emit(&[opc, modrm]);
+    }
+
+    /// Encode an offset-immediate instruction.
+    /// Register idx is encoded in the opcode.
+    pub(crate) fn encode_oi<T: Reg, U: Imm>(&mut self, opc: u8, op1: T, op2: U)
+    where
+        Self: EncodeR<T>,
+    {
+        let opc = opc + (op1.idx() & 0b111);
+        let prefix = <Self as EncodeR<T>>::legacy_prefix();
+        let rex = <Self as EncodeR<T>>::rex(op1);
+
+        self.emit_optional(&[prefix, rex]);
+        self.emit(&[opc]);
+        self.emit(op2.bytes());
+    }
+
+    /// Encode a register instruction.
+    pub(crate) fn encode_r<T: Reg>(&mut self, opc: u8, opc_ext: u8, op1: T)
+    where
+        Self: EncodeR<T>,
+    {
+        // M operand encoding.
+        //   op1           -> modrm.rm
+        //   opc extension -> modrm.reg
+        let modrm = modrm(
+            0b11,      /* mod */
+            opc_ext,   /* reg */
+            op1.idx(), /* rm */
+        );
+
+        let prefix = <Self as EncodeR<T>>::legacy_prefix();
+        let rex = <Self as EncodeR<T>>::rex(op1);
+
+        self.emit_optional(&[prefix, rex]);
+        self.emit(&[opc, modrm]);
+    }
+
+    /// Encode a memory-immediate instruction.
+    pub(crate) fn encode_mi<T: Imm>(&mut self, opc: u8, opc_ext: u8, op1: MemOp, op2: T)
+    where
+        Self: EncodeMI<T>,
+    {
+        // MI operand encoding.
+        //   op1 -> modrm.rm
+        //   op2 -> imm
+        let mode = match op1 {
+            MemOp::Indirect(..) => {
+                assert!(!op1.base().need_sib() && !op1.base().is_pc_rel());
+                0b00
+            }
+            MemOp::IndirectDisp(..) => {
+                assert!(!op1.base().need_sib());
+                0b10
+            }
+        };
+
+        let modrm = modrm(
+            mode,             /* mode */
+            opc_ext,          /* reg */
+            op1.base().idx(), /* rm */
+        );
+
+        let prefix = <Self as EncodeMI<T>>::legacy_prefix();
+        let rex = <Self as EncodeMI<T>>::rex(&op1);
+
+        self.emit_optional(&[prefix, rex]);
+        self.emit(&[opc, modrm]);
+        if let MemOp::IndirectDisp(_, disp) = op1 {
+            self.emit(&disp.to_ne_bytes());
+        }
+        self.emit(op2.bytes());
+    }
+
+    /// Encode a memory-register instruction.
+    pub(crate) fn encode_mr<T: Reg>(&mut self, opc: u8, op1: MemOp, op2: T)
+    where
+        Self: EncodeMR<T>,
+    {
+        // MR operand encoding.
+        //   op1 -> modrm.rm
+        //   op2 -> modrm.reg
+        let mode = match op1 {
+            MemOp::Indirect(..) => {
+                assert!(!op1.base().need_sib() && !op1.base().is_pc_rel());
+                0b00
+            }
+            MemOp::IndirectDisp(..) => {
+                assert!(!op1.base().need_sib());
+                0b10
+            }
+        };
+
+        let modrm = modrm(
+            mode,             /* mode */
+            op2.idx(),        /* reg */
+            op1.base().idx(), /* rm */
+        );
+        let prefix = <Self as EncodeMR<T>>::legacy_prefix();
+        let rex = <Self as EncodeMR<T>>::rex(&op1, op2);
+
+        self.emit_optional(&[prefix, rex]);
+        self.emit(&[opc, modrm]);
+        if let MemOp::IndirectDisp(_, disp) = op1 {
+            self.emit(&disp.to_ne_bytes());
+        }
+    }
+
+    /// Encode a register-memory instruction.
+    pub(crate) fn encode_rm<T: Reg>(&mut self, opc: u8, op1: T, op2: MemOp)
+    where
+        Self: EncodeMR<T>,
+    {
+        // RM operand encoding.
+        //   op1 -> modrm.reg
+        //   op2 -> modrm.rm
+        self.encode_mr(opc, op2, op1);
+    }
+
+    /// Encode a jump to label instruction.
+    pub(crate) 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.
+        // We currently only support disp32 jump targets.
+        self.emit(&[0u8; 4]);
+
+        // Resolve any pending relocations for the label.
+        self.resolve(op1);
+    }
+}
+
+// -- Encoder helper.
+
+/// Encode helper for register-register instructions.
+pub(crate) trait EncodeRR<T: Reg> {
+    fn legacy_prefix() -> Option<u8> {
+        None
+    }
+
+    fn rex(op1: T, op2: T) -> Option<u8> {
+        if op1.need_rex() || op2.need_rex() {
+            Some(rex(op1.rexw(), op2.idx(), 0, op1.idx()))
+        } else {
+            None
+        }
+    }
+}
+
+impl EncodeRR<Reg8> for Asm {}
+impl EncodeRR<Reg32> for Asm {}
+impl EncodeRR<Reg16> for Asm {
+    fn legacy_prefix() -> Option<u8> {
+        Some(0x66)
+    }
+}
+impl EncodeRR<Reg64> for Asm {}
+
+/// Encode helper for register instructions.
+pub(crate) trait EncodeR<T: Reg> {
+    fn legacy_prefix() -> Option<u8> {
+        None
+    }
+
+    fn rex(op1: T) -> Option<u8> {
+        if op1.need_rex() {
+            Some(rex(op1.rexw(), 0, 0, op1.idx()))
+        } else {
+            None
+        }
+    }
+}
+
+impl EncodeR<Reg8> for Asm {}
+impl EncodeR<Reg32> for Asm {}
+impl EncodeR<Reg16> for Asm {
+    fn legacy_prefix() -> Option<u8> {
+        Some(0x66)
+    }
+}
+impl EncodeR<Reg64> for Asm {}
+
+/// Encode helper for memory-register instructions.
+pub(crate) trait EncodeMR<T: Reg> {
+    fn legacy_prefix() -> Option<u8> {
+        None
+    }
+
+    fn rex(op1: &MemOp, op2: T) -> Option<u8> {
+        if op2.need_rex() || (op1.base().is_ext()) {
+            Some(rex(op2.rexw(), op2.idx(), 0, op1.base().idx()))
+        } else {
+            None
+        }
+    }
+}
+
+impl EncodeMR<Reg8> for Asm {}
+impl EncodeMR<Reg16> for Asm {
+    fn legacy_prefix() -> Option<u8> {
+        Some(0x66)
+    }
+}
+impl EncodeMR<Reg32> for Asm {}
+impl EncodeMR<Reg64> for Asm {}
+
+/// Encode helper for memory-immediate instructions.
+pub(crate) trait EncodeMI<T: Imm> {
+    fn legacy_prefix() -> Option<u8> {
+        None
+    }
+
+    fn rex(op1: &MemOp) -> Option<u8> {
+        if op1.base().is_ext() {
+            Some(rex(false, 0, 0, op1.base().idx()))
+        } else {
+            None
+        }
+    }
+}
+
+impl EncodeMI<Imm8> for Asm {}
+impl EncodeMI<Imm16> for Asm {
+    fn legacy_prefix() -> Option<u8> {
+        Some(0x66)
+    }
+}
+impl EncodeMI<Imm32> for Asm {}
+