Merge pull request #38 from akoshchiy/11270-keys-exists-api

feat: add exists_any_keys & exists_all_keys

Author: b41sh

src/functions.rs

@@ -16,6 +16,8 @@ use core::convert::TryInto;
 use std::borrow::Cow;
 use std::cmp::Ordering;
 use std::collections::VecDeque;
+use std::str::from_utf8;
+use std::str::from_utf8_unchecked;
 
 use crate::constants::*;
 use crate::error::*;
@@ -366,6 +368,120 @@ pub fn get_by_keypath<'a, I: Iterator<Item = &'a KeyPath<'a>>>(
         .map(|jentry| extract_by_jentry(&jentry, curr_jentry_encoded, curr_val_offset, value))
 }
 
+/// Checks whether all of the strings exist as top-level keys or array elements.
+pub fn exists_all_keys<'a, I: Iterator<Item = &'a [u8]>>(value: &[u8], keys: I) -> bool {
+    if !is_jsonb(value) {
+        match parse_value(value) {
+            Ok(val) => {
+                for key in keys {
+                    match from_utf8(key) {
+                        Ok(key) => {
+                            if !exists_value_key(&val, key) {
+                                return false;
+                            }
+                        }
+                        Err(_) => return false,
+                    }
+                }
+            }
+            Err(_) => return false,
+        };
+        return true;
+    }
+
+    let header = read_u32(value, 0).unwrap();
+
+    for key in keys {
+        match from_utf8(key) {
+            Ok(key) => {
+                if !exists_jsonb_key(value, header, key) {
+                    return false;
+                }
+            }
+            Err(_) => return false,
+        }
+    }
+    true
+}
+
+/// Checks whether any of the strings exist as top-level keys or array elements.
+pub fn exists_any_keys<'a, I: Iterator<Item = &'a [u8]>>(value: &[u8], keys: I) -> bool {
+    if !is_jsonb(value) {
+        if let Ok(val) = parse_value(value) {
+            for key in keys {
+                if let Ok(key) = from_utf8(key) {
+                    if exists_value_key(&val, key) {
+                        return true;
+                    }
+                }
+            }
+        }
+        return false;
+    }
+
+    let header = read_u32(value, 0).unwrap();
+
+    for key in keys {
+        if let Ok(key) = from_utf8(key) {
+            if exists_jsonb_key(value, header, key) {
+                return true;
+            }
+        }
+    }
+    false
+}
+
+fn exists_value_key(value: &Value, key: &str) -> bool {
+    match value {
+        Value::Array(arr) => {
+            let mut found = false;
+            for item in arr {
+                let matches = match item {
+                    Value::String(v) => key.eq(v),
+                    _ => false,
+                };
+                if matches {
+                    found = true;
+                    break;
+                }
+            }
+            found
+        }
+        Value::Object(obj) => obj.contains_key(key),
+        _ => false,
+    }
+}
+
+fn exists_jsonb_key(value: &[u8], header: u32, key: &str) -> bool {
+    match header & CONTAINER_HEADER_TYPE_MASK {
+        OBJECT_CONTAINER_TAG => {
+            let mut matches = false;
+            for obj_key in iteate_object_keys(value, header) {
+                if obj_key.eq(key) {
+                    matches = true;
+                    break;
+                }
+            }
+            matches
+        }
+        ARRAY_CONTAINER_TAG => {
+            let mut matches = false;
+            for (jentry, val) in iterate_array(value, header) {
+                if jentry.type_code != STRING_TAG {
+                    continue;
+                }
+                let val = unsafe { from_utf8_unchecked(val) };
+                if val.eq(key) {
+                    matches = true;
+                    break;
+                }
+            }
+            matches
+        }
+        _ => false,
+    }
+}
+
 fn get_jentry_by_name(
     value: &[u8],
     offset: usize,
@@ -1683,3 +1799,89 @@ fn read_u32(buf: &[u8], idx: usize) -> Result<u32, Error> {
         .unwrap();
     Ok(u32::from_be_bytes(bytes))
 }
+
+fn iterate_array(value: &[u8], header: u32) -> ArrayIterator<'_> {
+    let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
+    ArrayIterator {
+        value,
+        jentry_offset: 4,
+        val_offset: 4 * length + 4,
+        length,
+        idx: 0,
+    }
+}
+
+fn iteate_object_keys(value: &[u8], header: u32) -> ObjectKeyIterator<'_> {
+    let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
+    ObjectKeyIterator {
+        value,
+        jentry_offset: 4,
+        key_offset: 8 * length + 4,
+        length,
+        idx: 0,
+    }
+}
+
+struct ArrayIterator<'a> {
+    value: &'a [u8],
+    jentry_offset: usize,
+    val_offset: usize,
+    length: usize,
+    idx: usize,
+}
+
+impl<'a> Iterator for ArrayIterator<'a> {
+    type Item = (JEntry, &'a [u8]);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.idx >= self.length {
+            return None;
+        }
+        let encoded = read_u32(self.value, self.jentry_offset).unwrap();
+        let jentry = JEntry::decode_jentry(encoded);
+        let val_length = jentry.length as usize;
+
+        let item = (
+            jentry,
+            &self.value[self.val_offset..self.val_offset + val_length],
+        );
+
+        self.idx += 1;
+        self.val_offset += val_length;
+        self.jentry_offset += 4;
+
+        Some(item)
+    }
+}
+
+struct ObjectKeyIterator<'a> {
+    value: &'a [u8],
+    jentry_offset: usize,
+    key_offset: usize,
+    length: usize,
+    idx: usize,
+}
+
+impl<'a> Iterator for ObjectKeyIterator<'a> {
+    type Item = &'a str;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.idx >= self.length {
+            return None;
+        }
+
+        let encoded = read_u32(self.value, self.jentry_offset).unwrap();
+        let jentry = JEntry::decode_jentry(encoded);
+        let key_length = jentry.length as usize;
+
+        let key = unsafe {
+            from_utf8_unchecked(&self.value[self.key_offset..self.key_offset + key_length])
+        };
+
+        self.idx += 1;
+        self.key_offset += key_length;
+        self.jentry_offset += 4;
+
+        Some(key)
+    }
+}

tests/it/functions.rs

@@ -18,10 +18,11 @@ use std::collections::BTreeMap;
 
 use jsonb::{
     array_length, array_values, as_bool, as_null, as_number, as_str, build_array, build_object,
-    compare, convert_to_comparable, from_slice, get_by_index, get_by_keypath, get_by_name,
-    get_by_path, is_array, is_object, keypath::parse_key_paths, object_each, object_keys,
-    parse_value, path_exists, strip_nulls, to_bool, to_f64, to_i64, to_pretty_string, to_str,
-    to_string, to_u64, traverse_check_string, type_of, Number, Object, Value,
+    compare, convert_to_comparable, exists_all_keys, exists_any_keys, from_slice, get_by_index,
+    get_by_keypath, get_by_name, get_by_path, is_array, is_object, keypath::parse_key_paths,
+    object_each, object_keys, parse_value, path_exists, strip_nulls, to_bool, to_f64, to_i64,
+    to_pretty_string, to_str, to_string, to_u64, traverse_check_string, type_of, Number, Object,
+    Value,
 };
 
 use jsonb::jsonpath::parse_json_path;
@@ -1088,6 +1089,64 @@ fn test_get_by_keypath() {
     }
 }
 
+#[test]
+fn test_exists_all_keys() {
+    let sources = vec![
+        (r#"true"#, vec!["10", "20", "40"], false),
+        (r#"[]"#, vec!["10", "20", "40"], false),
+        (r#"{}"#, vec!["10", "20", "40"], false),
+        (r#"["10","20","30"]"#, vec!["10", "20", "40"], false),
+        (r#"["10","20","30"]"#, vec!["10", "20", "30"], true),
+        (r#"[10,20,30]"#, vec!["10", "20", "30"], false),
+        (r#"["10","20","30"]"#, vec!["10", "20", "20"], true),
+        (r#"{"a":1,"b":2,"c":3}"#, vec!["c", "b", "a"], true),
+        (r#"{"a":1,"b":2,"c":3}"#, vec!["a", "b", "a"], true),
+        (r#"{"a":1,"b":2,"c":3}"#, vec!["c", "f", "a"], false),
+    ];
+    for (json, keys, expected) in sources {
+        let keys = keys.iter().map(|k| k.as_bytes());
+        {
+            let json = parse_value(json.as_bytes()).unwrap().to_vec();
+            let result = exists_all_keys(&json, keys.clone());
+            assert_eq!(result, expected);
+        }
+        {
+            let json = json.as_bytes();
+            let result = exists_all_keys(json, keys.clone());
+            assert_eq!(result, expected);
+        }
+    }
+}
+
+#[test]
+fn test_exists_any_keys() {
+    let sources = vec![
+        (r#"true"#, vec!["10", "20", "40"], false),
+        (r#"[]"#, vec!["10", "20", "40"], false),
+        (r#"{}"#, vec!["10", "20", "40"], false),
+        (r#"[10,20,30]"#, vec!["10", "20", "30"], false),
+        (r#"["10","20","30"]"#, vec!["10", "20", "40"], true),
+        (r#"["10","20","30"]"#, vec!["10", "20", "30"], true),
+        (r#"["10","20","30"]"#, vec!["40", "50", "60"], false),
+        (r#"{"a":1,"b":2,"c":3}"#, vec!["c", "b", "a"], true),
+        (r#"{"a":1,"b":2,"c":3}"#, vec!["a", "b", "a"], true),
+        (r#"{"a":1,"b":2,"c":3}"#, vec!["z", "f", "x"], false),
+    ];
+    for (json, keys, expected) in sources {
+        let keys = keys.iter().map(|k| k.as_bytes());
+        {
+            let json = parse_value(json.as_bytes()).unwrap().to_vec();
+            let result = exists_any_keys(&json, keys.clone());
+            assert_eq!(result, expected);
+        }
+        {
+            let json = json.as_bytes();
+            let result = exists_any_keys(json, keys.clone());
+            assert_eq!(result, expected);
+        }
+    }
+}
+
 fn init_object<'a>(entries: Vec<(&str, Value<'a>)>) -> Value<'a> {
     let mut map = BTreeMap::new();
     for (key, val) in entries {