feat: Tracing JIT

- Profile loops: super simple right now, a simple counter
- When hits threshold, compile the node into a native function with no arguments
- Patch the bytecode to replace the loop with the function call

closes #134

src/Chunk.zig

@@ -118,14 +118,17 @@ pub const OpCode = enum(u8) {
     OP_LIST_APPEND,
 
     OP_MAP,
-    // FIXMEL delete and only use OP_SET_MAP_SUBSCRIPT
+    // FIXME: delete and only use OP_SET_MAP_SUBSCRIPT
     OP_SET_MAP,
 
     OP_EXPORT,
     OP_IMPORT,
 
     OP_TO_STRING,
     OP_TYPEOF,
+
+    OP_JIT,
+    OP_JIT_END,
 };
 
 /// A chunk of code to execute

src/Codegen.zig

@@ -57,8 +57,8 @@ jit: ?*JIT,
 
 reporter: Reporter,
 
-const generators = [_]NodeGen{
-    noGen, // AnonymousObjectType,
+const generators = [_]?NodeGen{
+    null, // AnonymousObjectType,
     generateAs, // As,
     generateAsyncCall, // AsyncCall,
     generateBinary, // Binary,
@@ -73,27 +73,27 @@ const generators = [_]NodeGen{
     generateEnum, // Enum,
     generateExport, // Export,
     generateExpression, // Expression,
-    noGen, // FiberType,
+    null, // FiberType,
     generateFloat, // Float,
     generateFor, // For,
     generateForceUnwrap, // ForceUnwrap,
     generateForEach, // ForEach,
     generateFunction, // Function,
-    noGen, // FunctionType,
+    null, // FunctionType,
     generateFunDeclaration, // FunDeclaration,
     generateGenericResolve, // GenericResolve,
-    noGen, // GenericResolveType,
-    noGen, // GenericType,
+    null, // GenericResolveType,
+    null, // GenericType,
     generateGrouping, // Grouping,
     generateIf, // If,
     generateImport, // Import,
     generateInteger, // Integer,
     generateIs, // Is,
     generateList, // List,
-    noGen, // ListType,
+    null, // ListType,
     generateMap, // Map,
-    noGen, // MapType,
-    noGen, // Namespace,
+    null, // MapType,
+    null, // Namespace,
     generateNamedVariable, // NamedVariable,
     generateNull, // Null,
     generateObjectDeclaration, // ObjectDeclaration,
@@ -105,7 +105,7 @@ const generators = [_]NodeGen{
     generateResolve, // Resolve,
     generateResume, // Resume,
     generateReturn, // Return,
-    noGen, // SimpleType,
+    null, // SimpleType,
     generateString, // String,
     generateStringLiteral, // StringLiteral,
     generateSubscript, // Subscript,
@@ -115,7 +115,7 @@ const generators = [_]NodeGen{
     generateTypeOfExpression, // TypeOfExpression,
     generateUnary, // Unary,
     generateUnwrap, // Unwrap,
-    noGen, // UserType,
+    null, // UserType,
     generateVarDeclaration, // VarDeclaration,
     generateVoid, // Void,
     generateWhile, // While,
@@ -408,7 +408,11 @@ inline fn generateNode(self: *Self, node: Ast.Node.Index, breaks: ?*std.ArrayLis
         return null;
     }
 
-    return try Self.generators[@intFromEnum(self.ast.nodes.items(.tag)[node])](self, node, breaks);
+    if (Self.generators[@intFromEnum(self.ast.nodes.items(.tag)[node])]) |generator| {
+        return generator(self, node, breaks);
+    }
+
+    return null;
 }
 
 fn nodeValue(self: *Self, node: Ast.Node.Index) Error!?Value {

src/Jit.zig

@@ -74,10 +74,14 @@ const Self = @This();
 vm: *VM,
 ctx: m.MIR_context_t,
 state: ?GenState = null,
-// List of closures being or already compiled
+// Set of closures being or already compiled
 compiled_closures: std.AutoHashMap(*o.ObjClosure, void),
-// Closure we can't compile (containing async call, or yield)
+// Closures we can't compile (containing async call, or yield)
 blacklisted_closures: std.AutoHashMap(*o.ObjClosure, void),
+// Set of hostpot being already compiled
+compiled_hotspots: std.AutoHashMap(Ast.Node.Index, void),
+// Hotspots we can't comiple (containing async call, or yield)
+blacklisted_hotspots: std.AutoHashMap(Ast.Node.Index, void),
 // MIR doesn't allow generating multiple functions at once, so we keep a set of function to compile
 // Once compiled, the value is set to an array of the native and raw native func_items
 functions_queue: std.AutoHashMap(Ast.Node.Index, ?[2]m.MIR_item_t),
@@ -97,7 +101,9 @@ pub fn init(vm: *VM) Self {
         .vm = vm,
         .ctx = m.MIR_init(),
         .compiled_closures = std.AutoHashMap(*o.ObjClosure, void).init(vm.gc.allocator),
+        .compiled_hotspots = std.AutoHashMap(Ast.Node.Index, void).init(vm.gc.allocator),
         .blacklisted_closures = std.AutoHashMap(*o.ObjClosure, void).init(vm.gc.allocator),
+        .blacklisted_hotspots = std.AutoHashMap(Ast.Node.Index, void).init(vm.gc.allocator),
         .functions_queue = std.AutoHashMap(Ast.Node.Index, ?[2]m.MIR_item_t).init(vm.gc.allocator),
         .objclosures_queue = std.AutoHashMap(*o.ObjClosure, void).init(vm.gc.allocator),
         .required_ext_api = std.AutoHashMap(ExternApi, void).init(vm.gc.allocator),
@@ -107,7 +113,9 @@ pub fn init(vm: *VM) Self {
 
 pub fn deinit(self: *Self) void {
     self.compiled_closures.deinit();
+    self.compiled_hotspots.deinit();
     self.blacklisted_closures.deinit();
+    self.blacklisted_hotspots.deinit();
     // std.debug.assert(self.functions_queue.count() == 0);
     self.functions_queue.deinit();
     // std.debug.assert(self.objclosures_queue.count() == 0);
@@ -138,28 +146,64 @@ pub fn compileFunction(self: *Self, ast: Ast, closure: *o.ObjClosure) Error!void
     // Build the function
     try self.buildFunction(ast, closure, function_node);
 
-    // Did we encountered other functions that need to be compiled?
-    var it = self.functions_queue.iterator();
-    while (it.next()) |kv| {
+    // Did we encounter other functions to compile?
+    try self.buildCollateralFunction(ast);
+
+    // Load modules
+    for (self.modules.items) |module| {
+        m.MIR_load_module(self.ctx, module);
+    }
+
+    // Load external functions
+    var it_ext = self.required_ext_api.iterator();
+    while (it_ext.next()) |kv| {
+        switch (kv.key_ptr.*) {
+            // TODO: don't mix those with actual api functions
+            .rawfn, .nativefn => {},
+            else => m.MIR_load_external(
+                self.ctx,
+                kv.key_ptr.*.name(),
+                kv.key_ptr.*.ptr(),
+            ),
+        }
+    }
+
+    // Link everything together
+    m.MIR_link(self.ctx, m.MIR_set_lazy_gen_interface, null);
+
+    m.MIR_gen_init(self.ctx, 1);
+    defer m.MIR_gen_finish(self.ctx);
+
+    // Generate all needed functions and set them in corresponding ObjFunctions
+    var it2 = self.functions_queue.iterator();
+    while (it2.next()) |kv| {
         const node = kv.key_ptr.*;
+        const items = kv.value_ptr.*.?;
 
-        if (kv.value_ptr.* == null) {
-            // Does it have an associated closure?
-            var it2 = self.objclosures_queue.iterator();
-            var sub_closure: ?*o.ObjClosure = null;
-            while (it2.next()) |kv2| {
-                if (kv2.key_ptr.*.function.node == node) {
-                    sub_closure = kv2.key_ptr.*;
-                    break;
-                }
-            }
-            try self.buildFunction(ast, sub_closure, node);
+        const native = m.MIR_gen(self.ctx, 0, items[0]);
+        const native_raw = m.MIR_gen(self.ctx, 0, items[1]);
 
-            // Building a new function might have added functions in the queue, so we reset the iterator
-            it = self.functions_queue.iterator();
+        // Find out if we need to set it in a ObjFunction
+        var it3 = self.objclosures_queue.iterator();
+        while (it3.next()) |kv2| {
+            if (kv2.key_ptr.*.function.node == node) {
+                kv2.key_ptr.*.function.native = native;
+                kv2.key_ptr.*.function.native_raw = native_raw;
+                break;
+            }
         }
     }
 
+    self.reset();
+}
+
+pub fn compileHotSpot(self: *Self, ast: Ast, hotspot_node: Ast.Node.Index) Error!*anyopaque {
+    // Build function surrounding the node
+    try self.buildFunction(ast, null, hotspot_node);
+
+    // Did we encounter other functions to compile?
+    try self.buildCollateralFunctions(ast);
+
     // Load modules
     for (self.modules.items) |module| {
         m.MIR_load_module(self.ctx, module);
@@ -187,6 +231,7 @@ pub fn compileFunction(self: *Self, ast: Ast, closure: *o.ObjClosure) Error!void
 
     // Generate all needed functions and set them in corresponding ObjFunctions
     var it2 = self.functions_queue.iterator();
+    var hotspot_native: ?*anyopaque = null;
     while (it2.next()) |kv| {
         const node = kv.key_ptr.*;
         const items = kv.value_ptr.*.?;
@@ -203,9 +248,39 @@ pub fn compileFunction(self: *Self, ast: Ast, closure: *o.ObjClosure) Error!void
                 break;
             }
         }
+
+        // If its the hotspot, return the NativeFn pointer
+        if (node == hotspot_node) {
+            hotspot_native = native;
+        }
     }
 
     self.reset();
+
+    return hotspot_native orelse Error.CantCompile;
+}
+
+fn buildCollateralFunctions(self: *Self, ast: Ast) Error!void {
+    var it = self.functions_queue.iterator();
+    while (it.next()) |kv| {
+        const node = kv.key_ptr.*;
+
+        if (kv.value_ptr.* == null) {
+            // Does it have an associated closure?
+            var it2 = self.objclosures_queue.iterator();
+            var sub_closure: ?*o.ObjClosure = null;
+            while (it2.next()) |kv2| {
+                if (kv2.key_ptr.*.function.node == node) {
+                    sub_closure = kv2.key_ptr.*;
+                    break;
+                }
+            }
+            try self.buildFunction(ast, sub_closure, node);
+
+            // Building a new function might have added functions in the queue, so we reset the iterator
+            it = self.functions_queue.iterator();
+        }
+    }
 }
 
 fn buildFunction(self: *Self, ast: Ast, closure: ?*o.ObjClosure, function_node: Ast.Node.Index) Error!void {

src/disassembler.zig

@@ -249,9 +249,10 @@ pub fn disassembleInstruction(chunk: *Chunk, offset: usize) !usize {
         .OP_TRY_END,
         .OP_GET_ENUM_CASE_FROM_VALUE,
         .OP_TYPEOF,
+        .OP_JIT_END,
         => simpleInstruction(instruction, offset),
 
-        .OP_SWAP => bytesInstruction(instruction, chunk, offset),
+        .OP_SWAP, .OP_JIT => bytesInstruction(instruction, chunk, offset),
 
         .OP_DEFINE_GLOBAL,
         .OP_GET_GLOBAL,