Add foreach and limit

Author: nicowilliams

builtin.c

@@ -963,6 +963,7 @@ static const char* const jq_builtins[] = {
   "     def _while: "
   "         if cond then ., (update | _while) else empty end; "
   "     _while;",
+  "def limit(n; exp): if n < 0 then exp else foreach exp as $item ([n, null]; if .[0] < 1 then empty else [.[0] -1, $item] end; .[1]) end;",
 };
 #undef LIBM_DD
 

compile.c

@@ -374,6 +374,41 @@ block gen_reduce(const char* varname, block source, block init, block body) {
                gen_op_bound(LOADVN, res_var));
 }
 
+block gen_foreach(const char* varname, block source, block init, block update, block extract) {
+  block output = gen_op_targetlater(JUMP);
+  block state_var = gen_op_var_fresh(STOREV, "foreach");
+  block loop = BLOCK(gen_op_simple(DUP),
+                     // get a value from the source expression:
+                     source,
+                     // bind the $varname to that value for all the code in
+                     // this block_bind() to see:
+                     block_bind(gen_op_unbound(STOREV, varname),
+                                // load the loop state variable
+                                BLOCK(gen_op_bound(LOADVN, state_var),
+                                      // generate updated state
+                                      update,
+                                      // save the updated state for value extraction
+                                      gen_op_simple(DUP),
+                                      // save new state
+                                      gen_op_bound(STOREV, state_var),
+                                      // extract an output...
+                                      extract,
+                                      // ...and output it
+                                      output),
+                                OP_HAS_VARIABLE));
+  block foreach = BLOCK(gen_op_simple(DUP),
+                        init,
+                        state_var,
+                        gen_op_target(FORK, loop),
+                        loop,
+                        // At this point `foreach`'s input will be on
+                        // top of the stack, and we don't want to output
+                        // it, so we backtrack.
+                        gen_op_simple(BACKTRACK));
+  inst_set_target(output, foreach);
+  return foreach;
+}
+
 block gen_definedor(block a, block b) {
   // var found := false
   block found_var = gen_op_var_fresh(STOREV, "found");

compile.h

@@ -32,6 +32,7 @@ block gen_subexp(block a);
 block gen_both(block a, block b);
 block gen_collect(block expr);
 block gen_reduce(const char* varname, block source, block init, block body);
+block gen_foreach(const char* varname, block source, block init, block update, block extract);
 block gen_definedor(block a, block b);
 block gen_condbranch(block iftrue, block iffalse);
 block gen_and(block a, block b);

docs/content/3.manual/manual.yml

@@ -1777,6 +1777,42 @@ sections:
             input: '[10,2,5,3]'
             output: ['20']
 
+      - title: `limit(n; exp)`
+        body: |
+
+          The `limit` function extracts up to `n` outputs from `exp`.
+
+        examples:
+          - program: '[limit(3;.[])]'
+            input: '[0,1,2,3,4,5,6,7,8,9]'
+            output: ['[0,1,2]']
+
+      - title: `foreach`
+        body: |
+
+          The `foreach` syntax is similar to `reduce`, but intended to
+          allow the construction of `limit` and reducers that produce
+          intermediate results (see example).
+
+          The form is `foreach EXP as $var (INIT; UPDATE; EXTRACT)`.
+          Like `reduce`, `INIT` is evaluated once to produce a state
+          value, then each output of `EXP` is bound to `$var`, `UPDATE`
+          is evaluated for each output of `EXP` with the current state
+          and with `$var` visible.  Each value output by `UPDATE`
+          replaces the previous state.  Finally, `EXTRACT` is evaluated
+          for each new state to extract an output of `foreach`.
+
+          This is mostly useful only for constructing `reduce`- and
+          `limit`-like functions.
+
+        examples:
+          - program: '[foreach .[] as $item
+                        ([[],[]];
+                        if $item == null then [[],.[0]] else [(.[0] + [$item]),[]] end;
+                        if $item == null then .[1] else empty end)]'
+            input: '[1,2,3,4,null,"a","b",null]'
+            output: ['[[1,2,3,4],["a","b"]]']
+
       - title: Recursion
         body: |
 

lexer.l

@@ -51,6 +51,7 @@ struct lexer_param;
 "or" { return OR; }
 "end" { return END; }
 "reduce" { return REDUCE; }
+"foreach" { return FOREACH; }
 "//" { return DEFINEDOR; }
 "try" { return TRY; }
 "catch" { return CATCH; }

parser.y

@@ -63,6 +63,7 @@ struct lexer_param;
 %token ELSE "else"
 %token ELSE_IF "elif"
 %token REDUCE "reduce"
+%token FOREACH "foreach"
 %token END "end"
 %token AND "and"
 %token OR "or"
@@ -234,6 +235,11 @@ Term "as" '$' IDENT '|' Exp {
   jv_free($5);
 } |
 
+"foreach" Term "as" '$' IDENT '(' Exp ';' Exp ';' Exp ')' {
+  $$ = gen_foreach(jv_string_value($5), $2, $7, $9, $11);
+  jv_free($5);
+} |
+
 "if" Exp "then" Exp ElseBody {
   $$ = gen_cond($2, $4, $5);
 } |

tests/all.test

@@ -226,6 +226,14 @@ null
 1
 [1,2,4,8,16,32,64]
 
+[foreach .[] as $item ([3, null]; if .[0] < 1 then empty else [.[0] -1, $item] end; .[1])]
+[11,22,33,44,55,66,77,88,99]
+[11,22,33]
+
+[limit(3; .[])]
+[11,22,33,44,55,66,77,88,99]
+[11,22,33]
+
 #
 # Slices
 #