fix(jit): Don't use hashmap to keep count of hotspots

src/Ast.zig

@@ -71,6 +71,9 @@ pub const Node = struct {
     /// To avoid generating a node const value multiple times
     value: ?Value = null,
 
+    /// How many time it was visited at runtime (used to decide wether its a hotspot that needs to be compiled)
+    count: usize = 0,
+
     pub const Index = u32;
 
     pub const Tag = enum(u8) {

src/Codegen.zig

@@ -53,9 +53,9 @@ current: ?*Frame = null,
 ast: Ast,
 gc: *GarbageCollector,
 flavor: RunFlavor,
-// Jump to patch at end of current expression with a optional unwrapping in the middle of it
+/// Jump to patch at end of current expression with a optional unwrapping in the middle of it
 opt_jumps: ?std.ArrayList(usize) = null,
-// Used to generate error messages
+/// Used to generate error messages
 parser: *Parser,
 jit: ?*JIT,
 
@@ -152,7 +152,7 @@ pub inline fn currentCode(self: *Self) usize {
     return self.current.?.function.?.chunk.code.items.len;
 }
 
-pub fn generate(self: *Self, ast: Ast) anyerror!?*obj.ObjFunction {
+pub fn generate(self: *Self, ast: Ast) Error!?*obj.ObjFunction {
     self.ast = ast;
     self.reporter.had_error = false;
     self.reporter.panic_mode = false;

src/Jit.zig

@@ -11,7 +11,6 @@ const ZigType = @import("zigtypes.zig").Type;
 const ExternApi = @import("jit_extern_api.zig").ExternApi;
 const api = @import("lib/buzz_api.zig");
 const io = @import("io.zig");
-const disassembler = @import("disassembler.zig");
 
 pub const Error = error{
     CantCompile,
@@ -41,7 +40,7 @@ const GenState = struct {
     module: m.MIR_module_t,
     prototypes: std.AutoHashMap(ExternApi, m.MIR_item_t),
 
-    // Root closure (not necessarily the one being compiled)
+    /// Root closure (not necessarily the one being compiled)
     closure: *o.ObjClosure,
     opt_jump: ?OptJump = null,
 
@@ -57,16 +56,16 @@ const GenState = struct {
     function_native: ?m.MIR_item_t = null,
     function_native_proto: ?m.MIR_item_t = null,
 
-    // Convenience registers
+    /// Convenience registers
     ctx_reg: ?m.MIR_reg_t = null,
     vm_reg: ?m.MIR_reg_t = null,
 
-    // Avoid register name collisions
+    /// Avoid register name collisions
     registers: std.AutoHashMap([*:0]const u8, usize),
 
-    // Label to jump to when breaking a loop without a label
+    /// Label to jump to when breaking a loop without a label
     break_label: m.MIR_insn_t = null,
-    // Label to jump to when continuing a loop whithout a label
+    /// Label to jump to when continuing a loop whithout a label
     continue_label: m.MIR_insn_t = null,
 
     breaks_label: Breaks,

src/Parser.zig

@@ -121,8 +121,6 @@ const crypto_api = if (is_wasm) std.StaticStringMap(buzz_api.NativeFn).initCompt
     },
 ) else void;
 
-// TODO: other libs
-
 const Self = @This();
 
 extern fn dlerror() [*:0]u8;
@@ -229,10 +227,10 @@ pub const Frame = struct {
     upvalues: [255]UpValue,
     upvalue_count: u8 = 0,
     scope_depth: u32 = 0,
-    // Keep track of the node that introduced the scope (useful for labeled break/continue statements)
+    /// Keep track of the node that introduced the scope (useful for labeled break/continue statements)
     scopes: std.ArrayList(?Ast.Node.Index),
-    // If false, `return` was omitted or within a conditionned block (if, loop, etc.)
-    // We only count `return` emitted within the scope_depth 0 of the current function or unconditionned else statement
+    /// If false, `return` was omitted or within a conditionned block (if, loop, etc.)
+    /// We only count `return` emitted within the scope_depth 0 of the current function or unconditionned else statement
     function_node: Ast.Node.Index,
     function: ?*obj.ObjFunction = null,
     generics: ?*std.AutoArrayHashMap(*obj.ObjString, *obj.ObjTypeDef) = null,
@@ -480,13 +478,13 @@ gc: *GarbageCollector,
 scanner: ?Scanner = null,
 current_token: ?Ast.TokenIndex = null,
 script_name: []const u8 = undefined,
-// If true the script is being imported
+/// If true the script is being imported
 imported: bool = false,
-// True when parsing a declaration inside an export statement
+/// True when parsing a declaration inside an export statement
 exporting: bool = false,
-// Cached imported functions (shared across instances of Parser)
+/// Cached imported functions (shared across instances of Parser)
 imports: *std.StringHashMap(ScriptImport),
-// Keep track of things imported by the current script
+/// Keep track of things imported by the current script
 script_imports: std.StringHashMap(LocalScriptImport),
 test_count: u64 = 0,
 // FIXME: use SinglyLinkedList instead of heap allocated ptrs
@@ -498,7 +496,7 @@ flavor: RunFlavor,
 ffi: FFI,
 reporter: Reporter,
 
-// Jump to patch at end of current expression with a optional unwrapping in the middle of it
+/// Jump to patch at end of current expression with a optional unwrapping in the middle of it
 opt_jumps: ?std.ArrayList(Precedence) = null,
 
 pub fn init(

src/vm.zig

@@ -450,7 +450,6 @@ pub const VM = struct {
     globals_count: usize = 0,
     import_registry: *ImportRegistry,
     jit: ?JIT = null,
-    hotspots: std.AutoHashMap(Ast.Node.Index, usize),
     hotspots_count: u128 = 0,
     flavor: RunFlavor,
     reporter: Reporter,
@@ -469,15 +468,13 @@ pub const VM = struct {
                 .allocator = gc.allocator,
             },
             .ffi = if (!is_wasm) FFI.init(gc) else {},
-            .hotspots = std.AutoHashMap(Ast.Node.Index, usize).init(gc.allocator),
         };
     }
 
     pub fn deinit(self: *Self) void {
         // TODO: we can't free this because exported closure refer to it
         // self.globals.deinit();
         self.ffi.deinit();
-        self.hotspots.deinit();
     }
 
     pub fn cliArgs(self: *Self, args: ?[][:0]u8) !*ObjList {
@@ -4276,17 +4273,8 @@ pub const VM = struct {
         if (is_wasm) unreachable;
 
         const node = self.readInstruction();
-        const count = (self.hotspots.get(node) orelse 0) + 1;
 
-        self.hotspots.put(
-            node,
-            count,
-        ) catch {
-            self.panic("Out of memory");
-            unreachable;
-        };
-
-        self.hotspots_count += 1;
+        self.current_ast.nodes.items(.count)[node] += 1;
 
         if (self.shouldCompileHotspot(node)) {
             var timer = std.time.Timer.start() catch unreachable;
@@ -4672,7 +4660,7 @@ pub const VM = struct {
                 io.print(
                     "Calling compiled version of function `{s}.{}.n{}`\n",
                     .{
-                        closure.function.name.string,
+                        closure.function.type_def.resolved_type.?.Function.name.string,
                         self.current_ast.nodes.items(.components)[closure.function.node].Function.id,
                         closure.function.node,
                     },
@@ -4755,7 +4743,7 @@ pub const VM = struct {
             io.print(
                 "Calling uncompiled version of function `{s}.{}.n{}`\n",
                 .{
-                    closure.function.name.string,
+                    closure.function.type_def.resolved_type.?.Function.name.string,
                     self.current_ast.nodes.items(.components)[closure.function.node].Function.id,
                     closure.function.node,
                 },
@@ -5155,20 +5143,19 @@ pub const VM = struct {
     }
 
     fn shouldCompileHotspot(self: *Self, node: Ast.Node.Index) bool {
-        if (!self.current_ast.nodes.items(.tag)[node].isHotspot()) {
-            return false;
-        }
-
-        if (self.jit != null and (self.jit.?.compiled_nodes.get(node) != null or self.jit.?.blacklisted_nodes.get(node) != null)) {
-            return false;
-        }
-
-        const count = self.hotspots.get(node) orelse 0;
+        const count = self.current_ast.nodes.items(.count)[node];
 
-        return BuildOptions.jit_always_on or
-            BuildOptions.jit_hotspot_always_on or
+        // JIT is on?
+        return self.jit != null and
+            // JIT compile all the thing?
+            (BuildOptions.jit_always_on or BuildOptions.jit_hotspot_always_on or
+            // Threshold reached
             (count > 10 and
-            (@as(f128, @floatFromInt(count)) / @as(f128, @floatFromInt(self.hotspots_count))) > BuildOptions.jit_prof_threshold);
+            (@as(f128, @floatFromInt(count)) / @as(f128, @floatFromInt(self.hotspots_count))) > BuildOptions.jit_prof_threshold)) and
+            // It qualifies has a hotspot
+            self.current_ast.nodes.items(.tag)[node].isHotspot() and
+            // It's not already done or blacklisted
+            (self.jit.?.compiled_nodes.get(node) == null and self.jit.?.blacklisted_nodes.get(node) == null);
     }
 
     fn patchHotspot(