Implement array indexing

src/grammar.pest

@@ -18,8 +18,9 @@ char = {
 }
 
 union =  { "[" ~ unionElement ~ ("," ~ unionElement)* ~ "]" }
-unionElement = _{ unionChild }
+unionElement = _{ unionChild | unionArrayIndex } // TODO: add unionArraySlice
 unionChild = { doubleQuotedString | singleQuotedString }
+unionArrayIndex = { "-" ? ~ ( "0" | ASCII_NONZERO_DIGIT ~ ASCII_DIGIT* ) }
 
 doubleQuotedString = _{ "\"" ~ doubleInner ~ "\"" }
 doubleInner = @{ doubleChar* }

src/matchers.rs

@@ -58,6 +58,35 @@ impl Matcher for Child {
     }
 }
 
+/// Selects an array item by index.
+///
+/// If the index is negative, it references element len-abs(index).
+pub struct ArrayIndex {
+    index: i64,
+}
+
+impl ArrayIndex {
+    pub fn new(index: i64) -> Self {
+        ArrayIndex { index }
+    }
+}
+
+impl Matcher for ArrayIndex {
+    fn select<'a>(&self, node: &'a Value) -> Iter<'a> {
+        let idx = if self.index >= 0 {
+            self.index as usize
+        } else {
+            let len = if let Value::Array(a) = node {
+                a.len() as i64
+            } else {
+                0
+            };
+            (len + self.index) as usize
+        };
+        Box::new(node.get(idx).into_iter())
+    }
+}
+
 /// Applies a sequence of selectors on the same node and returns
 /// a concatenation of the results.
 pub struct Union {
@@ -96,4 +125,52 @@ mod tests {
         let r: Vec<&Value> = s.select(&j).collect();
         assert_eq!(format!("{:?}", r), "[Number(1), Number(2)]");
     }
+
+    #[test]
+    fn array_index() {
+        let s = ArrayIndex::new(1);
+        let j = json!([1, 2]);
+        let r: Vec<&Value> = s.select(&j).collect();
+        assert_eq!(format!("{:?}", r), "[Number(2)]");
+    }
+
+    #[test]
+    fn array_index_zero() {
+        let s = ArrayIndex::new(0);
+        let j = json!([1, 2]);
+        let r: Vec<&Value> = s.select(&j).collect();
+        assert_eq!(format!("{:?}", r), "[Number(1)]");
+    }
+
+    #[test]
+    fn array_index_oob() {
+        let s = ArrayIndex::new(4);
+        let j = json!([1, 2]);
+        let r: Vec<&Value> = s.select(&j).collect();
+        assert_eq!(r.len(), 0);
+    }
+
+    #[test]
+    fn array_index_negative() {
+        let s = ArrayIndex::new(-1);
+        let j = json!([1, 2]);
+        let r: Vec<&Value> = s.select(&j).collect();
+        assert_eq!(format!("{:?}", r), "[Number(2)]");
+    }
+
+    #[test]
+    fn array_index_negative_extreme() {
+        let s = ArrayIndex::new(-2);
+        let j = json!([1, 2]);
+        let r: Vec<&Value> = s.select(&j).collect();
+        assert_eq!(format!("{:?}", r), "[Number(1)]");
+    }
+
+    #[test]
+    fn array_index_negative_oob() {
+        let s = ArrayIndex::new(-10);
+        let j = json!([1, 2]);
+        let r: Vec<&Value> = s.select(&j).collect();
+        assert_eq!(r.len(), 0);
+    }
 }

src/parser.rs

@@ -75,10 +75,18 @@ fn parse_dot_child_matcher(
 fn parse_union(matcher_rule: pest::iterators::Pair<Rule>) -> Vec<Box<dyn matchers::Matcher>> {
     let mut ms: Vec<Box<dyn matchers::Matcher>> = Vec::new();
     for r in matcher_rule.into_inner() {
-        if let Rule::unionChild = r.as_rule() {
-            for m in parse_union_child(r) {
-                ms.push(m)
+        match r.as_rule() {
+            Rule::unionChild => {
+                for m in parse_union_child(r) {
+                    ms.push(m)
+                }
+            }
+            Rule::unionArrayIndex => {
+                for m in parse_union_array_index(r) {
+                    ms.push(m)
+                }
             }
+            _ => {}
         }
     }
     vec![Box::new(matchers::Union::new(ms))]
@@ -102,6 +110,15 @@ fn parse_union_child(matcher_rule: pest::iterators::Pair<Rule>) -> Vec<Box<dyn m
     ms
 }
 
+fn parse_union_array_index(
+    matcher_rule: pest::iterators::Pair<Rule>,
+) -> Vec<Box<dyn matchers::Matcher>> {
+    let mut ms: Vec<Box<dyn matchers::Matcher>> = Vec::new();
+    let i = matcher_rule.as_str().parse().unwrap();
+    ms.push(Box::new(matchers::ArrayIndex::new(i)));
+    ms
+}
+
 const ESCAPED: &str = "\"'\\/bfnrt";
 const UNESCAPED: &str = "\"'\\/\u{0008}\u{000C}\u{000A}\u{000D}\u{0009}";
 

tests/cts.json

@@ -461,5 +461,48 @@
     "name": "union child, single quotes, incomplete escape",
     "selector": "$['\\']",
     "invalid_selector": true
+  }, {
+    "name": "union array access",
+    "selector": "$[0]",
+    "document": ["first", "second"],
+    "result": ["first"]
+  }, {
+    "name": "union array access, 1",
+    "selector": "$[1]",
+    "document": ["first", "second"],
+    "result": ["second"]
+  }, {
+    "name": "union array access, out of bound",
+    "selector": "$[2]",
+    "document": ["first", "second"],
+    "result": []
+  }, {
+    "name": "union array access, negative",
+    "selector": "$[-1]",
+    "document": ["first", "second"],
+    "result": ["second"]
+  }, {
+    "name": "union array access, more negative",
+    "selector": "$[-2]",
+    "document": ["first", "second"],
+    "result": ["first"]
+  }, {
+    "name": "union array access, negative out of bound",
+    "selector": "$[-3]",
+    "document": ["first", "second"],
+    "result": []
+  }, {
+    "name": "union array access, on object",
+    "selector": "$[0]",
+    "document": {"foo": 1},
+    "result": []
+  }, {
+    "name": "union array access, leading 0",
+    "selector": "$[01]",
+    "invalid_selector": true
+  }, {
+    "name": "union array access, leading -0",
+    "selector": "$[-01]",
+    "invalid_selector": true
   }
-]}
+]}