Implement indices natively. (#165)

jaq-core/src/lib.rs

@@ -232,6 +232,28 @@ fn to_csv(vs: &[Val]) -> Result<String, Error> {
     Ok(vs.iter().map(fr).collect::<Result<Vec<_>, _>>()?.join(","))
 }
 
+/// Return the indices of `y` in `x`.
+fn indices<'a>(x: &'a Val, y: &'a Val) -> Result<Box<dyn Iterator<Item = usize> + 'a>, Error> {
+    match (x, y) {
+        (Val::Str(_), Val::Str(y)) if y.is_empty() => Ok(Box::new(core::iter::empty())),
+        (Val::Arr(_), Val::Arr(y)) if y.is_empty() => Ok(Box::new(core::iter::empty())),
+        (Val::Str(x), Val::Str(y)) => {
+            let iw = x.as_bytes().windows(y.len()).enumerate();
+            let y = y.as_bytes();
+            Ok(Box::new(iw.filter_map(move |(i, w)| (w == y).then_some(i))))
+        }
+        (Val::Arr(x), Val::Arr(y)) => {
+            let iw = x.windows(y.len()).enumerate();
+            Ok(Box::new(iw.filter_map(|(i, w)| (w == **y).then_some(i))))
+        }
+        (Val::Arr(x), y) => {
+            let ix = x.iter().enumerate();
+            Ok(Box::new(ix.filter_map(move |(i, x)| (x == y).then_some(i))))
+        }
+        (x, y) => Err(Error::Index(x.clone(), y.clone())),
+    }
+}
+
 fn once_with<'a, T>(f: impl FnOnce() -> T + 'a) -> Box<dyn Iterator<Item = T> + 'a> {
     Box::new(core::iter::once_with(f))
 }
@@ -304,6 +326,12 @@ const CORE_RUN: &[(&str, usize, RunPtr)] = &[
         let vals = args.get(0).run(cv.clone());
         Box::new(vals.map(move |y| Ok(Val::Bool(cv.1.contains(&y?)))))
     }),
+    ("indices", 1, |args, cv| {
+        let vals = args.get(0).run(cv.clone());
+        let to_int = |i: usize| Val::Int(i.try_into().unwrap());
+        let f = move |v| indices(&cv.1, &v?).map(|idxs| Val::arr(idxs.map(to_int).collect()));
+        Box::new(vals.map(f))
+    }),
     ("split", 1, |args, cv| {
         let seps = args.get(0).run(cv.clone());
         Box::new(seps.map(move |sep| Ok(Val::arr(split(cv.1.as_str()?, sep?.as_str()?)))))

jaq-core/tests/tests.rs

@@ -157,6 +157,25 @@ fn has() {
     give(json!({"a": 1, "b": null}), r#"has("c")"#, json!(false));
 }
 
+yields!(indices_str, r#""a,b, cd, efg" | indices(", ")"#, [3, 7]);
+yields!(
+    indices_arr_num,
+    "[0, 1, 2, 1, 3, 1, 4] | indices(1)",
+    [1, 3, 5]
+);
+yields!(
+    indices_arr_arr,
+    "[0, 1, 2, 3, 1, 4, 2, 5, 1, 2, 6, 7] | indices([1, 2])",
+    [1, 8]
+);
+yields!(indices_arr_str, r#"["a", "b", "c"] | indices("b")"#, [1]);
+
+yields!(indices_arr_empty, "[0, 1] | indices([])", json!([]));
+yields!(indices_arr_larger, "[1, 2] | indices([1, 2, 3])", json!([]));
+
+yields!(indices_arr_overlap, "[0, 0, 0] | indices([0, 0])", [0, 1]);
+yields!(indices_str_overlap, r#""aaa" | indices("aa")"#, [0, 1]);
+
 #[test]
 fn json() {
     // TODO: correct this

jaq-std/src/std.jq

@@ -118,26 +118,8 @@ def in(xs)    : . as $x | xs | has     ($x);
 def inside(xs): . as $x | xs | contains($x);
 
 # Indexing
-def indices($i):
-  def enumerate:
-    . as $thing |
-    if type == "string"
-    then range(length) | [., $thing[.:.+1]]
-    else range(length) | [., $thing[.]]
-    end;
-
-  def windowed($size):
-    if $size <= 0 then empty
-    else . as $array | range(length - $size + 1) | $array[.:. + $size]
-    end;
-
-  if ($i | type) == "array" or (type == "string" and ($i | type) == "string")
-  then [[windowed($i | length)] | enumerate | select(.[1] == $i)[0]]
-  else [enumerate | select(.[1] == $i)[0]]
-  end;
-
-def index($i):  indices($i) | .[0];
-def rindex($i): indices($i) | .[-1:][0];
+def  index($i): indices($i)[ 0];
+def rindex($i): indices($i)[-1];
 
 # Walking
 def walk(f): def rec: (.[]? |= rec) | f; rec;

jaq-std/tests/std.rs

@@ -104,26 +104,6 @@ yields!(flatten_obj, "{a: 1} | flatten", json!([{"a": 1}]));
 // jq gives an error here
 yields!(flatten_num, "0 | flatten", [0]);
 
-#[test]
-fn indices() {
-    give(
-        json!("a,b, cd, efg, hijk"),
-        r#"indices(", ")"#,
-        json!([3, 7, 12]),
-    );
-    give(json!([0, 1, 2, 1, 3, 1, 4]), "indices(1)", json!([1, 3, 5]));
-    give(
-        json!([0, 1, 2, 3, 1, 4, 2, 5, 1, 2, 6, 7]),
-        "indices([1, 2])",
-        json!([1, 8]),
-    );
-    give(json!([]), "indices([])", json!([]));
-    give(json!([1, 2]), "indices([1, 2, 3])", json!([]));
-    give(json!(["a", "b", "c"]), r#"indices("b")"#, json!([1]));
-    give(json!([0, 0, 0]), "indices([0, 0])", json!([0, 1]));
-    give(json!("aaa"), r#"indices("aa")"#, json!([0, 1]));
-}
-
 #[test]
 fn inside() {
     give(