Flatten bitmap when there is only one element

We can avoid allocating a bitmap when the number of elements in the iseq
is fewer than the size of an iseq_bits_t

compile.c

@@ -2336,7 +2336,17 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
     // of the number of bits in an unsigned long.
 
     // Allocate enough room for the bitmask list
-    iseq_bits_t * mark_offset_bits = ZALLOC_N(iseq_bits_t, ISEQ_MBITS_BUFLEN(code_index));
+    iseq_bits_t * mark_offset_bits;
+    int code_size = code_index;
+
+    iseq_bits_t tmp[1] = {0};
+
+    if (ISEQ_MBITS_BUFLEN(code_index) == 1) {
+        mark_offset_bits = tmp;
+    }
+    else {
+        mark_offset_bits = ZALLOC_N(iseq_bits_t, ISEQ_MBITS_BUFLEN(code_index));
+    }
 
     list = FIRST_ELEMENT(anchor);
     insns_info_index = code_index = sp = 0;
@@ -2505,7 +2515,9 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
 			xfree(generated_iseq);
 			xfree(insns_info);
 			xfree(positions);
-			xfree(mark_offset_bits);
+                        if (ISEQ_MBITS_BUFLEN(code_size) > 1) {
+                            xfree(mark_offset_bits);
+                        }
 			debug_list(anchor, list);
 			COMPILE_ERROR(iseq, adjust->line_no,
 				      "iseq_set_sequence: adjust bug to %d %d < %d",
@@ -2525,7 +2537,13 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
     body->iseq_encoded = (void *)generated_iseq;
     body->iseq_size = code_index;
     body->stack_max = stack_max;
-    body->mark_offset_bits = mark_offset_bits;
+
+    if (ISEQ_MBITS_BUFLEN(body->iseq_size) == 1) {
+        body->mark_bits.single = mark_offset_bits[0];
+    }
+    else {
+        body->mark_bits.list = mark_offset_bits;
+    }
 
     /* get rid of memory leak when REALLOC failed */
     body->insns_info.body = insns_info;
@@ -11189,8 +11207,16 @@ ibf_load_code(const struct ibf_load *load, rb_iseq_t *iseq, ibf_offset_t bytecod
     struct rb_call_data *cd_entries = load_body->call_data;
     union iseq_inline_storage_entry *is_entries = load_body->is_entries;
 
-    iseq_bits_t * mark_offset_bits = ZALLOC_N(iseq_bits_t, ISEQ_MBITS_BUFLEN(iseq_size));
-    load_body->mark_offset_bits = mark_offset_bits;
+    iseq_bits_t * mark_offset_bits;
+
+    iseq_bits_t tmp[1] = {0};
+
+    if (ISEQ_MBITS_BUFLEN(iseq_size) == 1) {
+        mark_offset_bits = tmp;
+    }
+    else {
+        mark_offset_bits = ZALLOC_N(iseq_bits_t, ISEQ_MBITS_BUFLEN(iseq_size));
+    }
 
     unsigned int min_ic_index, min_ise_index, min_ivc_index;
     min_ic_index = min_ise_index = min_ivc_index = UINT_MAX;
@@ -11329,6 +11355,13 @@ ibf_load_code(const struct ibf_load *load, rb_iseq_t *iseq, ibf_offset_t bytecod
     load_body->iseq_encoded = code;
     load_body->iseq_size = code_index;
 
+    if (ISEQ_MBITS_BUFLEN(load_body->iseq_size) == 1) {
+        load_body->mark_bits.single = mark_offset_bits[0];
+    }
+    else {
+        load_body->mark_bits.list = mark_offset_bits;
+    }
+
     assert(code_index == iseq_size);
     assert(reading_pos == bytecode_offset + bytecode_size);
     return code;

iseq.c

@@ -193,7 +193,9 @@ rb_iseq_free(const rb_iseq_t *iseq)
 	}
 	ruby_xfree((void *)body->catch_table);
 	ruby_xfree((void *)body->param.opt_table);
-	ruby_xfree((void *)body->mark_offset_bits);
+        if (ISEQ_MBITS_BUFLEN(body->iseq_size) > 1) {
+            ruby_xfree((void *)body->mark_bits.list);
+        }
 
 	if (body->param.keyword != NULL) {
 	    ruby_xfree((void *)body->param.keyword->default_values);
@@ -313,6 +315,25 @@ iseq_extract_values(VALUE *code, size_t pos, iseq_value_itr_t * func, void *data
     return len;
 }
 
+static inline void
+iseq_scan_bits(unsigned int i, iseq_bits_t bits, VALUE *code, iseq_value_itr_t *func, void *data)
+{
+    unsigned int count = 0;
+
+    while(bits) {
+        if (bits & 0x1) {
+            unsigned int index = (i * ISEQ_MBITS_BITLENGTH) + count;
+            VALUE op = code[index];
+            VALUE newop = func(data, op);
+            if (newop != op) {
+                code[index] = newop;
+            }
+        }
+        bits >>= 1;
+        count++;
+    }
+}
+
 static void
 rb_iseq_each_value(const rb_iseq_t *iseq, iseq_value_itr_t * func, void *data)
 {
@@ -363,23 +384,13 @@ rb_iseq_each_value(const rb_iseq_t *iseq, iseq_value_itr_t * func, void *data)
     }
 
     // Embedded VALUEs
-    for (unsigned int i = 0; i < ISEQ_MBITS_BUFLEN(size); i++) {
-        iseq_bits_t bits = body->mark_offset_bits[i];
-        if (bits) {
-            unsigned int count = 0;
-
-            while(bits) {
-                if (bits & 0x1) {
-                    unsigned int index = (i * ISEQ_MBITS_BITLENGTH) + count;
-                    VALUE op = code[index];
-                    VALUE newop = func(data, op);
-                    if (newop != op) {
-                        code[index] = newop;
-                    }
-                }
-                bits >>= 1;
-                count++;
-            }
+    if (ISEQ_MBITS_BUFLEN(size) == 1) {
+        iseq_scan_bits(0, body->mark_bits.single, code, func, data);
+    }
+    else {
+        for (unsigned int i = 0; i < ISEQ_MBITS_BUFLEN(size); i++) {
+            iseq_bits_t bits = body->mark_bits.list[i];
+            iseq_scan_bits(i, bits, code, func, data);
         }
     }
 }

vm_core.h

@@ -465,7 +465,10 @@ struct rb_iseq_constant_body {
     unsigned int ivc_size; // Number of IVC and ICVARC caches
     unsigned int ci_size;
     unsigned int stack_max; /* for stack overflow check */
-    iseq_bits_t * mark_offset_bits; /* Find references for GC */
+    union {
+        iseq_bits_t * list; /* Find references for GC */
+        iseq_bits_t single;
+    } mark_bits;
 
     char catch_except_p; /* If a frame of this ISeq may catch exception, set TRUE */
     // If true, this ISeq is leaf *and* backtraces are not used, for example,