if-then-else and defined-or operators

7 files changed, 97 insertions, 3 deletions

diff --git a/c/compile.c b/c/compile.c
index b7fcc4f2..c83a5583 100644
--- a/c/compile.c
+++ b/c/compile.c
@@ -312,8 +312,51 @@ block gen_assign(block expr) {
   return c;
 }
 
-block gen_else(block a, block b) {
-  assert(0);
+block gen_definedor(block a, block b) {
+  // var found := false
+  block c = gen_op_simple(DUP);
+  block_append(&c, gen_op_const(LOADK, jv_false()));
+  block found_var = block_bind(gen_op_var_unbound(STOREV, "found"),
+                               gen_noop(), OP_HAS_VARIABLE);
+  block_append(&c, found_var);
+
+  // if found, backtrack. Otherwise execute b
+  block tail = gen_op_simple(DUP);
+  block_append(&tail, gen_op_var_bound(LOADV, found_var));
+  block backtrack = gen_op_simple(BACKTRACK);
+  block_append(&tail, gen_op_target(JUMP_F, backtrack));
+  block_append(&tail, backtrack);
+  block_append(&tail, gen_op_simple(POP));
+  block_append(&tail, b);
+
+  // try again
+  block if_notfound = gen_op_simple(BACKTRACK);
+
+  // found := true, produce result
+  block if_found = gen_op_simple(DUP);
+  block_append(&if_found, gen_op_const(LOADK, jv_true()));
+  block_append(&if_found, gen_op_var_bound(STOREV, found_var));
+  block_append(&if_found, gen_op_target(JUMP, tail));
+
+  block_append(&c, gen_op_target(FORK, if_notfound));
+  block_append(&c, a);
+  block_append(&c, gen_op_target(JUMP_F, if_found));
+  block_append(&c, if_found);
+  block_append(&c, if_notfound);
+  block_append(&c, tail);
+
+  return c;
+}
+
+block gen_cond(block cond, block iftrue, block iffalse) {
+  block b = gen_op_simple(DUP);
+  block_append(&b, cond);
+
+  block_append(&iftrue, gen_op_target(JUMP, iffalse));
+  block_append(&b, gen_op_target(JUMP_F, iftrue));
+  block_append(&b, block_join(gen_op_simple(POP), iftrue));
+  block_append(&b, block_join(gen_op_simple(POP), iffalse));
+  return b;
 }
 
 block gen_cbinding(struct symbol_table* t, block code) {
diff --git a/c/compile.h b/c/compile.h
index 5f0be89b..5a2faf7b 100644
--- a/c/compile.h
+++ b/c/compile.h
@@ -26,7 +26,9 @@ block gen_subexp(block a);
 block gen_both(block a, block b);
 block gen_collect(block expr);
 block gen_assign(block expr);
-block gen_else(block a, block b);
+block gen_definedor(block a, block b);
+
+block gen_cond(block cond, block iftrue, block iffalse);
 
 block gen_cbinding(struct symbol_table* functions, block b);
 
diff --git a/c/execute.c b/c/execute.c
index b7818d12..5e6a2c51 100644
--- a/c/execute.c
+++ b/c/execute.c
@@ -59,6 +59,7 @@ typedef struct {
 } data_stk_elem;
 
 void stack_push(stackval val) {
+  assert(jv_is_valid(val.value));
   data_stk_elem* s = forkable_stack_push(&data_stk, sizeof(data_stk_elem));
   s->sv = val;
 }
@@ -70,6 +71,7 @@ stackval stack_pop() {
     sv.value = jv_copy(sv.value);
   }
   forkable_stack_pop(&data_stk);
+  assert(jv_is_valid(sv.value));
   return sv;
 }
 
@@ -292,6 +294,17 @@ jv jq_next() {
       break;
     }
 
+    case JUMP_F: {
+      uint16_t offset = *pc++;
+      stackval t = stack_pop();
+      jv_kind kind = jv_get_kind(t.value);
+      if (kind == JV_KIND_FALSE || kind == JV_KIND_NULL) {
+        pc += offset;
+      }
+      stack_push(t); // FIXME do this better
+      break;
+    }
+
     case EACH: 
       stack_push(stackval_root(jv_number(0)));
       // fallthrough
diff --git a/c/lexer.l b/c/lexer.l
index 4acb026a..e3c140b6 100644
--- a/c/lexer.l
+++ b/c/lexer.l
@@ -13,6 +13,11 @@
 "as" { return AS; }
 "def" { return DEF; }
 "|=" { return SETPIPE; }
+"if" { return IF; }
+"then" { return THEN; }
+"else" { return ELSE; }
+"end" { return END; }
+"//" { return DEFINEDOR; }
 "."|"="|";"|"["|"]"|","|":"|"("|")"|"{"|"}"|"|"|"+"|"\$" { return yytext[0];}
 
 \"(\\.|[^\\"])*\" |
diff --git a/c/opcode_list.h b/c/opcode_list.h
index 7f2c144a..30e7203d 100644
--- a/c/opcode_list.h
+++ b/c/opcode_list.h
@@ -10,6 +10,7 @@ OP(YIELD, NONE, 1, 0)
 OP(EACH,  NONE,     1, 1)
 OP(FORK,  BRANCH,   0, 0)
 OP(JUMP,  BRANCH,   0, 0)
+OP(JUMP_F,BRANCH,   1, 0)
 OP(BACKTRACK, NONE, 0, 0)
 OP(APPEND, NONE,    2, 1)
 OP(INSERT, NONE,    4, 2)
diff --git a/c/parser.y b/c/parser.y
index 36292453..d4e8d8f4 100644
--- a/c/parser.y
+++ b/c/parser.y
@@ -25,13 +25,20 @@
 %left '|'
 %left ','
 %token EQ "=="
+%token DEFINEDOR "//"
 %token AS "as"
 %token DEF "def"
 %token SETPIPE "|="
+%token IF "if"
+%token THEN "then"
+%token ELSE "else"
+%token END "end"
+%right "//"
 %nonassoc '=' SETPIPE
 %nonassoc EQ
 %left '+'
 
+
 %type <blk> Exp Term MkDict MkDictPair ExpD
 
 %{
@@ -78,6 +85,10 @@ Term "as" '$' IDENT '|' Exp {
   jv_free($4);
 } |
 
+"if" Exp "then" Exp "else" Exp "end" {
+  $$ = gen_cond($2, $4, $6);
+} |
+
 Exp '=' Exp {
   block assign = gen_op_simple(DUP);
   block_append(&assign, $3);
@@ -87,6 +98,10 @@ Exp '=' Exp {
   $$ = gen_assign(assign);
 } |
 
+Exp "//" Exp {
+  $$ = gen_definedor($1, $3);
+} |
+
 Exp "|=" Exp {
   block assign = $1;
   block_append(&assign, gen_op_simple(DUP));
diff --git a/c/testdata b/c/testdata
index 6ddf09fd..888d8b74 100644
--- a/c/testdata
+++ b/c/testdata
@@ -234,3 +234,18 @@ def inc(x): x |= .+1; inc(.[].a)
 .foo[2].bar = 1
 {"foo":[11], "bar":42}
 {"foo":[11,null,{"bar":1}], "bar":42}
+
+#
+# Conditionals
+#
+
+[.[] | if .foo then "yep" else "nope" end]
+[{"foo":0},{"foo":1},{"foo":[]},{"foo":true},{"foo":false},{"foo":null},{"foo":"foo"},{}]
+["yep","yep","yep","yep","nope","nope","yep","nope"]
+
+# FIXME: define/test behaviour of 'if (.foo,.bar) then A else B end'
+
+[.[] | [.foo[] // .bar]]
+[{"foo":[1,2], "bar": 42}, {"foo":[1], "bar": null}, {"foo":[null,false,3], "bar": 18}, {"foo":[], "bar":42}, {"foo": [null,false,null], "bar": 41}]
+[[1,2], [1], [3], [42], [41]]
+