Merge pull request #25 from b41sh/feat-traverse_check_string

Feat: Support traverse_check_string function

Author: sundy-li

src/functions.rs

@@ -1275,6 +1275,59 @@ fn rand_scalar_value() -> Value<'static> {
     val
 }
 
+/// Traverse all the string fields in a jsonb value and check whether the conditions are met.
+pub fn traverse_check_string(value: &[u8], func: impl Fn(&[u8]) -> bool) -> bool {
+    if !is_jsonb(value) {
+        match parse_value(value) {
+            Ok(val) => {
+                let val_buf = val.to_vec();
+                return traverse_check_string(&val_buf, func);
+            }
+            Err(_) => {
+                return false;
+            }
+        }
+    }
+
+    let mut offsets = VecDeque::new();
+    offsets.push_back(0);
+
+    while let Some(offset) = offsets.pop_front() {
+        let header = read_u32(value, offset).unwrap();
+        let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
+
+        let size = match header & CONTAINER_HEADER_TYPE_MASK {
+            SCALAR_CONTAINER_TAG => 1,
+            ARRAY_CONTAINER_TAG => length,
+            OBJECT_CONTAINER_TAG => length * 2,
+            _ => unreachable!("invalid jsonb value"),
+        };
+
+        let mut jentry_offset = offset + 4;
+        let mut val_offset = offset + 4 + 4 * size;
+        for _ in 0..size {
+            let encoded = read_u32(value, jentry_offset).unwrap();
+            let jentry = JEntry::decode_jentry(encoded);
+            match jentry.type_code {
+                CONTAINER_TAG => {
+                    offsets.push_back(val_offset);
+                }
+                STRING_TAG => {
+                    let val_length = jentry.length as usize;
+                    if func(&value[val_offset..val_offset + val_length]) {
+                        return true;
+                    }
+                }
+                _ => {}
+            }
+            jentry_offset += 4;
+            val_offset += jentry.length as usize;
+        }
+    }
+
+    false
+}
+
 // Check whether the value is `JSONB` format,
 // for compatibility with previous `JSON` string.
 fn is_jsonb(value: &[u8]) -> bool {

tests/it/functions.rs

@@ -19,7 +19,7 @@ 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_name, get_by_path, is_array,
     is_object, object_keys, parse_value, to_bool, to_f64, to_i64, to_str, to_string, to_u64,
-    Number, Object, Value,
+    traverse_check_string, Number, Object, Value,
 };
 
 use jsonb::jsonpath::parse_json_path;
@@ -781,3 +781,39 @@ fn test_to_string() {
         buf.clear();
     }
 }
+
+#[test]
+fn test_traverse_check_string() {
+    let sources = vec![
+        (r#"null"#, false),
+        (r#"11"#, false),
+        (r#""a""#, false),
+        (r#""c""#, true),
+        (r#"[1,2,3,4,"b"]"#, false),
+        (r#"[1,2,[3,[4,"c"]]]"#, true),
+        (r#"[true,false,[1,2,3],{"a":"b"}]"#, false),
+        (
+            r#"{"a":true,"b":{"b1":"v1","b2":11},"c":[true,12,"c1","c2"]}"#,
+            true,
+        ),
+        (
+            r#"{"a0":true,"b0":{"b1":"v1","b2":11},"c0":[true,12,"c1","c"]}"#,
+            true,
+        ),
+        (
+            r#"{"a0":true,"b0":{"b1":"v1","b2":11},"c0":[true,12,"c1","c2"]}"#,
+            false,
+        ),
+    ];
+    let mut buf: Vec<u8> = Vec::new();
+    for (s, expect) in sources {
+        let value = parse_value(s.as_bytes()).unwrap();
+        value.write_to_vec(&mut buf);
+        let res = traverse_check_string(&buf, |v| {
+            let s = unsafe { std::str::from_utf8_unchecked(v) };
+            s == "c"
+        });
+        assert_eq!(res, expect);
+        buf.clear();
+    }
+}