Merge pull request #3065 from kou/support-backtrace-after-method-calls

Support backtrace after method calls

5 files changed, 307 insertions, 25 deletions

diff --git a/src/backtrace.c b/src/backtrace.c
index 80a5e2935..b581083df 100644
--- a/src/backtrace.c
+++ b/src/backtrace.c
@@ -14,6 +14,15 @@
 #include <mruby/error.h>
 #include <mruby/numeric.h>
 
+struct backtrace_location_raw {
+  int i;
+  int lineno;
+  const char *filename;
+  mrb_sym method_id;
+  const char *sep;
+  struct RClass *klass;
+};
+
 struct backtrace_location {
   int i;
   int lineno;
@@ -23,6 +32,7 @@ struct backtrace_location {
   const char *class_name;
 };
 
+typedef void (*each_backtrace_func)(mrb_state*, struct backtrace_location_raw*, void*);
 typedef void (*output_stream_func)(mrb_state*, struct backtrace_location*, void*);
 
 #ifndef MRB_DISABLE_STDIO
@@ -89,7 +99,7 @@ get_backtrace_i(mrb_state *mrb, struct backtrace_location *loc, void *data)
 }
 
 static void
-output_backtrace(mrb_state *mrb, mrb_int ciidx, mrb_code *pc0, output_stream_func func, void *data)
+each_backtrace(mrb_state *mrb, mrb_int ciidx, mrb_code *pc0, each_backtrace_func func, void *data)
 {
   int i;
 
@@ -97,7 +107,7 @@ output_backtrace(mrb_state *mrb, mrb_int ciidx, mrb_code *pc0, output_stream_fun
     ciidx = 10; /* ciidx is broken... */
 
   for (i = ciidx; i >= 0; i--) {
-    struct backtrace_location loc;
+    struct backtrace_location_raw loc;
     mrb_callinfo *ci;
     mrb_irep *irep;
     mrb_code *pc;
@@ -134,13 +144,40 @@ output_backtrace(mrb_state *mrb, mrb_int ciidx, mrb_code *pc0, output_stream_fun
       loc.filename = "(unknown)";
     }
 
-    loc.method = mrb_sym2name(mrb, ci->mid);
-    loc.class_name = mrb_class_name(mrb, ci->proc->target_class);
+    loc.method_id = ci->mid;
+    loc.klass = ci->proc->target_class;
     loc.i = i;
     func(mrb, &loc, data);
   }
 }
 
+struct output_backtrace_args {
+  output_stream_func func;
+  void *data;
+};
+
+static void
+output_backtrace_i(mrb_state *mrb, struct backtrace_location_raw *loc_raw, void *data)
+{
+  struct backtrace_location loc;
+  struct output_backtrace_args *args = data;
+
+  loc.i          = loc_raw->i;
+  loc.lineno     = loc_raw->lineno;
+  loc.filename   = loc_raw->filename;
+  loc.method     = mrb_sym2name(mrb, loc_raw->method_id);
+  loc.sep        = loc_raw->sep;
+  loc.class_name = mrb_class_name(mrb, loc_raw->klass);
+
+  args->func(mrb, &loc, args->data);
+}
+
+static void
+output_backtrace(mrb_state *mrb, mrb_int ciidx, mrb_code *pc0, output_stream_func func, void *data)
+{
+  each_backtrace(mrb, ciidx, pc0, output_backtrace_i, data);
+}
+
 static void
 exc_output_backtrace(mrb_state *mrb, struct RObject *exc, output_stream_func func, void *stream)
 {
@@ -167,18 +204,76 @@ exc_output_backtrace(mrb_state *mrb, struct RObject *exc, output_stream_func fun
 
 #ifndef MRB_DISABLE_STDIO
 
+static void
+print_backtrace(mrb_state *mrb, mrb_value backtrace)
+{
+  int i, n;
+  FILE *stream = stderr;
+
+  fprintf(stream, "trace:\n");
+
+  n = RARRAY_LEN(backtrace);
+  for (i = 0; i < n; i++) {
+    mrb_value entry = RARRAY_PTR(backtrace)[i];
+
+    fprintf(stream, "\t[%d] %.*s\n", i, RSTRING_LEN(entry), RSTRING_PTR(entry));
+  }
+}
+
+static void
+print_backtrace_saved(mrb_state *mrb)
+{
+  int i;
+  FILE *stream = stderr;
+
+  fprintf(stream, "trace:\n");
+  for (i = 0; i < mrb->backtrace.n; i++) {
+    mrb_backtrace_entry *entry;
+
+    entry = &(mrb->backtrace.entries[i]);
+    fprintf(stream, "\t[%d] %s:%d", i, entry->filename, entry->lineno);
+
+    if (entry->method_id != 0) {
+      const char *method_name;
+
+      method_name = mrb_sym2name(mrb, entry->method_id);
+      if (entry->klass) {
+        fprintf(stream, ":in %s%s%s",
+                mrb_class_name(mrb, entry->klass),
+                entry->sep,
+                method_name);
+      }
+      else {
+        fprintf(stream, ":in %s", method_name);
+      }
+    }
+
+    fprintf(stream, "\n");
+  }
+}
+
 MRB_API void
 mrb_print_backtrace(mrb_state *mrb)
 {
-  struct print_backtrace_args args;
+  mrb_value backtrace;
 
   if (!mrb->exc || mrb_obj_is_kind_of(mrb, mrb_obj_value(mrb->exc), E_SYSSTACK_ERROR)) {
     return;
   }
 
-  args.stream = stderr;
-  args.tracehead = TRUE;
-  exc_output_backtrace(mrb, mrb->exc, print_backtrace_i, (void*)&args);
+  backtrace = mrb_obj_iv_get(mrb, mrb->exc, mrb_intern_lit(mrb, "backtrace"));
+  if (!mrb_nil_p(backtrace)) {
+    print_backtrace(mrb, backtrace);
+  }
+  else if (mrb->backtrace.n > 0) {
+    print_backtrace_saved(mrb);
+  }
+  else {
+    struct print_backtrace_args args;
+    args.stream = stderr;
+    args.tracehead = TRUE;
+    exc_output_backtrace(mrb, mrb->exc, print_backtrace_i, (void*)&args);
+  }
 }
 
 #else
@@ -215,3 +310,112 @@ mrb_get_backtrace(mrb_state *mrb)
 
   return ary;
 }
+
+void
+mrb_free_backtrace(mrb_state *mrb)
+{
+  mrb->backtrace.exc = 0;
+  mrb->backtrace.n = 0;
+  mrb->backtrace.n_allocated = 0;
+  mrb_free(mrb, mrb->backtrace.entries);
+}
+
+static void
+save_backtrace_i(mrb_state *mrb,
+                 struct backtrace_location_raw *loc_raw,
+                 void *data)
+{
+  mrb_backtrace_entry *entry;
+
+  if (loc_raw->i >= mrb->backtrace.n_allocated) {
+    int new_n_allocated;
+    if (mrb->backtrace.n_allocated == 0) {
+      new_n_allocated = 8;
+    }
+    else {
+      new_n_allocated = mrb->backtrace.n_allocated * 2;
+    }
+    mrb->backtrace.entries =
+      mrb_realloc(mrb,
+                  mrb->backtrace.entries,
+                  sizeof(mrb_backtrace_entry) * new_n_allocated);
+    mrb->backtrace.n_allocated = new_n_allocated;
+  }
+
+  entry = &mrb->backtrace.entries[mrb->backtrace.n];
+  entry->filename  = loc_raw->filename;
+  entry->lineno    = loc_raw->lineno;
+  entry->klass     = loc_raw->klass;
+  entry->sep       = loc_raw->sep;
+  entry->method_id = loc_raw->method_id;
+
+  mrb->backtrace.n++;
+}
+
+void
+mrb_save_backtrace(mrb_state *mrb)
+{
+  mrb_value lastpc;
+  mrb_code *code;
+  mrb_int ciidx;
+
+  mrb->backtrace.n = 0;
+  mrb->backtrace.exc = 0;
+
+  if (!mrb->exc)
+    return;
+
+  mrb->backtrace.exc = mrb->exc;
+
+  lastpc = mrb_obj_iv_get(mrb, mrb->exc, mrb_intern_lit(mrb, "lastpc"));
+  if (mrb_nil_p(lastpc)) {
+    code = NULL;
+  }
+  else {
+    code = (mrb_code*)mrb_cptr(lastpc);
+  }
+
+  ciidx = mrb_fixnum(mrb_obj_iv_get(mrb, mrb->exc, mrb_intern_lit(mrb, "ciidx")));
+
+  each_backtrace(mrb, ciidx, code, save_backtrace_i, NULL);
+}
+
+mrb_value
+mrb_restore_backtrace(mrb_state *mrb)
+{
+  int i;
+  mrb_value backtrace;
+
+  backtrace = mrb_ary_new(mrb);
+  for (i = 0; i < mrb->backtrace.n; i++) {
+    int ai;
+    mrb_backtrace_entry *entry;
+    mrb_value mrb_entry;
+
+    ai = mrb_gc_arena_save(mrb);
+    entry = &(mrb->backtrace.entries[i]);
+
+    mrb_entry = mrb_str_new_cstr(mrb, entry->filename);
+    mrb_str_cat_lit(mrb, mrb_entry, ":");
+    mrb_str_concat(mrb, mrb_entry,
+                   mrb_fixnum_to_str(mrb,
+                                     mrb_fixnum_value(entry->lineno),
+                                     10));
+    if (entry->method_id != 0) {
+      mrb_str_cat_lit(mrb, mrb_entry, ":in ");
+
+      if (entry->klass) {
+        mrb_str_cat_cstr(mrb, mrb_entry, mrb_class_name(mrb, entry->klass));
+        mrb_str_cat_cstr(mrb, mrb_entry, entry->sep);
+      }
+
+      mrb_str_cat_cstr(mrb, mrb_entry, mrb_sym2name(mrb, entry->method_id));
+    }
+
+    mrb_ary_push(mrb, backtrace, mrb_entry);
+
+    mrb_gc_arena_restore(mrb, ai);
+  }
+
+  return backtrace;
+}
diff --git a/src/error.c b/src/error.c
index 15a969d93..14e4ab4d3 100644
--- a/src/error.c
+++ b/src/error.c
@@ -174,6 +174,42 @@ exc_inspect(mrb_state *mrb, mrb_value exc)
   return str;
 }
 
+void mrb_save_backtrace(mrb_state *mrb);
+mrb_value mrb_restore_backtrace(mrb_state *mrb);
+
+static mrb_value
+exc_get_backtrace(mrb_state *mrb, mrb_value exc)
+{
+  mrb_sym attr_name;
+  mrb_value backtrace;
+
+  attr_name = mrb_intern_lit(mrb, "backtrace");
+  backtrace = mrb_iv_get(mrb, exc, attr_name);
+  if (mrb_nil_p(backtrace)) {
+    if (mrb_obj_ptr(exc) == mrb->backtrace.exc && mrb->backtrace.n > 0) {
+      backtrace = mrb_restore_backtrace(mrb);
+      mrb->backtrace.n = 0;
+      mrb->backtrace.exc = 0;
+    }
+    else {
+      backtrace = mrb_exc_backtrace(mrb, exc);
+    }
+    mrb_iv_set(mrb, exc, attr_name, backtrace);
+  }
+
+  return backtrace;
+}
+
+static mrb_value
+exc_set_backtrace(mrb_state *mrb, mrb_value exc)
+{
+  mrb_value backtrace;
+
+  mrb_get_args(mrb, "o", &backtrace);
+  mrb_iv_set(mrb, exc, mrb_intern_lit(mrb, "backtrace"), backtrace);
+
+  return backtrace;
+}
 
 static void
 exc_debug_info(mrb_state *mrb, struct RObject *exc)
@@ -202,12 +238,52 @@ exc_debug_info(mrb_state *mrb, struct RObject *exc)
   }
 }
 
+static void
+set_backtrace(mrb_state *mrb, mrb_value info, mrb_value bt)
+{
+  mrb_funcall(mrb, info, "set_backtrace", 1, bt);
+}
+
+static mrb_bool
+have_backtrace(mrb_state *mrb, struct RObject *exc)
+{
+  return !mrb_nil_p(mrb_obj_iv_get(mrb, exc, mrb_intern_lit(mrb, "backtrace")));
+}
+
+void
+mrb_exc_set(mrb_state *mrb, mrb_value exc)
+{
+  if (!mrb->gc.out_of_memory && mrb->backtrace.n > 0) {
+    mrb_value target_exc = mrb_nil_value();
+    if ((mrb->exc && !have_backtrace(mrb, mrb->exc))) {
+      target_exc = mrb_obj_value(mrb->exc);
+    }
+    else if (!mrb_nil_p(exc) && mrb_obj_ptr(exc) == mrb->backtrace.exc) {
+      target_exc = exc;
+    }
+    if (!mrb_nil_p(target_exc)) {
+      mrb_value backtrace;
+      backtrace = mrb_restore_backtrace(mrb);
+      set_backtrace(mrb, target_exc, backtrace);
+    }
+  }
+
+  mrb->backtrace.n = 0;
+  if (mrb_nil_p(exc)) {
+    mrb->exc = 0;
+  }
+  else {
+    mrb->exc = mrb_obj_ptr(exc);
+  }
+}
+
 MRB_API mrb_noreturn void
 mrb_exc_raise(mrb_state *mrb, mrb_value exc)
 {
-  mrb->exc = mrb_obj_ptr(exc);
+  mrb_exc_set(mrb, exc);
   if (!mrb->gc.out_of_memory) {
     exc_debug_info(mrb, mrb->exc);
+    mrb_save_backtrace(mrb);
   }
   if (!mrb->jmp) {
     mrb_p(mrb, exc);
@@ -337,12 +413,6 @@ mrb_bug(mrb_state *mrb, const char *fmt, ...)
   exit(EXIT_FAILURE);
 }
 
-static void
-set_backtrace(mrb_state *mrb, mrb_value info, mrb_value bt)
-{
-  mrb_funcall(mrb, info, "set_backtrace", 1, bt);
-}
-
 static mrb_value
 make_exception(mrb_state *mrb, int argc, const mrb_value *argv, mrb_bool isstr)
 {
@@ -449,7 +519,8 @@ mrb_init_exception(mrb_state *mrb)
   mrb_define_method(mrb, exception, "to_s",            exc_to_s,          MRB_ARGS_NONE());
   mrb_define_method(mrb, exception, "message",         exc_message,       MRB_ARGS_NONE());
   mrb_define_method(mrb, exception, "inspect",         exc_inspect,       MRB_ARGS_NONE());
-  mrb_define_method(mrb, exception, "backtrace",       mrb_exc_backtrace, MRB_ARGS_NONE());
+  mrb_define_method(mrb, exception, "backtrace",       exc_get_backtrace, MRB_ARGS_NONE());
+  mrb_define_method(mrb, exception, "set_backtrace",   exc_set_backtrace, MRB_ARGS_REQ(1));
 
   mrb->eStandardError_class = mrb_define_class(mrb, "StandardError", mrb->eException_class); /* 15.2.23 */
   runtime_error = mrb_define_class(mrb, "RuntimeError", mrb->eStandardError_class);          /* 15.2.28 */
diff --git a/src/load.c b/src/load.c
index da88f0d3a..0b47f2b89 100644
--- a/src/load.c
+++ b/src/load.c
@@ -614,10 +614,12 @@ mrb_read_irep(mrb_state *mrb, const uint8_t *bin)
   return read_irep(mrb, bin, flags);
 }
 
+void mrb_exc_set(mrb_state *mrb, mrb_value exc);
+
 static void
 irep_error(mrb_state *mrb)
 {
-  mrb->exc = mrb_obj_ptr(mrb_exc_new_str_lit(mrb, E_SCRIPT_ERROR, "irep load error"));
+  mrb_exc_set(mrb, mrb_exc_new_str_lit(mrb, E_SCRIPT_ERROR, "irep load error"));
 }
 
 MRB_API mrb_value
diff --git a/src/state.c b/src/state.c
index c8c0658e4..898e77336 100644
--- a/src/state.c
+++ b/src/state.c
@@ -215,6 +215,8 @@ mrb_str_pool(mrb_state *mrb, mrb_value str)
   return mrb_obj_value(ns);
 }
 
+void mrb_free_backtrace(mrb_state *mrb);
+
 MRB_API void
 mrb_free_context(mrb_state *mrb, struct mrb_context *c)
 {
@@ -242,6 +244,7 @@ mrb_close(mrb_state *mrb)
 
   /* free */
   mrb_gc_free_gv(mrb);
+  mrb_free_backtrace(mrb);
   mrb_free_context(mrb, mrb->root_c);
   mrb_free_symtbl(mrb);
   mrb_alloca_free(mrb);
diff --git a/src/vm.c b/src/vm.c
index 391646017..b2149b250 100644
--- a/src/vm.c
+++ b/src/vm.c
@@ -258,6 +258,8 @@ cipop(mrb_state *mrb)
   c->ci--;
 }
 
+void mrb_exc_set(mrb_state *mrb, mrb_value exc);
+
 static void
 ecall(mrb_state *mrb, int i)
 {
@@ -669,7 +671,7 @@ localjump_error(mrb_state *mrb, localjump_error_kind kind)
   mrb_str_cat(mrb, msg, lead, sizeof(lead) - 1);
   mrb_str_cat(mrb, msg, kind_str[kind], kind_str_len[kind]);
   exc = mrb_exc_new_str(mrb, E_LOCALJUMP_ERROR, msg);
-  mrb->exc = mrb_obj_ptr(exc);
+  mrb_exc_set(mrb, exc);
 }
 
 static void
@@ -688,7 +690,7 @@ argnum_error(mrb_state *mrb, mrb_int num)
                   mrb_fixnum_value(mrb->c->ci->argc), mrb_fixnum_value(num));
   }
   exc = mrb_exc_new_str(mrb, E_ARGUMENT_ERROR, str);
-  mrb->exc = mrb_obj_ptr(exc);
+  mrb_exc_set(mrb, exc);
 }
 
 #define ERR_PC_SET(mrb, pc) mrb->c->ci->err = pc;
@@ -1007,7 +1009,7 @@ RETRY_TRY_BLOCK:
 
     CASE(OP_RAISE) {
       /* A      raise(R(A)) */
-      mrb->exc = mrb_obj_ptr(regs[GETARG_A(i)]);
+      mrb_exc_set(mrb, regs[GETARG_A(i)]);
       goto L_RAISE;
     }
 
@@ -1241,7 +1243,7 @@ RETRY_TRY_BLOCK:
         mrb_value exc;
 
         exc = mrb_exc_new_str_lit(mrb, E_NOMETHOD_ERROR, "super called outside of method");
-        mrb->exc = mrb_obj_ptr(exc);
+        mrb_exc_set(mrb, exc);
         goto L_RAISE;
       }
       recv = regs[0];
@@ -1331,7 +1333,7 @@ RETRY_TRY_BLOCK:
           mrb_value exc;
 
           exc = mrb_exc_new_str_lit(mrb, E_NOMETHOD_ERROR, "super called outside of method");
-          mrb->exc = mrb_obj_ptr(exc);
+          mrb_exc_set(mrb, exc);
           goto L_RAISE;
         }
         stack = e->stack + 1;
@@ -1561,7 +1563,7 @@ RETRY_TRY_BLOCK:
             }
             if (mrb->c->prev->ci == mrb->c->prev->cibase) {
               mrb_value exc = mrb_exc_new_str_lit(mrb, E_FIBER_ERROR, "double resume");
-              mrb->exc = mrb_obj_ptr(exc);
+              mrb_exc_set(mrb, exc);
               goto L_RAISE;
             }
             /* automatic yield at the end */
@@ -2323,7 +2325,7 @@ RETRY_TRY_BLOCK:
       /* A      R(A) := target_class */
       if (!mrb->c->ci->target_class) {
         mrb_value exc = mrb_exc_new_str_lit(mrb, E_TYPE_ERROR, "no target class or module");
-        mrb->exc = mrb_obj_ptr(exc);
+        mrb_exc_set(mrb, exc);
         goto L_RAISE;
       }
       regs[GETARG_A(i)] = mrb_obj_value(mrb->c->ci->target_class);
@@ -2381,7 +2383,7 @@ RETRY_TRY_BLOCK:
       else {
         exc = mrb_exc_new_str(mrb, E_LOCALJUMP_ERROR, msg);
       }
-      mrb->exc = mrb_obj_ptr(exc);
+      mrb_exc_set(mrb, exc);
       goto L_RAISE;
     }
   }