├── .gitignore ├── README.md └── src ├── main.zig ├── interpreter.zig └── jit.zig /.gitignore: -------------------------------------------------------------------------------- 1 | zig-cache 2 | zig-out 3 | build 4 | out 5 | tmp 6 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # TinyJIT 2 | 3 | An educational JIT compiler that emits ARM64 machine code for a simple expression language. 4 | 5 | This was made on a series of discord live streams, 6 | and is (very obviously) not intended for production use. 7 | 8 | I'm trying to write a small book about JIT compilation, and this repository 9 | serves as the reference implementation for that. 10 | 11 | -------------------------------------------------------------------------------- /src/main.zig: -------------------------------------------------------------------------------- 1 | const std = @import("std"); 2 | const builtin = @import("builtin"); 3 | const interp = @import("interpreter.zig"); 4 | 5 | const CodeBlock = interp.CodeBlock; 6 | const Interpreter = interp.Interpreter; 7 | const Op = interp.Op; 8 | 9 | fn runProgram(allocator: std.mem.Allocator, program: []const CodeBlock, jit: bool) !void { 10 | var interpreter = try Interpreter.init(allocator, program); 11 | interpreter.is_jit_enabled = jit; 12 | 13 | defer { 14 | interpreter.deinit(); 15 | allocator.destroy(interpreter); 16 | } 17 | 18 | try interpreter.run(); 19 | } 20 | 21 | pub fn main() !void { 22 | var gpa = std.heap.GeneralPurposeAllocator(.{}){}; 23 | const allocator = gpa.allocator(); 24 | 25 | // 0 is x, 1 is i 26 | const b0 = CodeBlock{ 27 | .instructions = &[_]u8{ 28 | Op(.push), 0, // x = 0 29 | Op(.push), 0, // i = 0 30 | Op(.jump), 1, // jump to while loop (block 1 31 | }, 32 | 33 | .constants = &[_]i64{0}, 34 | }; 35 | 36 | const b1 = CodeBlock{ 37 | .instructions = &[_]u8{ 38 | // while i < 1_000_000 39 | Op(.load_var), 1, // load i 40 | Op(.push), 0, // load 1_000_000 41 | Op(.eq), 42 | Op(.jump_nz), 2, // jump to exit if i == 1_000_000 43 | 44 | // x = x + 1 45 | Op(.load_var), 0, // load x 46 | Op(.load_var), 1, // load i 47 | Op(.add), 48 | Op(.store_var), 0, // x = x + i 49 | 50 | // i += 1 51 | Op(.load_var), 1, // load i 52 | Op(.push), 1, // load 1 53 | Op(.add), 54 | Op(.store_var), 1, // i = i + 1 55 | Op(.jump), 1, // jump to start 56 | }, 57 | .constants = &[_]i64{ 5_000_000, 1 }, 58 | }; 59 | 60 | const b2 = CodeBlock{ 61 | .instructions = &[_]u8{ 62 | Op(.load_var), 0, // load x 63 | }, 64 | .constants = &[_]i64{0}, 65 | }; 66 | 67 | const program = [_]CodeBlock{ b0, b1, b2 }; 68 | 69 | var is_jit_enabled = false; 70 | if (std.os.argv.len > 0) { 71 | const arg = std.os.argv[std.os.argv.len - 1]; 72 | var len: usize = 0; 73 | while (arg[len] != 0) len += 1; 74 | 75 | const jit_flag = "--jit"; 76 | is_jit_enabled = std.mem.eql(u8, arg[0..len], jit_flag); 77 | } 78 | 79 | const n_runs = 100; 80 | const before = std.time.milliTimestamp(); 81 | for (0..n_runs) |_| { 82 | try runProgram(allocator, &program, is_jit_enabled); 83 | } 84 | const after = std.time.milliTimestamp(); 85 | 86 | const dt: f64 = @floatFromInt(after - before); 87 | const avg_time = dt / @as(f64, @floatFromInt(n_runs)); 88 | 89 | const msg = try std.fmt.allocPrintZ( 90 | allocator, 91 | "Time per run (JIT = {s}): {d} ms\n", 92 | .{ if (is_jit_enabled) "ON" else "OFF", avg_time }, 93 | ); 94 | defer allocator.free(msg); 95 | 96 | const stdout = std.io.getStdOut(); 97 | defer stdout.close(); 98 | 99 | _ = try stdout.write(msg); 100 | } 101 | -------------------------------------------------------------------------------- /src/interpreter.zig: -------------------------------------------------------------------------------- 1 | const std = @import("std"); 2 | const jit = @import("jit.zig"); 3 | 4 | pub const Opcode = enum(u8) { 5 | push, 6 | add, 7 | eq, 8 | jump_nz, 9 | jump, 10 | load_var, 11 | store_var, 12 | }; 13 | 14 | /// Shorthand to convert an Opcode to a u8 15 | pub inline fn Op(op: Opcode) u8 { 16 | return @intFromEnum(op); 17 | } 18 | 19 | pub const CodeBlock = struct { 20 | /// A list of `Opcode`s cast to u8 21 | instructions: []const u8, 22 | /// Constants used within this block 23 | constants: []const i64, 24 | }; 25 | 26 | pub const Interpreter = struct { 27 | stack: [32_000]i64 = undefined, 28 | stack_pos: u64 = 0, 29 | 30 | allocator: std.mem.Allocator, 31 | 32 | blocks: []const CodeBlock, 33 | jit_blocks: []?jit.CompiledFunction, 34 | 35 | current_block: *const CodeBlock = undefined, 36 | pc: usize = 0, 37 | 38 | jit_compiler: jit.JITCompiler, 39 | is_jit_enabled: bool = false, 40 | 41 | pub fn init(allocator: std.mem.Allocator, blocks: []const CodeBlock) !*Interpreter { 42 | const self = try allocator.create(Interpreter); 43 | self.* = Interpreter{ 44 | .blocks = blocks, 45 | .current_block = &blocks[0], 46 | .jit_compiler = jit.JITCompiler.init(allocator, self), 47 | .jit_blocks = try allocator.alloc(?jit.CompiledFunction, blocks.len), 48 | .allocator = allocator, 49 | }; 50 | 51 | for (self.jit_blocks) |*jit_block| { 52 | jit_block.* = null; 53 | } 54 | 55 | return self; 56 | } 57 | 58 | pub fn deinit(self: *const Interpreter) void { 59 | self.jit_compiler.deinit(); 60 | 61 | // unmap all the JIT functions 62 | for (self.jit_blocks) |maybe_jit_block| { 63 | if (maybe_jit_block) |*jit_block| { 64 | jit_block.deinit(); 65 | } 66 | } 67 | 68 | self.allocator.free(self.jit_blocks); 69 | } 70 | 71 | inline fn loadConst(self: *Interpreter) i64 { 72 | const index = self.current_block.instructions[self.pc]; 73 | const constant = self.current_block.constants[index]; 74 | 75 | self.pc += 1; 76 | return constant; 77 | } 78 | 79 | inline fn top(self: *Interpreter) i64 { 80 | return self.stack[self.stack_pos - 1]; 81 | } 82 | 83 | inline fn pop(self: *Interpreter) i64 { 84 | self.stack_pos -= 1; 85 | return self.stack[self.stack_pos]; 86 | } 87 | 88 | inline fn push(self: *Interpreter, value: i64) void { 89 | self.stack[self.stack_pos] = value; 90 | self.stack_pos += 1; 91 | } 92 | 93 | inline fn nextOp(self: *Interpreter) u8 { 94 | const op = self.current_block.instructions[self.pc]; 95 | self.pc += 1; 96 | return op; 97 | } 98 | 99 | pub fn run(self: *Interpreter) !void { 100 | while (self.pc < self.current_block.instructions.len) { 101 | const instr: Opcode = @enumFromInt(self.nextOp()); 102 | 103 | switch (instr) { 104 | .push => self.push(self.loadConst()), 105 | .add => { 106 | const a = self.pop(); 107 | const b = self.pop(); 108 | self.push(a + b); 109 | }, 110 | .eq => { 111 | const a = self.pop(); 112 | const b = self.pop(); 113 | self.push(if (a == b) 1 else 0); 114 | }, 115 | 116 | .jump => try self.jump(), 117 | .jump_nz => { 118 | if (self.pop() != 0) { 119 | try self.jump(); 120 | } else { 121 | self.pc += 1; 122 | } 123 | }, 124 | 125 | .load_var => { 126 | const index = self.nextOp(); 127 | self.push(self.stack[index]); 128 | }, 129 | 130 | .store_var => { 131 | // [value, index] 132 | const index = self.nextOp(); 133 | self.stack[index] = self.pop(); 134 | }, 135 | } 136 | } 137 | } 138 | 139 | /// Call a JIT compiled function 140 | inline fn callJit(self: *Interpreter, compiled: *const jit.CompiledFunction) void { 141 | var next_block_idx: usize = 0; // inout parameter for JITted function 142 | compiled.func( 143 | (&self.stack).ptr, 144 | self.current_block.instructions.ptr, 145 | &self.stack_pos, 146 | &self.pc, 147 | &next_block_idx, 148 | self.current_block.constants.ptr, 149 | ); 150 | 151 | self.current_block = &self.blocks[next_block_idx]; 152 | // self.pc = 0; 153 | } 154 | 155 | fn doJit(self: *Interpreter, block_index: usize) !void { 156 | const block = &self.blocks[block_index]; 157 | const compiled = try self.jit_compiler.compileBlock(block_index, block); 158 | 159 | self.jit_blocks[block_index] = compiled; 160 | self.callJit(&compiled); 161 | } 162 | 163 | fn jump(self: *Interpreter) !void { 164 | // block index is the next "instruction". 165 | const block_idx = self.nextOp(); 166 | self.pc = 0; // start from first instr in the next block 167 | const dst_block = &self.blocks[block_idx]; 168 | 169 | if (!self.is_jit_enabled) { 170 | self.current_block = dst_block; 171 | return; 172 | } 173 | 174 | // check if the block has been JIT compiled before. 175 | if (self.jit_blocks[block_idx]) |*compiled| { 176 | self.callJit(compiled); 177 | return; 178 | } 179 | 180 | if (self.current_block == dst_block) { 181 | // self-referencing block. potentially a loop. 182 | // JIT compile this. 183 | try self.doJit(block_idx); 184 | return; 185 | } 186 | 187 | // Not a self-referencing block, so do regular execution. 188 | self.current_block = dst_block; 189 | } 190 | }; 191 | -------------------------------------------------------------------------------- /src/jit.zig: -------------------------------------------------------------------------------- 1 | const std = @import("std"); 2 | const interp = @import("interpreter.zig"); 3 | const mman = @cImport(@cInclude("sys/mman.h")); 4 | const pthread = @cImport(@cInclude("pthread.h")); 5 | 6 | const CodeBlock = interp.CodeBlock; 7 | const Interpreter = interp.Interpreter; 8 | const Opcode = interp.Opcode; 9 | const Op = interp.Op; 10 | 11 | const Armv8a = struct { 12 | pub const ret: u32 = 0xd65f03c0; 13 | 14 | // GPRs in AArch64. We only ever use 12 of them. 15 | const Reg = enum(u32) { x0 = 0, x1, x2, x3, x4, x5, x6, x7, x8, x9, x10, x11, x12 }; 16 | 17 | pub inline fn ldrReg(dst_reg: Reg, src_reg: Reg) u32 { 18 | const dst = @intFromEnum(dst_reg); 19 | const src = @intFromEnum(src_reg); 20 | return (0x3E5000 << 10) | (src << 5) | dst; 21 | } 22 | 23 | pub inline fn ldrRegScaled(dst_reg: Reg, base_reg: Reg, offset_reg: Reg) u32 { 24 | const dst = @intFromEnum(dst_reg); 25 | const base = @intFromEnum(base_reg); 26 | const offset = @intFromEnum(offset_reg); 27 | return 0xF8607800 | (offset << 16) | (base << 5) | dst; 28 | } 29 | 30 | pub inline fn ldrRegUnscaled(dst_reg: Reg, base_reg: Reg, offset_reg: Reg) u32 { 31 | const dst = @intFromEnum(dst_reg); 32 | const base = @intFromEnum(base_reg); 33 | const off = @intFromEnum(offset_reg); 34 | return 0xF8606800 | (off << 16) | (base << 5) | dst; 35 | } 36 | 37 | pub inline fn subRegImm(dst_reg: Reg, src_reg: Reg, imm: u32) u32 { 38 | std.debug.assert(imm <= 0b111111_111111); 39 | const src = @intFromEnum(src_reg); 40 | const dst = @intFromEnum(dst_reg); 41 | 42 | return 0xD1000000 | (imm << 10) | (src << 5) | dst; 43 | } 44 | 45 | pub inline fn addRegImm(dst_reg: Reg, src_reg: Reg, imm: u32) u32 { 46 | const dst = @intFromEnum(dst_reg); 47 | const src = @intFromEnum(src_reg); 48 | std.debug.assert(imm <= 0b111111_111111); 49 | return 0x91000400 | (imm << 10) | (src << 5) | dst; 50 | } 51 | 52 | pub inline fn addRegs(dst_reg: Reg, reg_a: Reg, reg_b: Reg) u32 { 53 | const a = @intFromEnum(reg_a); 54 | const b = @intFromEnum(reg_b); 55 | const dst = @intFromEnum(dst_reg); 56 | 57 | return 0x8b000000 | (b << 16) | (a << 5) | dst; 58 | } 59 | 60 | pub inline fn strReg(src_reg: Reg, dst_reg: Reg, offset: u32) u32 { 61 | const dst = @intFromEnum(src_reg); 62 | const src = @intFromEnum(dst_reg); 63 | std.debug.assert(offset <= 0b111111_111111); 64 | return 0xF9000000 | (offset << 10) | (src << 5) | dst; 65 | } 66 | 67 | pub inline fn strRegScaled(src_reg: Reg, base_reg: Reg, offset_reg: Reg) u32 { 68 | const src = @intFromEnum(src_reg); 69 | const base = @intFromEnum(base_reg); 70 | const offset = @intFromEnum(offset_reg); 71 | 72 | return 0xF8207800 | (offset << 16) | (base << 5) | src; 73 | } 74 | 75 | pub inline fn mult8(dst_reg: Reg, src_reg: Reg) u32 { 76 | const dst = @intFromEnum(dst_reg); 77 | const src = @intFromEnum(src_reg); 78 | 79 | return 0xD37DF000 | (src << 5) | dst; 80 | } 81 | 82 | pub inline fn ldrByteReg(dst_reg: Reg, base_reg: Reg, offset_reg: Reg) u32 { 83 | const dst = @intFromEnum(dst_reg); 84 | const base = @intFromEnum(base_reg); 85 | const offset = @intFromEnum(offset_reg); 86 | 87 | return 0x38606800 | (offset << 16) | (base << 5) | dst; 88 | } 89 | 90 | pub inline fn cmpRegs(a_reg: Reg, b_reg: Reg) u32 { 91 | const a = @intFromEnum(a_reg); 92 | const b = @intFromEnum(b_reg); 93 | return 0xEB00001F | (b << 16) | (a << 5); 94 | } 95 | 96 | /// Same as the `cset [reg] eq` instruction in ARM 97 | pub inline fn setIfStatusEq(dst_reg: Reg) u32 { 98 | const dst = @intFromEnum(dst_reg); 99 | return 0x9A9F17E0 | dst; 100 | } 101 | 102 | pub inline fn movRegImm(dst_reg: Reg, imm: u32) u32 { 103 | const dst = @intFromEnum(dst_reg); 104 | return 0xD2800000 | (imm << 5) | dst; 105 | } 106 | 107 | pub inline fn cmpImmediate(reg: Reg, value: u32) u32 { 108 | const r = @intFromEnum(reg); 109 | return 0xF100001F | (value << 10) | (r << 5); 110 | } 111 | 112 | pub inline fn branchIfEq(branch_offset: i32) u32 { 113 | const mask: i32 = std.math.maxInt(u19); 114 | const offset: u32 = @bitCast(branch_offset & mask); 115 | return 0x54000000 | (offset << 5); 116 | } 117 | }; 118 | 119 | pub const JitFunction = *fn ( 120 | stack: [*]i64, // x0 121 | instructions: [*]const u8, // x1 122 | stack_ptr: *usize, // x2 123 | instr_ptr: *usize, // x3 124 | current_block_index: *usize, // x5 125 | constants: [*]const i64, // x6 126 | ) callconv(.C) void; 127 | 128 | pub const CompiledFunction = struct { 129 | func: JitFunction, // both func and buf point to the same data. 130 | buf: [*]u32, 131 | size: usize, 132 | 133 | pub fn init(func: JitFunction, buf: [*]u32, size: usize) CompiledFunction { 134 | return .{ .func = func, .buf = buf, .size = size }; 135 | } 136 | 137 | pub fn deinit(self: *const CompiledFunction) void { 138 | if (mman.munmap(self.buf, self.size) != 0) { 139 | std.debug.panic("munmap failed\n", .{}); 140 | } 141 | } 142 | }; 143 | 144 | pub const JITCompiler = struct { 145 | allocator: std.mem.Allocator, 146 | 147 | interpreter: *const Interpreter, 148 | machine_code: std.ArrayList(u32), 149 | 150 | const ArgReg = struct { 151 | const stackAddr = .x0; 152 | const instructionsAddr = .x1; 153 | const stackIndexPtr = .x2; 154 | const instrIndexPtr = .x3; 155 | const currentBlockNumber = .x4; 156 | const constantsAddr = .x5; 157 | }; 158 | 159 | const VarReg = struct { 160 | const stackIndex = .x8; 161 | const instrIndex = .x12; 162 | const tempA = .x9; 163 | const tempB = .x10; 164 | const tempC = .x11; 165 | }; 166 | 167 | pub fn init(allocator: std.mem.Allocator, interpreter: *Interpreter) JITCompiler { 168 | return .{ 169 | .allocator = allocator, 170 | .interpreter = interpreter, 171 | .machine_code = std.ArrayList(u32).init(allocator), 172 | }; 173 | } 174 | 175 | pub fn deinit(self: *const JITCompiler) void { 176 | self.machine_code.deinit(); 177 | } 178 | 179 | fn allocJitBuf(nbytes: usize) [*]u32 { 180 | const buf: *anyopaque = mman.mmap( 181 | null, 182 | nbytes, 183 | mman.PROT_WRITE | mman.PROT_EXEC, 184 | mman.MAP_PRIVATE | mman.MAP_ANONYMOUS | mman.MAP_JIT, 185 | -1, 186 | 0, 187 | ) orelse unreachable; // TODO: return error 188 | 189 | if (buf == mman.MAP_FAILED) { 190 | std.debug.panic("mmap failed\n", .{}); // return error 191 | } 192 | 193 | return @ptrCast(@alignCast(buf)); 194 | } 195 | 196 | fn deallocJitBuf(buf: [*]u32, size: usize) void { 197 | if (mman.munmap(buf, size) != 0) { 198 | std.debug.panic("munmap failed\n", .{}); 199 | } 200 | } 201 | 202 | /// Take all the machine code instructions emitted so far, 203 | /// compile them into a function, then return the function. 204 | fn getCompiledFunction(self: *JITCompiler) CompiledFunction { 205 | const num_instructions = self.machine_code.items.len; 206 | const bufsize = num_instructions; 207 | const buf = allocJitBuf(bufsize); 208 | 209 | pthread.pthread_jit_write_protect_np(0); 210 | @memcpy(buf, self.machine_code.items); 211 | pthread.pthread_jit_write_protect_np(1); 212 | 213 | const func: JitFunction = @ptrCast(buf); 214 | return CompiledFunction.init(func, buf, bufsize); 215 | } 216 | 217 | inline fn emit(self: *JITCompiler, instr: u32) !void { 218 | try self.machine_code.append(instr); 219 | } 220 | 221 | fn emitPrelude(self: *JITCompiler) !void { 222 | // deref the stack pointer, store it in a local var 223 | try self.emit(Armv8a.ldrReg(VarReg.stackIndex, ArgReg.stackIndexPtr)); 224 | // deref the instruction pointer, store it in a local var 225 | try self.emit(Armv8a.ldrReg(VarReg.instrIndex, ArgReg.instrIndexPtr)); 226 | } 227 | 228 | fn emitEpilogue(self: *JITCompiler) !void { 229 | // Restore the stack and instruction pointers 230 | try self.emit(Armv8a.strReg(VarReg.stackIndex, ArgReg.stackIndexPtr, 0)); 231 | try self.emit(Armv8a.strReg(VarReg.instrIndex, ArgReg.instrIndexPtr, 0)); 232 | } 233 | 234 | fn emitPop(self: *JITCompiler) !void { 235 | try self.emit(Armv8a.subRegImm( 236 | VarReg.stackIndex, 237 | VarReg.stackIndex, 238 | 1, 239 | )); 240 | } 241 | 242 | fn emitPushReg(self: *JITCompiler, reg: Armv8a.Reg) !void { 243 | try self.emit(Armv8a.strRegScaled( 244 | reg, 245 | ArgReg.stackAddr, 246 | VarReg.stackIndex, 247 | )); 248 | 249 | try self.emit(Armv8a.addRegImm( 250 | VarReg.stackIndex, 251 | VarReg.stackIndex, 252 | 1, 253 | )); 254 | } 255 | 256 | inline fn emitReturn(self: *JITCompiler) !void { 257 | try self.emitEpilogue(); 258 | try self.emit(Armv8a.ret); 259 | } 260 | 261 | inline fn emitAdvanceIp(self: *JITCompiler) !void { 262 | try self.emit(Armv8a.addRegImm( 263 | VarReg.instrIndex, 264 | VarReg.instrIndex, 265 | 1, 266 | )); 267 | } 268 | 269 | inline fn readInstruction(self: *JITCompiler, dst_reg: Armv8a.Reg) !void { 270 | try self.emit(Armv8a.ldrByteReg( 271 | dst_reg, 272 | ArgReg.instructionsAddr, 273 | VarReg.instrIndex, 274 | )); 275 | 276 | try self.emitAdvanceIp(); 277 | } 278 | 279 | inline fn readStackTop(self: *JITCompiler, dst_reg: Armv8a.Reg) !void { 280 | try self.emit(Armv8a.ldrRegScaled( 281 | dst_reg, 282 | ArgReg.stackAddr, 283 | VarReg.stackIndex, 284 | )); 285 | } 286 | 287 | pub fn compileBlock(self: *JITCompiler, block_number: usize, block: *const CodeBlock) !CompiledFunction { 288 | try self.emitPrelude(); 289 | 290 | var i: usize = 0; 291 | while (i < block.instructions.len) { 292 | const instruction = block.instructions[i]; 293 | const op: Opcode = @enumFromInt(instruction); 294 | 295 | // ip += 1; 296 | try self.emit(Armv8a.addRegImm( 297 | VarReg.instrIndex, 298 | VarReg.instrIndex, 299 | 1, 300 | )); 301 | 302 | i += 1; 303 | 304 | switch (op) { 305 | .load_var => { 306 | // a = instructions[ip]; ip++; 307 | try self.readInstruction(VarReg.tempA); 308 | i += 1; 309 | 310 | try self.emit(Armv8a.ldrRegScaled( 311 | VarReg.tempB, 312 | ArgReg.stackAddr, 313 | VarReg.tempA, 314 | )); // b = stack[a] 315 | try self.emitPushReg(VarReg.tempB); // push(stack[a]); 316 | }, 317 | 318 | .eq => { 319 | try self.emitPop(); 320 | try self.readStackTop(VarReg.tempA); 321 | 322 | try self.emitPop(); 323 | try self.readStackTop(VarReg.tempB); 324 | 325 | try self.emit(Armv8a.cmpRegs(VarReg.tempA, VarReg.tempB)); 326 | try self.emit(Armv8a.setIfStatusEq(VarReg.tempC)); 327 | try self.emitPushReg(VarReg.tempC); 328 | }, 329 | 330 | .store_var => { 331 | // a = instructions[ip]; ip++; 332 | try self.readInstruction(VarReg.tempA); 333 | i += 1; 334 | 335 | // b = pop(); 336 | try self.emitPop(); 337 | try self.emit(Armv8a.ldrRegScaled( 338 | VarReg.tempB, 339 | ArgReg.stackAddr, 340 | VarReg.stackIndex, 341 | )); 342 | 343 | try self.emit(Armv8a.strRegScaled( 344 | VarReg.tempB, 345 | ArgReg.stackAddr, 346 | VarReg.tempA, 347 | )); // stack[a] = b; 348 | }, 349 | 350 | .push => { 351 | try self.readInstruction(VarReg.tempA); 352 | i += 1; 353 | 354 | try self.emit(Armv8a.ldrRegScaled( 355 | VarReg.tempB, 356 | ArgReg.constantsAddr, 357 | VarReg.tempA, 358 | )); // b = constants[a] 359 | 360 | try self.emitPushReg(VarReg.tempB); // push(p) 361 | }, 362 | 363 | .jump => { 364 | try self.readInstruction(VarReg.tempA); 365 | const dst_block = block.instructions[i]; 366 | i += 1; 367 | 368 | if (dst_block == block_number) { 369 | // we're jumping back the same block. 370 | try self.emit(Armv8a.movRegImm(VarReg.instrIndex, 0)); // reset ip 371 | const distance: i32 = @intCast(self.machine_code.items.len); 372 | try self.emit(Armv8a.branchIfEq(-distance)); // jump backwards. 373 | } else { 374 | try self.emit( 375 | Armv8a.strReg(VarReg.tempA, ArgReg.currentBlockNumber, 0), 376 | ); 377 | try self.emit(Armv8a.movRegImm(VarReg.instrIndex, 0)); // ip = 0 378 | try self.emitReturn(); 379 | } 380 | }, 381 | 382 | .jump_nz => { 383 | // a = pop(); 384 | try self.emitPop(); 385 | try self.readStackTop(VarReg.tempA); 386 | 387 | const dst_block = VarReg.tempB; 388 | 389 | // block_index = instructions[ip]; ip++; 390 | try self.readInstruction(dst_block); 391 | i += 1; 392 | 393 | // if (a == 0) 394 | try self.emit(Armv8a.cmpImmediate(VarReg.tempA, 0)); 395 | try self.emit(Armv8a.setIfStatusEq(VarReg.tempC)); 396 | 397 | try self.emit(0); // dummy instr, this will be patched below 398 | const jmp_instr_index = self.machine_code.items.len - 1; 399 | 400 | // current_block = block_index; 401 | try self.emit(Armv8a.strReg(dst_block, ArgReg.currentBlockNumber, 0)); 402 | try self.emit(Armv8a.movRegImm(VarReg.instrIndex, 0)); // ip = 0 403 | try self.emitReturn(); 404 | 405 | // patch the dummy jump instruction we emitted above 406 | const offset = self.machine_code.items.len - jmp_instr_index; 407 | self.machine_code.items[jmp_instr_index] = 408 | Armv8a.branchIfEq(@intCast(offset)); 409 | }, 410 | 411 | .add => { 412 | // A = pop() 413 | try self.emitPop(); 414 | try self.readStackTop(VarReg.tempA); 415 | 416 | // B = pop() 417 | try self.emitPop(); 418 | try self.readStackTop(VarReg.tempB); 419 | 420 | try self.emit(Armv8a.addRegs( 421 | VarReg.tempA, 422 | VarReg.tempA, 423 | VarReg.tempB, 424 | )); // a = a + b 425 | 426 | try self.emitPushReg(VarReg.tempA); 427 | }, 428 | } 429 | } 430 | 431 | try self.emitReturn(); 432 | return self.getCompiledFunction(); 433 | } 434 | }; 435 | 436 | test "ARMv8a code generation" { 437 | const ldr_x8_x2 = Armv8a.ldrReg(.x8, .x2); 438 | try std.testing.expectEqual(0xf9400048, ldr_x8_x2); 439 | // ldr x9, [x0, x8, lsl #3] 440 | const ldr_scaled_offset = Armv8a.ldrRegScaled(.x9, .x0, .x8); 441 | try std.testing.expectEqual(0xf8687809, ldr_scaled_offset); 442 | // ldr x11, [x0, x10] 443 | const ldr_unscaled_offset = Armv8a.ldrRegUnscaled(.x11, .x0, .x10); 444 | try std.testing.expectEqual(0xf86a680b, ldr_unscaled_offset); 445 | // add x12, x12, #1 446 | try std.testing.expectEqual(0x9100058c, Armv8a.addRegImm(.x12, .x12, 1)); 447 | // sub x8, x8, #1 448 | try std.testing.expectEqual(0xd1000508, Armv8a.subRegImm(.x8, .x8, 1)); 449 | // str x8, [x2] 450 | try std.testing.expectEqual(0xf9000048, Armv8a.strReg(.x8, .x2, 0)); 451 | // lsl x10, x8, #3 452 | try std.testing.expectEqual(0xd37df10a, Armv8a.mult8(.x10, .x8)); 453 | // add x9, x11, x9 454 | try std.testing.expectEqual(0x8b090169, Armv8a.addRegs(.x9, .x11, .x9)); 455 | // str x8, [x9, x10, lsl #3] 456 | try std.testing.expectEqual(0xf82a7928, Armv8a.strRegScaled(.x8, .x9, .x10)); 457 | // ldrb w8, [x1, x8] 458 | try std.testing.expectEqual(0x38686828, Armv8a.ldrByteReg(.x8, .x1, .x8)); 459 | // cmp x1, x2 460 | try std.testing.expectEqual(0xeb02003f, Armv8a.cmpRegs(.x1, .x2)); 461 | // cset x3, eq 462 | try std.testing.expectEqual(0x9A9F17E3, Armv8a.setIfStatusEq(.x3)); 463 | // mov x1, #9 464 | try std.testing.expectEqual(0xd2800121, Armv8a.movRegImm(.x1, 9)); 465 | // cmp x1, #1 466 | try std.testing.expectEqual(0xf100043f, Armv8a.cmpImmediate(.x1, 1)); 467 | // b.eq (ip + 2) 468 | try std.testing.expectEqual(0x54000040, Armv8a.branchIfEq(2)); 469 | // b.eq (ip - 3) 470 | try std.testing.expectEqual(0x54ffffa0, Armv8a.branchIfEq(-3)); 471 | } 472 | 473 | test "JITCompiler" { 474 | const allocator = std.testing.allocator; 475 | const program = [_]CodeBlock{.{ 476 | .constants = &[_]i64{ 30, 12 }, 477 | .instructions = &[_]u8{ 478 | Op(.add), 479 | Op(.push), // x = 30 480 | 0, 481 | Op(.load_var), // push(x) 482 | 0, 483 | Op(.push), 484 | 1, 485 | Op(.store_var), 486 | 0, 487 | 488 | Op(.push), 489 | 1, 490 | Op(.push), 491 | 0, 492 | Op(.eq), 493 | 494 | Op(.jump_nz), 495 | 33, 496 | 497 | Op(.push), 498 | 1, 499 | Op(.push), 500 | 0, 501 | Op(.eq), 502 | 503 | Op(.push), 504 | 1, 505 | Op(.push), 506 | 1, 507 | Op(.eq), 508 | 509 | Op(.jump_nz), 510 | 22, 511 | }, 512 | }}; 513 | 514 | var interpreter = try Interpreter.init(allocator, &program); 515 | defer { 516 | interpreter.deinit(); 517 | allocator.destroy(interpreter); 518 | } 519 | 520 | const compiled = try interpreter.jit_compiler.compileBlock(0, &program[0]); 521 | 522 | var stack = [_]i64{ 2, 3, 0, 0, 0, 0, 0, 0 }; 523 | var s_ptr: usize = 2; 524 | var i_ptr: usize = 0; 525 | var instructions = program[0].instructions; 526 | var current_block_index: usize = 0; 527 | 528 | compiled.func( 529 | (&stack).ptr, 530 | (&instructions).ptr, 531 | &s_ptr, 532 | &i_ptr, 533 | ¤t_block_index, 534 | program[0].constants.ptr, 535 | ); 536 | 537 | try std.testing.expectEqual(4, s_ptr); 538 | try std.testing.expectEqual(0, i_ptr); 539 | try std.testing.expectEqualSlices(i64, &[_]i64{ 12, 30, 5, 0 }, stack[0..s_ptr]); 540 | try std.testing.expectEqual(22, current_block_index); 541 | } 542 | --------------------------------------------------------------------------------