Add `indices(s)`, improve `index(s)`, `rindex(s)`

Now these deal with arrays as input and `s` being an array or a scalar.

6 files changed, 61 insertions, 4 deletions

diff --git a/builtin.c b/builtin.c
index ad4145a5..274576bc 100644
--- a/builtin.c
+++ b/builtin.c
@@ -688,8 +688,9 @@ static const char* const jq_builtins[] = {
   "def from_entries: map({(.key): .value}) | add;",
   "def with_entries(f): to_entries | map(f) | from_entries;",
   "def reverse: [.[length - 1 - range(0;length)]];",
-  "def index(i): .[i][0];",
-  "def rindex(i): .[i][-1:][0];",
+  "def indices(i): if type == \"array\" and (i|type) == \"array\" then .[i] elif type == \"array\" then .[[i]] else .[i] end;",
+  "def index(i):   if type == \"array\" and (i|type) == \"array\" then .[i] elif type == \"array\" then .[[i]] else .[i] end | .[0];",
+  "def rindex(i):  if type == \"array\" and (i|type) == \"array\" then .[i] elif type == \"array\" then .[[i]] else .[i] end | .[-1:][0];",
   "def paths: path(recurse(if (type|. == \"array\" or . == \"object\") then .[] else empty end))|select(length > 0);",
   "def leaf_paths: . as $dot|paths|select(. as $p|$dot|getpath($p)|type|. != \"array\" and . != \"object\");",
   "def any: reduce .[] as $i (false; . or $i);",
diff --git a/docs/content/3.manual/manual.yml b/docs/content/3.manual/manual.yml
index 5a77b5e1..0852f25b 100644
--- a/docs/content/3.manual/manual.yml
+++ b/docs/content/3.manual/manual.yml
@@ -998,6 +998,25 @@ sections:
             input: '{"foo": 12, "bar":[1,2,{"barp":12, "blip":13}]}'
             output: ['false']
 
+      - title: `indices(s)`
+        body: |
+
+          Outputs an array containing the indices in `.` where `s`
+          occurs.  The input may be an array, in which case if `s` is an
+          array then the indices output will be those where all elements
+          in `.` match those of `s`.
+
+        examples:
+          - program: 'indices(", ")'
+            input: '"a,b, cd, efg, hijk"'
+            output: ['[3,7,12]']
+          - program: 'indices(1)'
+            input: '[0,1,2,1,3,1,4]'
+            output: ['[1,3,5]']
+          - program: 'indices([1,2])'
+            input: '[0,1,2,3,1,4,2,5,1,2,6,7]'
+            output: ['[1,8]']
+
       - title: `index(s)`, `rindex(s)`
         body: |
 
diff --git a/jv.c b/jv.c
index 24971af6..2da8ecb0 100644
--- a/jv.c
+++ b/jv.c
@@ -364,6 +364,28 @@ int jv_array_contains(jv a, jv b) {
   return r;
 }
 
+jv jv_array_indexes(jv a, jv b) {
+  jv res = jv_array();
+  int idx = -1;
+  jv_array_foreach(a, ai, aelem) {
+    jv_array_foreach(b, bi, belem) {
+      // quieten compiler warnings about aelem not being used... by
+      // using it
+      if ((bi == 0 && !jv_equal(jv_copy(aelem), jv_copy(belem))) ||
+          (bi > 0 && !jv_equal(jv_array_get(jv_copy(a), ai + bi), jv_copy(belem))))
+        idx = -1;
+      else if (bi == 0 && idx == -1)
+        idx = ai;
+    }
+    if (idx > -1)
+      res = jv_array_append(res, jv_number(idx));
+    idx = -1;
+  }
+  jv_free(a);
+  jv_free(b);
+  return res;
+}
+
 
 /*
  * Strings (internal helpers)
diff --git a/jv.h b/jv.h
index 01c1bb62..b8087adb 100644
--- a/jv.h
+++ b/jv.h
@@ -71,6 +71,7 @@ jv jv_array_set(jv, int, jv);
 jv jv_array_append(jv, jv);
 jv jv_array_concat(jv, jv);
 jv jv_array_slice(jv, int, int);
+jv jv_array_indexes(jv, jv);
 #define jv_array_foreach(a, i, x) \
   for (int jv_len__ = jv_array_length(jv_copy(a)), i=0, jv_j__ = 1;     \
        jv_j__; jv_j__ = 0)                                              \
diff --git a/jv_aux.c b/jv_aux.c
index ec2aaa18..3c5750bd 100644
--- a/jv_aux.c
+++ b/jv_aux.c
@@ -85,6 +85,8 @@ jv jv_get(jv t, jv k) {
     }
   } else if (jv_get_kind(t) == JV_KIND_STRING && jv_get_kind(k) == JV_KIND_STRING) {
     v = jv_string_indexes(t, k);
+  } else if (jv_get_kind(t) == JV_KIND_ARRAY && jv_get_kind(k) == JV_KIND_ARRAY) {
+    v = jv_array_indexes(t, k);
   } else if (jv_get_kind(t) == JV_KIND_NULL && 
              (jv_get_kind(k) == JV_KIND_STRING || 
               jv_get_kind(k) == JV_KIND_NUMBER || 
diff --git a/tests/all.test b/tests/all.test
index 92a8b0e8..b45132ea 100644
--- a/tests/all.test
+++ b/tests/all.test
@@ -664,9 +664,21 @@ def inc(x): x |= .+1; inc(.[].a)
 "a,b,, c, d,ef, , ghi, jklmn, o"
 [4,7,13,15,20,27]
 
-[(index(","), rindex(","))]
+[(index(","), rindex(",")), indices(",")]
 "a,bc,def,ghij,klmno"
-[1,13]
+[1,13,[1,4,8,13]]
+
+indices(1)
+[0,1,1,2,3,4,1,5]
+[1,2,6]
+
+indices([1,2])
+[0,1,2,3,1,4,2,5,1,2,6,7]
+[1,8]
+
+indices(", ")
+"a,b, cd,e, fgh, ijkl"
+[3,9,14]
 
 [.[]|split(",")]
 ["a, bc, def, ghij, jklmn, a,b, c,d, e,f", "a,b,c,d, e,f,g,h"]