Fixup API to get closer to a libjq

8 files changed, 203 insertions, 75 deletions

diff --git a/bytecode.c b/bytecode.c
index b5703a9a..184cf983 100644
--- a/bytecode.c
+++ b/bytecode.c
@@ -14,6 +14,22 @@ static int bytecode_operation_length(uint16_t* codeptr) {
   return length;
 }
 
+static void dump_code(int indent, struct bytecode* bc) {
+  int pc = 0;
+  while (pc < bc->codelen) {
+    printf("%*s", indent, "");
+    dump_operation(bc, bc->code + pc);
+    printf("\n");
+    pc += bytecode_operation_length(bc->code + pc);
+  }
+}
+
+static void symbol_table_free(struct symbol_table* syms) {
+  jv_mem_free(syms->cfunctions);
+  jv_free(syms->cfunc_names);
+  jv_mem_free(syms);
+}
+
 void dump_disassembly(int indent, struct bytecode* bc) {
   if (bc->nclosures > 0) {
     printf("%*s[params: ", indent, "");
@@ -37,16 +53,6 @@ void dump_disassembly(int indent, struct bytecode* bc) {
   }
 }
 
-void dump_code(int indent, struct bytecode* bc) {
-  int pc = 0;
-  while (pc < bc->codelen) {
-    printf("%*s", indent, "");
-    dump_operation(bc, bc->code + pc);
-    printf("\n");
-    pc += bytecode_operation_length(bc->code + pc);
-  }
-}
-
 static struct bytecode* getlevel(struct bytecode* bc, int level) {
   while (level > 0) {
     bc = bc->parent;
@@ -109,13 +115,9 @@ void dump_operation(struct bytecode* bc, uint16_t* codeptr) {
   }  
 }
 
-void symbol_table_free(struct symbol_table* syms) {
-  jv_mem_free(syms->cfunctions);
-  jv_free(syms->cfunc_names);
-  jv_mem_free(syms);
-}
-
 void bytecode_free(struct bytecode* bc) {
+  if (!bc)
+    return;
   jv_mem_free(bc->code);
   jv_free(bc->constants);
   for (int i=0; i<bc->nsubfunctions; i++)
diff --git a/bytecode.h b/bytecode.h
index c742cdd9..70929c91 100644
--- a/bytecode.h
+++ b/bytecode.h
@@ -39,10 +39,8 @@ struct bytecode {
 };
 
 void dump_disassembly(int, struct bytecode* code);
-void dump_code(int, struct bytecode* code);
 void dump_operation(struct bytecode* bc, uint16_t* op);
 
-void symbol_table_free(struct symbol_table* syms);
 void bytecode_free(struct bytecode* bc);
 
 #endif
diff --git a/execute.c b/execute.c
index 44cf2ba9..b5168b6c 100644
--- a/execute.c
+++ b/execute.c
@@ -23,6 +23,9 @@ struct jq_state {
   struct forkable_stack data_stk;
   struct forkable_stack frame_stk;
   struct forkable_stack fork_stk;
+  void (*nomem_handler)(void *);
+  void *nomem_handler_data;
+  struct bytecode* bc;
   jv path;
   int subexp_nest;
   int debug_trace_enabled;
@@ -118,6 +121,17 @@ uint16_t* stack_restore(jq_state *jq){
   return retaddr;
 }
 
+static void jq_reset(jq_state *jq) {
+  while (stack_restore(jq)) {}
+  assert(forkable_stack_empty(&jq->fork_stk));
+  assert(forkable_stack_empty(&jq->data_stk));
+  assert(forkable_stack_empty(&jq->frame_stk));
+
+  if (jv_get_kind(jq->path) != JV_KIND_INVALID)
+    jv_free(jq->path);
+  jq->path = jv_null();
+}
+
 
 static struct closure make_closure(struct forkable_stack* stk, frame_ptr fr, uint16_t* pc) {
   uint16_t level = *pc++;
@@ -147,6 +161,8 @@ void print_error(jv value) {
 jv jq_next(jq_state *jq) {
   jv cfunc_input[MAX_CFUNCTION_ARGS];
 
+  jv_nomem_handler(jq->nomem_handler, jq->nomem_handler_data);
+
   uint16_t* pc = stack_restore(jq);
   assert(pc);
 
@@ -552,28 +568,41 @@ jv jq_next(jq_state *jq) {
   }
 }
 
+jq_state *jq_init(void) {
+  jq_state *jq;
+  jq = jv_mem_alloc_unguarded(sizeof(*jq));
+  if (jq == NULL)
+    return NULL;
+  memset(jq, 0, sizeof(*jq));
+  jq->debug_trace_enabled = 0;
+  jq->initial_execution = 1;
+  forkable_stack_init(&jq->data_stk, sizeof(data_stk_elem) * 100);
+  forkable_stack_init(&jq->frame_stk, 1024);
+  forkable_stack_init(&jq->fork_stk, 1024);
+  jq->path = jv_null();
+  return jq;
+}
+
+void jq_set_nomem_handler(jq_state *jq, void (*nomem_handler)(void *), void *data) {
+  jv_nomem_handler(nomem_handler, data);
+  jq->nomem_handler = nomem_handler;
+  jq->nomem_handler_data = data;
+}
 
-void jq_init(struct bytecode* bc, jv input, jq_state **jq, int flags) {
-  jq_state *new_jq;
-  new_jq = jv_mem_alloc(sizeof(*new_jq));
-  memset(new_jq, 0, sizeof(*new_jq));
-  new_jq->path = jv_null();
-  forkable_stack_init(&new_jq->data_stk, sizeof(data_stk_elem) * 100);
-  forkable_stack_init(&new_jq->frame_stk, 1024);
-  forkable_stack_init(&new_jq->fork_stk, 1024);
-  
-  stack_push(new_jq, input);
-  struct closure top = {bc, -1};
-  frame_push(&new_jq->frame_stk, top, 0);
-  stack_save(new_jq, bc->code);
-  stack_switch(new_jq);
+void jq_start(jq_state *jq, jv input, int flags) {
+  jv_nomem_handler(jq->nomem_handler, jq->nomem_handler_data);
+  jq_reset(jq);
+  stack_push(jq, input);
+  struct closure top = {jq->bc, -1};
+  frame_push(&jq->frame_stk, top, 0);
+  stack_save(jq, jq->bc->code);
+  stack_switch(jq);
   if (flags & JQ_DEBUG_TRACE) {
-    new_jq->debug_trace_enabled = 1;
+    jq->debug_trace_enabled = 1;
   } else {
-    new_jq->debug_trace_enabled = 0;
+    jq->debug_trace_enabled = 0;
   }
-  new_jq->initial_execution = 1;
-  *jq = new_jq;
+  jq->initial_execution = 1;
 }
 
 void jq_teardown(jq_state **jq) {
@@ -582,25 +611,26 @@ void jq_teardown(jq_state **jq) {
     return;
   *jq = NULL;
 
-  while (stack_restore(old_jq)) {}
-
-  assert(forkable_stack_empty(&old_jq->fork_stk));
-  assert(forkable_stack_empty(&old_jq->data_stk));
-  assert(forkable_stack_empty(&old_jq->frame_stk));
+  jq_reset(old_jq);
   forkable_stack_free(&old_jq->fork_stk);
   forkable_stack_free(&old_jq->data_stk);
   forkable_stack_free(&old_jq->frame_stk);
-
-  jv_free(old_jq->path);
+  bytecode_free(old_jq->bc);
+  old_jq->bc = 0;
   jv_mem_free(old_jq);
 }
 
-struct bytecode* jq_compile_args(const char* str, jv args) {
+int jq_compile_args(jq_state *jq, const char* str, jv args) {
+  jv_nomem_handler(jq->nomem_handler, jq->nomem_handler_data);
   assert(jv_get_kind(args) == JV_KIND_ARRAY);
   struct locfile locations;
   locfile_init(&locations, str, strlen(str));
   block program;
-  struct bytecode* bc = 0;
+  jq_reset(jq);
+  if (jq->bc) {
+    bytecode_free(jq->bc);
+    jq->bc = 0;
+  }
   int nerrors = jq_parse(&locations, &program);
   if (nerrors == 0) {
     for (int i=0; i<jv_array_length(jv_copy(args)); i++) {
@@ -613,16 +643,20 @@ struct bytecode* jq_compile_args(const char* str, jv args) {
     jv_free(args);
     nerrors = builtins_bind(&program);
     if (nerrors == 0) {
-      nerrors = block_compile(program, &locations, &bc);
+      nerrors = block_compile(program, &locations, &jq->bc);
     }
   }
   if (nerrors) {
     fprintf(stderr, "%d compile %s\n", nerrors, nerrors > 1 ? "errors" : "error");
   }
   locfile_free(&locations);
-  return bc;
+  return jq->bc != NULL;
+}
+
+int jq_compile(jq_state *jq, const char* str) {
+  return jq_compile_args(jq, str, jv_array());
 }
 
-struct bytecode* jq_compile(const char* str) {
-  return jq_compile_args(str, jv_array());
+void jq_dump_disassembly(jq_state *jq, int indent) {
+  dump_disassembly(indent, jq->bc);
 }
diff --git a/execute.h b/execute.h
index 004b66ba..0cb0b36c 100644
--- a/execute.h
+++ b/execute.h
@@ -2,19 +2,20 @@
 #define EXECUTE_H
 #include "bytecode.h"
 
+typedef struct jq_state jq_state;
 
-
-struct bytecode* jq_compile(const char* str);
+int jq_compile(jq_state *, const char* str);
 
 /* args must be an array of the form [{name:"foo", value:"thing"}, {name:"bar",value:3}] */
-struct bytecode* jq_compile_args(const char* str, jv args);
+int jq_compile_args(jq_state *, const char* str, jv args);
 
-typedef struct jq_state jq_state;
 enum {JQ_DEBUG_TRACE = 1};
 
-void jq_init(struct bytecode* bc, jv value, jq_state **, int flags);
+jq_state *jq_init(void);
+void jq_set_nomem_handler(jq_state *, void (*)(void *), void *);
+void jq_start(jq_state *, jv value, int flags);
 jv jq_next(jq_state *);
 void jq_teardown(jq_state **);
-
+void jq_dump_disassembly(jq_state *, int);
 
 #endif
diff --git a/jq_test.c b/jq_test.c
index 30124a8b..199a3bc4 100644
--- a/jq_test.c
+++ b/jq_test.c
@@ -27,6 +27,9 @@ static void run_jq_tests(FILE *testdata) {
   int tests = 0, passed = 0, invalid = 0;
   jq_state *jq = NULL;
 
+  jq = jq_init();
+  assert(jq);
+
   while (1) {
     if (!fgets(buf, sizeof(buf), testdata)) break;
     if (skipline(buf)) continue;
@@ -34,15 +37,15 @@ static void run_jq_tests(FILE *testdata) {
     printf("Testing %s\n", buf);
     int pass = 1;
     tests++;
-    struct bytecode* bc = jq_compile(buf);
-    if (!bc) {invalid++; continue;}
+    int compiled = jq_compile(jq, buf);
+    if (!compiled) {invalid++; continue;}
     printf("Disassembly:\n");
-    dump_disassembly(2, bc);
+    jq_dump_disassembly(jq, 2);
     printf("\n");
     if (!fgets(buf, sizeof(buf), testdata)) { invalid++; break; }
     jv input = jv_parse(buf);
     if (!jv_is_valid(input)){ invalid++; continue; }
-    jq_init(bc, input, &jq, JQ_DEBUG_TRACE);
+    jq_start(jq, input, JQ_DEBUG_TRACE);
 
     while (fgets(buf, sizeof(buf), testdata)) {
       if (skipline(buf)) break;
@@ -81,10 +84,9 @@ static void run_jq_tests(FILE *testdata) {
         jv_free(extra);
       }
     }
-    jq_teardown(&jq);
-    bytecode_free(bc);
     passed+=pass;
   }
+  jq_teardown(&jq);
   printf("%d of %d tests passed (%d malformed)\n", passed,tests,invalid);
   if (passed != tests) exit(1);
 }
diff --git a/jv_alloc.c b/jv_alloc.c
index 4fd36a60..03e9a7ce 100644
--- a/jv_alloc.c
+++ b/jv_alloc.c
@@ -2,10 +2,88 @@
 #include <stdio.h>
 #include "jv_alloc.h"
 
+struct nomem_handler {
+    jv_nomem_handler_f handler;
+    void *data;
+};
+
+#ifndef USE_PTHREAD_KEY
+#ifdef _MSC_VER
+static __declspec(thread) struct nomem_handler nomem_handler;
+#else
+static __thread struct nomem_handler nomem_handler;
+#endif
+
+void jv_nomem_handler(jv_nomem_handler_f handler, void *data) {
+  nomem_handler.handler = handler;
+}
+
 static void memory_exhausted() {
+  if (nomem_handler.handler)
+    nomem_handler.handler(nomem_handler.data); // Maybe handler() will longjmp() to safety
+  // Or not
   fprintf(stderr, "error: cannot allocate memory\n");
   abort();
 }
+#else
+#ifdef HAVE_PTHREAD_KEY_CREATE
+#include <pthread.h>
+
+pthread_key_t nomem_handler_key;
+pthread_once_t mem_once = PTHREAD_ONCE_INIT;
+
+static void tsd_init(void) {
+  if (pthread_key_create(&nomem_handler_key, NULL) != 0) {
+    fprintf(stderr, "error: cannot create thread specific key");
+    abort();
+  }
+}
+
+void jv_nomem_handler(jv_nomem_handler_f handler, void *data) {
+  pthread_once(&mem_once, tsd_init); // cannot fail
+  struct nomem_handler *nomem_handler = calloc(1, sizeof (nomem_handler));
+  if (nomem_handler == NULL) {
+    handler(data);
+    fprintf(stderr, "error: cannot allocate memory\n");
+    abort();
+  }
+  nomem_handler.handler = handler;
+  nomem_handler.data = data;
+  if (pthread_setspecific(nomem_handler_key, nomem_handler) != 0) {
+    handler(data);
+    fprintf(stderr, "error: cannot set thread specific data");
+    abort();
+  }
+}
+
+static void memory_exhausted() {
+  jv_nomem_handler_f handler;
+
+  pthread_once(&mem_once, tsd_init);
+  handler = pthread_getspecific(nomem_handler_key);
+  if (handler)
+    handler(); // Maybe handler() will longjmp() to safety
+  // Or not
+  fprintf(stderr, "error: cannot allocate memory\n");
+  abort();
+}
+
+#else
+
+static struct nomem_handler nomem_handler;
+void jv_nomem_handler(jv_nomem_handler_f handler, void *data) {
+  nomem_handler.handler = handler;
+  nomem_handler.data = data;
+}
+
+static void memory_exhausted() {
+  fprintf(stderr, "error: cannot allocate memory\n");
+  abort();
+}
+
+#endif /* HAVE_PTHREAD_KEY_CREATE */
+#endif /* !USE_PTHREAD_KEY */
+
 
 void* jv_mem_alloc(size_t sz) {
   void* p = malloc(sz);
@@ -15,6 +93,10 @@ void* jv_mem_alloc(size_t sz) {
   return p;
 }
 
+void* jv_mem_alloc_unguarded(size_t sz) {
+  return malloc(sz);
+}
+
 void jv_mem_free(void* p) {
   free(p);
 }
diff --git a/jv_alloc.h b/jv_alloc.h
index 9cf1d34e..cbcf9548 100644
--- a/jv_alloc.h
+++ b/jv_alloc.h
@@ -14,7 +14,10 @@ static void jv_mem_invalidate(void* mem, size_t n) {
 #endif
 }
 
+typedef void (*jv_nomem_handler_f)(void *);
+void jv_nomem_handler(jv_nomem_handler_f, void *);
 void* jv_mem_alloc(size_t);
+void* jv_mem_alloc_unguarded(size_t);
 void jv_mem_free(void*);
 __attribute__((warn_unused_result)) void* jv_mem_realloc(void*, size_t);
 
diff --git a/main.c b/main.c
index 221ddba0..02399641 100644
--- a/main.c
+++ b/main.c
@@ -65,11 +65,9 @@ enum {
   DUMP_DISASM = 2048,
 };
 static int options = 0;
-static struct bytecode* bc;
 
-static void process(jv value, int flags) {
-  jq_state *jq = NULL;
-  jq_init(bc, value, &jq, flags);
+static void process(jq_state *jq, jv value, int flags) {
+  jq_start(jq, value, flags);
   jv result;
   while (jv_is_valid(result = jq_next(jq))) {
     if ((options & RAW_OUTPUT) && jv_get_kind(result) == JV_KIND_STRING) {
@@ -94,7 +92,6 @@ static void process(jv value, int flags) {
     printf("\n");
   }
   jv_free(result);
-  jq_teardown(&jq);
 }
 
 FILE* current_input;
@@ -126,13 +123,22 @@ static int read_more(char* buf, size_t size) {
 }
 
 int main(int argc, char* argv[]) {
+  jq_state *jq;
   int ret = 0;
+  int compiled = 0;
+
   if (argc) progname = argv[0];
 
   if (argc > 1 && !strcmp(argv[1], "--run-tests")) {
     return jq_testsuite(argc - 1, argv + 1);
   }
 
+  jq = jq_init();
+  if (jq == NULL) {
+    perror("malloc");
+    return 1;
+  }
+
   const char* program = 0;
   input_filenames = jv_mem_alloc(sizeof(const char*) * argc);
   ninput_files = 0;
@@ -231,20 +237,20 @@ int main(int argc, char* argv[]) {
       jv_free(data);
       return 1;
     }
-    bc = jq_compile_args(jv_string_value(data), program_arguments);
+    compiled = jq_compile_args(jq, jv_string_value(data), program_arguments);
     jv_free(data);
   } else {
-    bc = jq_compile_args(program, program_arguments);
+    compiled = jq_compile_args(jq, program, program_arguments);
   }
-  if (!bc) return 1;
+  if (!compiled) return 1;
 
   if (options & DUMP_DISASM) {
-    dump_disassembly(0, bc);
+    jq_dump_disassembly(jq, 0);
     printf("\n");
   }
 
   if (options & PROVIDE_NULL) {
-    process(jv_null(), jq_flags);
+    process(jq, jv_null(), jq_flags);
   } else {
     jv slurped;
     if (options & SLURP) {
@@ -265,7 +271,7 @@ int main(int argc, char* argv[]) {
             slurped = jv_string_concat(slurped, jv_string(buf));
           } else {
             if (buf[len-1] == '\n') buf[len-1] = 0;
-            process(jv_string(buf), jq_flags);
+            process(jq, jv_string(buf), jq_flags);
           }
         }
       } else {
@@ -275,7 +281,7 @@ int main(int argc, char* argv[]) {
           if (options & SLURP) {
             slurped = jv_array_append(slurped, value);
           } else {
-            process(value, jq_flags);
+            process(jq, value, jq_flags);
           }
         }
         if (jv_invalid_has_msg(jv_copy(value))) {
@@ -293,11 +299,11 @@ int main(int argc, char* argv[]) {
     if (ret != 0)
       goto out;
     if (options & SLURP) {
-      process(slurped, jq_flags);
+      process(jq, slurped, jq_flags);
     }
   }
 out:
   jv_mem_free(input_filenames);
-  bytecode_free(bc);
+  jq_teardown(&jq);
   return ret;
 }