Store iseq node_id in a side array

The node id lived in each iseq_insn_info_entry, which prevented add_insn_info
from merging adjacent entries that only differ in node_id. Since most
neighboring instructions come from different AST nodes, this significantly
increases memory usage.

On Shopify's monolith, this change decreases boot heap memsize by 40MB.

compile.c

@@ -2630,14 +2630,8 @@ add_insn_info(struct iseq_insn_info_entry *insns_info, unsigned int *positions,
 {
     if (insns_info_index == 0 ||
         insns_info[insns_info_index-1].line_no != iobj->insn_info.line_no ||
-#ifdef USE_ISEQ_NODE_ID
-        insns_info[insns_info_index-1].node_id != iobj->insn_info.node_id ||
-#endif
         insns_info[insns_info_index-1].events  != iobj->insn_info.events) {
         insns_info[insns_info_index].line_no    = iobj->insn_info.line_no;
-#ifdef USE_ISEQ_NODE_ID
-        insns_info[insns_info_index].node_id    = iobj->insn_info.node_id;
-#endif
         insns_info[insns_info_index].events     = iobj->insn_info.events;
         positions[insns_info_index]             = code_index;
         return TRUE;
@@ -2650,7 +2644,6 @@ add_adjust_info(struct iseq_insn_info_entry *insns_info, unsigned int *positions
                 int insns_info_index, int code_index, const ADJUST *adjust)
 {
     insns_info[insns_info_index].line_no    = adjust->line_no;
-    insns_info[insns_info_index].node_id    = -1;
     insns_info[insns_info_index].events     = 0;
     positions[insns_info_index]             = code_index;
     return TRUE;
@@ -2777,6 +2770,10 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
 
     const int positions_size = insn_num;
     positions = ALLOC_N(unsigned int, insn_num);
+#ifdef USE_ISEQ_NODE_ID
+    int *node_ids = ALLOC_N(int, insn_num);
+    unsigned int node_ids_size = 0;
+#endif
     if (ISEQ_IS_SIZE(body)) {
         body->is_entries = ZALLOC_N(union iseq_inline_storage_entry, ISEQ_IS_SIZE(body));
     }
@@ -2954,6 +2951,9 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
                     }
                 }
                 if (add_insn_info(insns_info, positions, insns_info_index, code_index, iobj)) insns_info_index++;
+#ifdef USE_ISEQ_NODE_ID
+                node_ids[node_ids_size++] = iobj->insn_info.node_id;
+#endif
                 code_index += len;
                 break;
             }
@@ -2989,6 +2989,9 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
                                           "iseq_set_sequence: adjust bug (ISEQ_ELEMENT_ADJUST must not be the first in iseq)");
                         }
                         if (add_adjust_info(insns_info, positions, insns_info_index, code_index, adjust)) insns_info_index++;
+#ifdef USE_ISEQ_NODE_ID
+                        node_ids[node_ids_size++] = -1;
+#endif
                     }
                     if (diff > 1) {
                         generated_iseq[code_index++] = BIN(adjuststack);
@@ -3002,6 +3005,9 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
                         SIZED_FREE_N(generated_iseq, generated_iseq_size);
                         SIZED_FREE_N(insns_info, insns_info_size);
                         SIZED_FREE_N(positions, positions_size);
+#ifdef USE_ISEQ_NODE_ID
+                        SIZED_FREE_N(node_ids, insns_info_size);
+#endif
                         if (ISEQ_MBITS_BUFLEN(code_size) > 1) {
                             SIZED_FREE_N(mark_offset_bits, ISEQ_MBITS_BUFLEN(code_index));
                         }
@@ -3042,12 +3048,19 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
     /* get rid of memory leak when REALLOC failed */
     body->insns_info.body = insns_info;
     body->insns_info.positions_or_succ_index_table.positions = positions;
+#ifdef USE_ISEQ_NODE_ID
+    body->insns_info.node_ids = node_ids;
+#endif
 
     SIZED_REALLOC_N(insns_info, struct iseq_insn_info_entry, insns_info_index, insns_info_size);
     body->insns_info.body = insns_info;
     SIZED_REALLOC_N(positions, unsigned int, insns_info_index, positions_size);
     body->insns_info.positions_or_succ_index_table.positions = positions;
     body->insns_info.size = insns_info_index;
+#ifdef USE_ISEQ_NODE_ID
+    body->insns_info.node_ids = node_ids;
+    body->insns_info.node_ids_size = node_ids_size;
+#endif
 
     return COMPILE_OK;
 }
@@ -12787,7 +12800,7 @@ typedef uint32_t ibf_offset_t;
 
 #define IBF_MAJOR_VERSION ISEQ_MAJOR_VERSION
 #ifdef RUBY_DEVEL
-#define IBF_DEVEL_VERSION 9
+#define IBF_DEVEL_VERSION 10
 #define IBF_MINOR_VERSION (ISEQ_MINOR_VERSION * 10000 + IBF_DEVEL_VERSION)
 #else
 #define IBF_MINOR_VERSION ISEQ_MINOR_VERSION
@@ -13502,36 +13515,53 @@ static ibf_offset_t
 ibf_dump_insns_info_body(struct ibf_dump *dump, const rb_iseq_t *iseq)
 {
     ibf_offset_t offset = ibf_dump_pos(dump);
-    const struct iseq_insn_info_entry *entries = ISEQ_BODY(iseq)->insns_info.body;
+    const struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
+    const struct iseq_insn_info_entry *entries = body->insns_info.body;
 
-    unsigned int i;
-    for (i = 0; i < ISEQ_BODY(iseq)->insns_info.size; i++) {
-        ibf_dump_write_small_value(dump, entries[i].line_no);
 #ifdef USE_ISEQ_NODE_ID
-        ibf_dump_write_small_value(dump, entries[i].node_id);
+    ibf_dump_write_small_value(dump, body->insns_info.node_ids_size);
 #endif
+    unsigned int i;
+    for (i = 0; i < body->insns_info.size; i++) {
+        ibf_dump_write_small_value(dump, entries[i].line_no);
         ibf_dump_write_small_value(dump, entries[i].events);
     }
+#ifdef USE_ISEQ_NODE_ID
+    if (body->insns_info.node_ids_size > 0) {
+        ibf_dump_write(dump, body->insns_info.node_ids, sizeof(int) * body->insns_info.node_ids_size);
+    }
+#endif
 
     return offset;
 }
 
-static struct iseq_insn_info_entry *
-ibf_load_insns_info_body(const struct ibf_load *load, ibf_offset_t body_offset, unsigned int size)
+static void
+ibf_load_insns_info_body(const struct ibf_load *load, ibf_offset_t body_offset, unsigned int size,
+                         struct rb_iseq_constant_body *load_body)
 {
     ibf_offset_t reading_pos = body_offset;
     struct iseq_insn_info_entry *entries = ALLOC_N(struct iseq_insn_info_entry, size);
 
+#ifdef USE_ISEQ_NODE_ID
+    const unsigned int node_ids_size = (unsigned int)ibf_load_small_value(load, &reading_pos);
+#endif
     unsigned int i;
     for (i = 0; i < size; i++) {
         entries[i].line_no = (int)ibf_load_small_value(load, &reading_pos);
-#ifdef USE_ISEQ_NODE_ID
-        entries[i].node_id = (int)ibf_load_small_value(load, &reading_pos);
-#endif
         entries[i].events = (rb_event_flag_t)ibf_load_small_value(load, &reading_pos);
     }
-
-    return entries;
+#ifdef USE_ISEQ_NODE_ID
+    if (node_ids_size > 0) {
+        if ((size_t)node_ids_size > (load->current_buffer->size - reading_pos) / sizeof(int)) {
+            rb_raise(rb_eRuntimeError, "invalid byte sequence");
+        }
+        int *node_ids = ALLOC_N(int, node_ids_size);
+        memcpy(node_ids, load->current_buffer->buff + reading_pos, sizeof(int) * node_ids_size);
+        load_body->insns_info.node_ids = node_ids;
+        load_body->insns_info.node_ids_size = node_ids_size;
+    }
+#endif
+    load_body->insns_info.body = entries;
 }
 
 static ibf_offset_t
@@ -14222,7 +14252,7 @@ ibf_load_iseq_each(struct ibf_load *load, rb_iseq_t *iseq, ibf_offset_t offset)
     load_body->param.opt_table      = ibf_load_param_opt_table(load, param_opt_table_offset, param_opt_num);
     load_body->param.keyword        = ibf_load_param_keyword(load, param_keyword_offset);
     load_body->param.flags.has_kw   = (param_flags >> 4) & 1;
-    load_body->insns_info.body      = ibf_load_insns_info_body(load, insns_info_body_offset, insns_info_size);
+                                      ibf_load_insns_info_body(load, insns_info_body_offset, insns_info_size, load_body);
     load_body->insns_info.positions_or_succ_index_table.positions = ibf_load_insns_info_positions(load, insns_info_positions_offset, insns_info_size);
     load_body->local_table          = ibf_load_local_table(load, local_table_offset, local_table_size);
     ibf_load_lvar_states(load, load_body, lvar_states_offset, local_table_size, load_body->local_table);

iseq.c

@@ -204,6 +204,9 @@ rb_iseq_free(const rb_iseq_t *iseq)
 #endif
         SIZED_FREE_N(body->iseq_encoded, body->iseq_size);
         SIZED_FREE_N(body->insns_info.body, body->insns_info.size);
+#ifdef USE_ISEQ_NODE_ID
+        SIZED_FREE_N(body->insns_info.node_ids, body->insns_info.node_ids_size);
+#endif
 #if VM_INSN_INFO_TABLE_IMPL == 2
         ruby_xfree(body->insns_info.positions_or_succ_index_table.succ_index_table);
 #else
@@ -593,6 +596,9 @@ rb_iseq_memsize(const rb_iseq_t *iseq)
         if (body->variable) size += sizeof(struct rb_iseq_variable);
         size += body->iseq_size * sizeof(VALUE);
         size += body->insns_info.size * (sizeof(struct iseq_insn_info_entry) + sizeof(unsigned int));
+#ifdef USE_ISEQ_NODE_ID
+        size += body->insns_info.node_ids_size * sizeof(int);
+#endif
         size += body->local_table_size * sizeof(ID); // body->local_table
         if (!ISEQ_LVAR_STATES_EMBED_P(body->local_table_size) && body->lvar_states.list) {
             size += ISEQ_LVAR_STATES_BUFLEN(body->local_table_size) * sizeof(uint8_t);
@@ -2575,14 +2581,27 @@ rb_iseq_line_no(const rb_iseq_t *iseq, size_t pos)
 int
 rb_iseq_node_id(const rb_iseq_t *iseq, size_t pos)
 {
-    const struct iseq_insn_info_entry *entry = get_insn_info(iseq, pos);
-
-    if (entry) {
-        return entry->node_id;
-    }
-    else {
+    const struct rb_iseq_constant_body *const body = ISEQ_BODY(iseq);
+    if (body->insns_info.node_ids_size == 0) {
         return 0;
     }
+
+    const VALUE *code = rb_iseq_original_iseq(iseq);
+    const int *node_ids = body->insns_info.node_ids;
+    unsigned int i = 0;
+    size_t code_pos = 0;
+
+    while (i < body->insns_info.node_ids_size && code_pos < body->iseq_size) {
+        size_t len = (size_t)insn_len(code[code_pos]);
+        if (pos < code_pos + len) {
+            return node_ids[i];
+        }
+        code_pos += len;
+        i++;
+    }
+
+    VM_ASSERT(0); /* node_ids out of sync with the instruction sequence */
+    return -1;
 }
 #endif
 
@@ -3777,7 +3796,9 @@ iseq_data_to_ary(const rb_iseq_t *iseq)
 
         info = get_insn_info(iseq, pos);
 #ifdef USE_ISEQ_NODE_ID
-        rb_ary_push(node_ids, INT2FIX(info->node_id));
+        VM_ASSERT(l < (long)iseq_body->insns_info.node_ids_size);
+        rb_ary_push(node_ids, INT2FIX(l < (long)iseq_body->insns_info.node_ids_size ?
+                                      iseq_body->insns_info.node_ids[l] : -1));
 #endif
 
         if (prev_insn_info != info) {

iseq.h

@@ -46,10 +46,6 @@ iseq_lvar_state_set(uint8_t *buf, unsigned int i, enum lvar_state state)
     *byte = (*byte & ~(((1 << ISEQ_LVAR_STATE_BITS) - 1) << shift)) | ((uint8_t)state << shift);
 }
 
-#ifndef USE_ISEQ_NODE_ID
-#define USE_ISEQ_NODE_ID 1
-#endif
-
 #ifndef rb_iseq_t
 typedef struct rb_iseq_struct rb_iseq_t;
 #define rb_iseq_t rb_iseq_t
@@ -337,9 +333,6 @@ struct rb_compile_option_struct {
 
 struct iseq_insn_info_entry {
     int line_no;
-#ifdef USE_ISEQ_NODE_ID
-    int node_id;
-#endif
     rb_event_flag_t events;
 };
 

vm_core.h

@@ -140,6 +140,10 @@ If `rb_vm_check_ints()` is called between the `RUBY_ASSERT_CRITICAL_SECTION_ENTE
 # define VM_INSN_INFO_TABLE_IMPL 2
 #endif
 
+#ifndef USE_ISEQ_NODE_ID
+#define USE_ISEQ_NODE_ID 1
+#endif
+
 #if defined(NSIG_MAX)           /* POSIX issue 8 */
 # undef NSIG
 # define NSIG NSIG_MAX
@@ -528,6 +532,10 @@ struct rb_iseq_constant_body {
 #endif
         } positions_or_succ_index_table;
         unsigned int size;
+#ifdef USE_ISEQ_NODE_ID
+        const int *node_ids;
+        unsigned int node_ids_size;
+#endif
     } insns_info;
 
     const ID *local_table;		/* must free */

zjit/src/cruby_bindings.inc.rs

@@ -1276,6 +1276,8 @@ pub struct rb_iseq_constant_body_iseq_insn_info {
     pub body: *const iseq_insn_info_entry,
     pub positions_or_succ_index_table: rb_iseq_constant_body_iseq_insn_info__bindgen_ty_1,
     pub size: ::std::os::raw::c_uint,
+    pub node_ids: *const ::std::os::raw::c_int,
+    pub node_ids_size: ::std::os::raw::c_uint,
 }
 #[repr(C)]
 #[derive(Copy, Clone)]
@@ -2050,7 +2052,6 @@ pub struct iseq_compile_data__bindgen_ty_3 {
 #[derive(Debug, Copy, Clone)]
 pub struct iseq_insn_info_entry {
     pub line_no: ::std::os::raw::c_int,
-    pub node_id: ::std::os::raw::c_int,
     pub events: rb_event_flag_t,
 }
 pub const CATCH_TYPE_RESCUE: rb_catch_type = 3;
@@ -2114,7 +2115,7 @@ pub type rb_iseq_callback = ::std::option::Option<
     unsafe extern "C" fn(arg1: *const rb_iseq_t, arg2: *mut ::std::os::raw::c_void),
 >;
 pub const ISEQ_BODY_OFFSET_PARAM: zjit_struct_offsets = 16;
-pub const ISEQ_BODY_OFFSET_OUTER_VARIABLES: zjit_struct_offsets = 240;
+pub const ISEQ_BODY_OFFSET_OUTER_VARIABLES: zjit_struct_offsets = 256;
 pub const RUBY_OFFSET_THREAD_RACTOR: zjit_struct_offsets = 24;
 pub type zjit_struct_offsets = u32;
 #[repr(C)]