diff --git a/src/functions.rs b/src/functions.rs index 49cb824..cdbaad1 100644 --- a/src/functions.rs +++ b/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>>( .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>(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>(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 { .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 { + 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 { + 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) + } +} diff --git a/tests/it/functions.rs b/tests/it/functions.rs index bbf4100..d3f4a2b 100644 --- a/tests/it/functions.rs +++ b/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 {