Replace `tempirep` with `RProc`

Previously I used the `RData` object to avoid a memory leak in `mrb_irep` if `src/load.c` failed.
ref: https://github.com/mruby/mruby/pull/4250
commit: f1523d24042ca3416dc5b9be7b3fc220ddaed896

Considering that the `RProc` object will be created in the subsequent process, it is preferable to create the `RProc` object from the beginning.
Along with this, the inside of `read_irep()` is replaced with the processing centered on the `RProc` object.

The global function that returns the `mrb_irep` pointer is still provided for compatibility.

2 files changed, 82 insertions, 77 deletions

diff --git a/src/load.c b/src/load.c
index 48496ba48..1ab8c1ad6 100644
--- a/src/load.c
+++ b/src/load.c
@@ -26,6 +26,17 @@
 
 #define SIZE_ERROR_MUL(nmemb, size) ((size_t)(nmemb) > SIZE_MAX / (size))
 
+#define DEFINE_READ_IREP_FUNC(funcdecl, basecall) \
+  funcdecl \
+  { \
+    int ai = mrb_gc_arena_save(mrb); \
+    struct RProc *proc = basecall; \
+    struct mrb_irep *irep = (mrb_irep*)(proc ? proc->body.irep : NULL); \
+    if (irep) proc->body.irep = NULL; \
+    mrb_gc_arena_restore(mrb, ai); \
+    return irep; \
+  }
+
 static size_t
 offset_crc_body(void)
 {
@@ -58,29 +69,20 @@ str_to_double(mrb_state *mrb, const char *p)
 
 mrb_value mrb_str_len_to_inum(mrb_state *mrb, const char *str, size_t len, mrb_int base, int badcheck);
 
-static void
-tempirep_free(mrb_state *mrb, void *p)
-{
-  if (p) mrb_irep_decref(mrb, (mrb_irep *)p);
-}
-
-static const mrb_data_type tempirep_type = { "temporary irep", tempirep_free };
-
-static mrb_irep*
-read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flags)
+static mrb_bool
+read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flags, mrb_irep **irepp)
 {
   int i;
   const uint8_t *src = bin;
   ptrdiff_t diff;
   uint16_t tt, pool_data_len, snl;
   int plen;
-  struct RData *irep_obj = mrb_data_object_alloc(mrb, mrb->object_class, NULL, &tempirep_type);
   mrb_pool_value *pool;
   mrb_sym *syms;
   int ai = mrb_gc_arena_save(mrb);
   mrb_irep *irep = mrb_add_irep(mrb);
 
-  irep_obj->data = irep;
+  *irepp = irep;
 
   /* skip record size */
   src += sizeof(uint32_t);
@@ -109,7 +111,7 @@ read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flag
                       sizeof(struct mrb_irep_catch_handler) * irep->clen;
     mrb_static_assert1(sizeof(struct mrb_irep_catch_handler) == 13);
     if (SIZE_ERROR_MUL(irep->ilen, sizeof(mrb_code))) {
-      return NULL;
+      return FALSE;
     }
     if ((flags & FLAG_SRC_MALLOC) == 0) {
       irep->iseq = (mrb_code*)src;
@@ -128,7 +130,7 @@ read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flag
   src += sizeof(uint16_t);
   if (plen > 0) {
     if (SIZE_ERROR_MUL(plen, sizeof(mrb_value))) {
-      return NULL;
+      return FALSE;
     }
     irep->pool = pool = (mrb_pool_value*)mrb_calloc(mrb, sizeof(mrb_pool_value), plen);
 
@@ -162,7 +164,7 @@ read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flag
         }
         break;
 #else
-        return NULL;            /* INT64 not supported on MRB_32BIT */
+        return FALSE;           /* INT64 not supported on MRB_32BIT */
 #endif
 
       case IREP_TT_FLOAT:
@@ -172,7 +174,7 @@ read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flag
         src += sizeof(double);
         break;
 #else
-        return NULL;            /* MRB_NO_FLOAT */
+        return FALSE;           /* MRB_NO_FLOAT */
 #endif
 
       case IREP_TT_STR:
@@ -194,7 +196,7 @@ read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flag
 
       default:
         /* should not happen */
-        return NULL;
+        return FALSE;
       }
       irep->plen = i+1;
     }
@@ -205,7 +207,7 @@ read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flag
   src += sizeof(uint16_t);
   if (irep->slen > 0) {
     if (SIZE_ERROR_MUL(irep->slen, sizeof(mrb_sym))) {
-      return NULL;
+      return FALSE;
     }
     irep->syms = syms = (mrb_sym *)mrb_malloc(mrb, sizeof(mrb_sym) * irep->slen);
 
@@ -233,54 +235,60 @@ read_irep_record_1(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flag
   mrb_assert_int_fit(ptrdiff_t, diff, size_t, SIZE_MAX);
   *len = (size_t)diff;
 
-  irep_obj->data = NULL;
-
-  return irep;
+  return TRUE;
 }
 
-static mrb_irep*
-read_irep_record(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flags)
+static mrb_bool
+read_irep_record(mrb_state *mrb, const uint8_t *bin, size_t *len, uint8_t flags, mrb_irep **irepp)
 {
-  struct RData *irep_obj = mrb_data_object_alloc(mrb, mrb->object_class, NULL, &tempirep_type);
   int ai = mrb_gc_arena_save(mrb);
-  mrb_irep *irep = read_irep_record_1(mrb, bin, len, flags);
+  mrb_bool readsuccess = read_irep_record_1(mrb, bin, len, flags, irepp);
   mrb_irep **reps;
   int i;
 
   mrb_gc_arena_restore(mrb, ai);
-  if (irep == NULL) {
-    return NULL;
+  if (!readsuccess) {
+    return FALSE;
   }
 
-  reps = (mrb_irep**)mrb_calloc(mrb, irep->rlen, sizeof(mrb_irep*));
-  irep->reps = (const mrb_irep**)reps;
-  irep_obj->data = irep;
+  reps = (mrb_irep**)mrb_calloc(mrb, (*irepp)->rlen, sizeof(mrb_irep*));
+  (*irepp)->reps = (const mrb_irep**)reps;
 
   bin += *len;
-  for (i=0; i<irep->rlen; i++) {
+  for (i=0; i<(*irepp)->rlen; i++) {
     size_t rlen;
 
-    reps[i] = read_irep_record(mrb, bin, &rlen, flags);
+    readsuccess = read_irep_record(mrb, bin, &rlen, flags, &reps[i]);
     mrb_gc_arena_restore(mrb, ai);
-    if (reps[i] == NULL) {
-      return NULL;
+    if (!readsuccess) {
+      return FALSE;
     }
     bin += rlen;
     *len += rlen;
   }
 
-  irep_obj->data = NULL;
-
-  return irep;
+  return TRUE;
 }
 
 static mrb_irep*
-read_section_irep(mrb_state *mrb, const uint8_t *bin, uint8_t flags)
+read_section_irep(mrb_state *mrb, const uint8_t *bin, uint8_t flags, struct RProc **proc)
 {
   size_t len;
 
+  /*
+   * This proc object keeps all the data in progress to avoid memory leaks
+   * if something goes wrong while reading irep.
+   */
+  *proc = mrb_proc_new(mrb, NULL);
+
+  mrb_irep **irepp = (mrb_irep**)&(*proc)->body.irep;
   bin += sizeof(struct rite_section_irep_header);
-  return read_irep_record(mrb, bin, &len, flags);
+  if (read_irep_record(mrb, bin, &len, flags, irepp)) {
+    return *irepp;
+  }
+  else {
+    return NULL;
+  }
 }
 
 static int
@@ -544,11 +552,11 @@ read_binary_header(const uint8_t *bin, size_t bufsize, size_t *bin_size, uint16_
   return MRB_DUMP_OK;
 }
 
-static mrb_irep*
+static struct RProc*
 read_irep(mrb_state *mrb, const uint8_t *bin, size_t bufsize, uint8_t flags)
 {
   int result;
-  struct RData *irep_obj = NULL;
+  struct RProc *proc = NULL;
   mrb_irep *irep = NULL;
   const struct rite_section_header *section_header;
   uint16_t crc;
@@ -569,15 +577,12 @@ read_irep(mrb_state *mrb, const uint8_t *bin, size_t bufsize, uint8_t flags)
     return NULL;
   }
 
-  irep_obj = mrb_data_object_alloc(mrb, mrb->object_class, NULL, &tempirep_type);
-
   bin += sizeof(struct rite_binary_header);
   do {
     section_header = (const struct rite_section_header *)bin;
     if (memcmp(section_header->section_ident, RITE_SECTION_IREP_IDENT, sizeof(section_header->section_ident)) == 0) {
-      irep = read_section_irep(mrb, bin, flags);
+      irep = read_section_irep(mrb, bin, flags, &proc);
       if (!irep) return NULL;
-      irep_obj->data = irep;
     }
     else if (memcmp(section_header->section_ident, RITE_SECTION_DEBUG_IDENT, sizeof(section_header->section_ident)) == 0) {
       if (!irep) return NULL;   /* corrupted data */
@@ -596,13 +601,11 @@ read_irep(mrb_state *mrb, const uint8_t *bin, size_t bufsize, uint8_t flags)
     bin += bin_to_uint32(section_header->section_size);
   } while (memcmp(section_header->section_ident, RITE_BINARY_EOF, sizeof(section_header->section_ident)) != 0);
 
-  irep_obj->data = NULL;
-
-  return irep;
+  return proc;
 }
 
-mrb_irep*
-mrb_read_irep(mrb_state *mrb, const uint8_t *bin)
+static struct RProc*
+mrb_proc_read_irep(mrb_state *mrb, const uint8_t *bin)
 {
 #if defined(MRB_USE_LINK_TIME_RO_DATA_P) || defined(MRB_USE_CUSTOM_RO_DATA_P)
   uint8_t flags = mrb_ro_data_p((char*)bin) ? FLAG_SRC_STATIC : FLAG_SRC_MALLOC;
@@ -613,12 +616,20 @@ mrb_read_irep(mrb_state *mrb, const uint8_t *bin)
   return read_irep(mrb, bin, (size_t)-1, flags);
 }
 
-MRB_API mrb_irep*
-mrb_read_irep_buf(mrb_state *mrb, const void *buf, size_t bufsize)
+DEFINE_READ_IREP_FUNC(
+  mrb_irep *mrb_read_irep(mrb_state *mrb, const uint8_t *bin),
+  mrb_proc_read_irep(mrb, bin))
+
+static struct RProc*
+mrb_proc_read_irep_buf(mrb_state *mrb, const void *buf, size_t bufsize)
 {
   return read_irep(mrb, (const uint8_t *)buf, bufsize, FLAG_SRC_MALLOC);
 }
 
+DEFINE_READ_IREP_FUNC(
+  MRB_API mrb_irep *mrb_read_irep_buf(mrb_state *mrb, const void *buf, size_t bufsize),
+  mrb_proc_read_irep_buf(mrb, buf, bufsize))
+
 void mrb_exc_set(mrb_state *mrb, mrb_value exc);
 
 static void
@@ -630,17 +641,13 @@ irep_error(mrb_state *mrb)
 void mrb_codedump_all(mrb_state*, struct RProc*);
 
 static mrb_value
-load_irep(mrb_state *mrb, mrb_irep *irep, mrbc_context *c)
+load_irep(mrb_state *mrb, struct RProc *proc, mrbc_context *c)
 {
-  struct RProc *proc;
-
-  if (!irep) {
+  if (!proc || !proc->body.irep) {
     irep_error(mrb);
     return mrb_nil_value();
   }
-  proc = mrb_proc_new(mrb, irep);
   proc->c = NULL;
-  mrb_irep_decref(mrb, irep);
   if (c && c->dump_result) mrb_codedump_all(mrb, proc);
   if (c && c->no_exec) return mrb_obj_value(proc);
   return mrb_top_run(mrb, proc, mrb_top_self(mrb), 0);
@@ -649,23 +656,15 @@ load_irep(mrb_state *mrb, mrb_irep *irep, mrbc_context *c)
 MRB_API mrb_value
 mrb_load_irep_cxt(mrb_state *mrb, const uint8_t *bin, mrbc_context *c)
 {
-  struct RData *irep_obj = mrb_data_object_alloc(mrb, mrb->object_class, NULL, &tempirep_type);
-  mrb_irep *irep = mrb_read_irep(mrb, bin);
-  mrb_value ret;
-
-  if (!irep) return mrb_undef_value();
-  irep_obj->data = irep;
-  mrb_irep_incref(mrb, irep);
-  ret = load_irep(mrb, irep, c);
-  irep_obj->data = NULL;
-  mrb_irep_decref(mrb, irep);
-  return ret;
+  struct RProc *proc = mrb_proc_read_irep(mrb, bin);
+  if (!proc) return mrb_undef_value();
+  return load_irep(mrb, proc, c);
 }
 
 MRB_API mrb_value
 mrb_load_irep_buf_cxt(mrb_state *mrb, const void *buf, size_t bufsize, mrbc_context *c)
 {
-  return load_irep(mrb, mrb_read_irep_buf(mrb, buf, bufsize), c);
+  return load_irep(mrb, mrb_proc_read_irep_buf(mrb, buf, bufsize), c);
 }
 
 MRB_API mrb_value
@@ -688,10 +687,10 @@ mrb_load_proc(mrb_state *mrb, const struct RProc *proc)
 
 #ifndef MRB_NO_STDIO
 
-mrb_irep*
-mrb_read_irep_file(mrb_state *mrb, FILE* fp)
+static struct RProc*
+mrb_proc_read_irep_file(mrb_state *mrb, FILE *fp)
 {
-  mrb_irep *irep = NULL;
+  struct RProc *proc = NULL;
   uint8_t *buf;
   const size_t header_size = sizeof(struct rite_binary_header);
   size_t buf_size = 0;
@@ -715,17 +714,21 @@ mrb_read_irep_file(mrb_state *mrb, FILE* fp)
   if (fread(buf+header_size, buf_size-header_size, 1, fp) == 0) {
     goto irep_exit;
   }
-  irep = read_irep(mrb, buf, (size_t)-1, FLAG_SRC_MALLOC);
+  proc = read_irep(mrb, buf, (size_t)-1, FLAG_SRC_MALLOC);
 
 irep_exit:
   mrb_free(mrb, buf);
-  return irep;
+  return proc;
 }
 
+DEFINE_READ_IREP_FUNC(
+  mrb_irep *mrb_read_irep_file(mrb_state *mrb, FILE *fp),
+  mrb_proc_read_irep_file(mrb, fp))
+
 MRB_API mrb_value
 mrb_load_irep_file_cxt(mrb_state *mrb, FILE* fp, mrbc_context *c)
 {
-  return load_irep(mrb, mrb_read_irep_file(mrb, fp), c);
+  return load_irep(mrb, mrb_proc_read_irep_file(mrb, fp), c);
 }
 
 MRB_API mrb_value
diff --git a/src/proc.c b/src/proc.c
index d58d9d0df..68ca941c3 100644
--- a/src/proc.c
+++ b/src/proc.c
@@ -52,7 +52,9 @@ mrb_proc_new(mrb_state *mrb, const mrb_irep *irep)
     p->e.target_class = tc;
   }
   p->body.irep = irep;
-  mrb_irep_incref(mrb, (mrb_irep*)irep);
+  if (irep) {
+    mrb_irep_incref(mrb, (mrb_irep*)irep);
+  }
 
   return p;
 }