From ce80a2dfdf8d5edf9dc8ed58a1e458628ec6b48e Mon Sep 17 00:00:00 2001
From: Yuichi Nishiwaki <yuichi.nishiwaki@gmail.com>
Date: Thu, 20 Apr 2017 13:22:28 -0700
Subject: [PATCH] add cxt.ai

---
 lib/error.c          | 8 +++++---
 lib/include/picrin.h | 4 ----
 lib/proc.c           | 7 +++----
 lib/state.c          | 2 ++
 lib/state.h          | 1 +
 5 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/lib/error.c b/lib/error.c
index 2d88ef38..6694eb90 100644
--- a/lib/error.c
+++ b/lib/error.c
@@ -65,7 +65,8 @@ pic_enter_try(pic_state *pic)
 {
   pic_value cont, handler;
   pic_value var, env;
-  size_t ai = pic_enter(pic);
+
+  pic->cxt->ai = pic->ai;
 
   /* call/cc */
   cont = pic_make_cont(pic, pic->cxt, pic_invalid_value(pic), pic->dyn_env);
@@ -76,7 +77,7 @@ pic_enter_try(pic_state *pic)
   pic_weak_set(pic, env, var, pic_cons(pic, handler, pic_call(pic, var, 0)));
   pic->dyn_env = pic_cons(pic, env, pic->dyn_env);
 
-  pic_leave(pic, ai);
+  pic_leave(pic, pic->cxt->ai);
 }
 
 void
@@ -86,6 +87,7 @@ pic_exit_try(pic_state *pic)
   pic->dyn_env = pic_cdr(pic, pic->dyn_env);
   pic->cxt = cxt->prev;
   pic_free(pic, cxt);
+  /* don't rewind ai here */
 }
 
 pic_value
@@ -93,7 +95,7 @@ pic_abort_try(pic_state *pic)
 {
   struct context *cxt = pic->cxt;
   pic_value err = cxt->sp->regs[1];
-  pic->cxt = pic->cxt->prev;
+  pic->cxt = cxt->prev;
   pic_free(pic, cxt);
   pic_protect(pic, err);
   return err;
diff --git a/lib/include/picrin.h b/lib/include/picrin.h
index 1b09adb3..bcdcf070 100644
--- a/lib/include/picrin.h
+++ b/lib/include/picrin.h
@@ -304,18 +304,14 @@ PIC_NORETURN void pic_error(pic_state *, const char *msg, int n, ...);
     extern void pic_enter_try(pic_state *);                             \
     extern void pic_exit_try(pic_state *);                              \
     extern pic_value pic_abort_try(pic_state *);                        \
-    size_t pic_try_ai_ = pic_enter(pic);                                \
     PIC_JMPBUF *jmp = pic_prepare_try(pic);                             \
     if (PIC_SETJMP(*jmp) == 0) {                                        \
       pic_enter_try(pic);
 #define pic_catch(e) pic_catch_(e, PIC_GENSYM(label))
 #define pic_catch_(e, label)                              \
       pic_exit_try(pic);                                  \
-      /* don't rewind ai here */                          \
     } else {                                              \
       e = pic_abort_try(pic);                             \
-      pic_leave(pic, pic_try_ai_);                        \
-      pic_protect(pic, e);                                \
       goto label;                                         \
     }                                                     \
   } while (0);                                            \
diff --git a/lib/proc.c b/lib/proc.c
index 9fe96568..69dc8c64 100644
--- a/lib/proc.c
+++ b/lib/proc.c
@@ -364,7 +364,6 @@ pic_value
 pic_apply(pic_state *pic, pic_value proc, int argc, pic_value *argv)
 {
   struct context cxt;
-  size_t arena_base = pic->ai;
 
 #define MKCALL(argc) (cxt.tmpcode[0] = OP_CALL, cxt.tmpcode[1] = (argc), cxt.tmpcode)
 
@@ -383,10 +382,10 @@ pic_apply(pic_state *pic, pic_value proc, int argc, pic_value *argv)
   cxt.prev = pic->cxt;
   pic->cxt = &cxt;
 
-#define SAVE (pic->ai = arena_base)
+#define SAVE (pic->ai = cxt.ai)
 
-  if (PIC_SETJMP(cxt.jmp) != 0) {
-    SAVE;
+  if (PIC_SETJMP(cxt.jmp) == 0) {
+    cxt.ai = pic->ai;
   }
 
 #define A (cxt.pc[1])
diff --git a/lib/state.c b/lib/state.c
index a861cf51..fe625aa7 100644
--- a/lib/state.c
+++ b/lib/state.c
@@ -171,6 +171,7 @@ pic_open(pic_allocf allocf, void *userdata)
   pic->userdata = userdata;
 
   /* context */
+  pic->default_cxt.ai = 0;
   pic->default_cxt.pc = NULL;
   pic->default_cxt.fp = NULL;
   pic->default_cxt.sp = NULL;
@@ -256,6 +257,7 @@ pic_close(pic_state *pic)
   pic->features = pic_invalid_value(pic);
   pic->dyn_env = pic_invalid_value(pic);
 
+  assert(pic->cxt->ai == 0);
   assert(pic->cxt->pc == NULL);
   assert(pic->cxt->fp == NULL);
   assert(pic->cxt->sp == NULL);
diff --git a/lib/state.h b/lib/state.h
index 60fd9907..a2d9f809 100644
--- a/lib/state.h
+++ b/lib/state.h
@@ -16,6 +16,7 @@ KHASH_DECLARE(oblist, struct string *, struct symbol *)
 
 struct context {
   PIC_JMPBUF jmp;
+  size_t ai;
 
   /* vm */
   const code_t *pc;