Remove the YIELD opcode (use RET instead)

3 files changed, 22 insertions, 18 deletions

diff --git a/compile.c b/compile.c
index de9c14c3..d305c5c8 100644
--- a/compile.c
+++ b/compile.c
@@ -507,13 +507,7 @@ static int compile(struct locfile* locations, struct bytecode* bc, block b) {
   int var_frame_idx = 0;
   bc->nsubfunctions = 0;
   errors += expand_call_arglist(locations, &b);
-  if (bc->parent) {
-    // functions should end in a return
-    b = BLOCK(b, gen_op_simple(RET));
-  } else {
-    // the toplevel should YIELD;BACKTRACK; when it finds an answer
-    b = BLOCK(b, gen_op_simple(YIELD), gen_op_simple(BACKTRACK));
-  }
+  b = BLOCK(b, gen_op_simple(RET));
   for (inst* curr = b.first; curr; curr = curr->next) {
     if (!curr->next) assert(curr == b.last);
     int length = opcode_describe(curr->op)->length;
diff --git a/execute.c b/execute.c
index b9d1d31d..f5bd033e 100644
--- a/execute.c
+++ b/execute.c
@@ -26,6 +26,7 @@ struct jq_state {
   jv path;
   int subexp_nest;
   int debug_trace_enabled;
+  int initial_execution;
 };
 
 typedef struct {
@@ -149,7 +150,8 @@ jv jq_next(jq_state *jq) {
   uint16_t* pc = stack_restore(jq);
   assert(pc);
 
-  int backtracking = 0;
+  int backtracking = !jq->initial_execution;
+  jq->initial_execution = 0;
   while (1) {
     uint16_t opcode = *pc;
 
@@ -465,13 +467,6 @@ jv jq_next(jq_state *jq) {
       pc += offset;
       break;
     }
-
-    case YIELD: {
-      jv value = stack_pop(jq);
-      stack_save(jq, pc);
-      stack_switch(jq);
-      return value;
-    }
       
     case CALL_BUILTIN: {
       int nargs = *pc++;
@@ -510,10 +505,25 @@ jv jq_next(jq_state *jq) {
     }
 
     case RET: {
-      pc = *frame_current_retaddr(&jq->frame_stk);
-      frame_pop(&jq->frame_stk);
+      uint16_t* retaddr = *frame_current_retaddr(&jq->frame_stk);
+      if (retaddr) {
+        // function return
+        pc = retaddr;
+        frame_pop(&jq->frame_stk);
+      } else {
+        // top-level return, yielding value
+        jv value = stack_pop(jq);
+        stack_save(jq, pc - 1);
+        stack_push(jq, jv_null());
+        stack_switch(jq);
+        return value;
+      }
       break;
     }
+    case ON_BACKTRACK(RET): {
+      // resumed after top-level return
+      goto do_backtrack;
+    }
     }
   }
 }
@@ -538,6 +548,7 @@ void jq_init(struct bytecode* bc, jv input, jq_state **jq, int flags) {
   } else {
     new_jq->debug_trace_enabled = 0;
   }
+  new_jq->initial_execution = 1;
   *jq = new_jq;
 }
 
diff --git a/opcode_list.h b/opcode_list.h
index 38206871..14450c51 100644
--- a/opcode_list.h
+++ b/opcode_list.h
@@ -5,7 +5,6 @@ OP(POP,   NONE,     1, 0)
 OP(LOADV, VARIABLE, 1, 1)
 OP(STOREV, VARIABLE, 1, 0)
 OP(INDEX, NONE,     2, 1)
-OP(YIELD, NONE, 1, 0)
 OP(EACH,  NONE,     1, 1)
 OP(FORK,  BRANCH,   0, 0)
 OP(JUMP,  BRANCH,   0, 0)