Skip to content

feat: add get_by_keypath #36

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Oct 19, 2023
Merged

feat: add get_by_keypath #36

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3d52e63
feat: add get_by_keypath
akoshchiy Oct 17, 2023
0ec23ed
review fixes
akoshchiy Oct 18, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
278 changes: 174 additions & 104 deletions src/functions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ use core::convert::TryInto;
use std::borrow::Cow;
use std::cmp::Ordering;
use std::collections::VecDeque;
use std::str::from_utf8;

use crate::constants::*;
use crate::error::*;
Expand Down Expand Up @@ -232,36 +233,9 @@ pub fn get_by_index(value: &[u8], index: usize) -> Option<Vec<u8>> {
let header = read_u32(value, 0).unwrap();
match header & CONTAINER_HEADER_TYPE_MASK {
ARRAY_CONTAINER_TAG => {
let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
if index >= length {
return None;
}
let mut jentry_offset = 4;
let mut val_offset = 4 * length + 4;
for i in 0..length {
let encoded = read_u32(value, jentry_offset).unwrap();
let jentry = JEntry::decode_jentry(encoded);
let val_length = jentry.length as usize;
if i < index {
jentry_offset += 4;
val_offset += val_length;
continue;
}
let val = match jentry.type_code {
CONTAINER_TAG => value[val_offset..val_offset + val_length].to_vec(),
_ => {
let mut buf = Vec::with_capacity(8 + val_length);
buf.extend_from_slice(&SCALAR_CONTAINER_TAG.to_be_bytes());
buf.extend_from_slice(&encoded.to_be_bytes());
if jentry.length > 0 {
buf.extend_from_slice(&value[val_offset..val_offset + val_length]);
}
buf
}
};
return Some(val);
}
None
get_jentry_by_index(value, 0, header, index).map(|(jentry, encoded, val_offset)| {
extract_by_jentry(&jentry, encoded, val_offset, value)
})
}
_ => None,
}
Expand All @@ -288,64 +262,166 @@ pub fn get_by_name(value: &[u8], name: &str, ignore_case: bool) -> Option<Vec<u8

let header = read_u32(value, 0).unwrap();
match header & CONTAINER_HEADER_TYPE_MASK {
OBJECT_CONTAINER_TAG => {
let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
let mut jentry_offset = 4;
let mut val_offset = 8 * length + 4;

let mut key_jentries: VecDeque<JEntry> = VecDeque::with_capacity(length);
for _ in 0..length {
let encoded = read_u32(value, jentry_offset).unwrap();
let key_jentry = JEntry::decode_jentry(encoded);

jentry_offset += 4;
val_offset += key_jentry.length as usize;
key_jentries.push_back(key_jentry);
}
OBJECT_CONTAINER_TAG => get_jentry_by_name(value, 0, header, name, ignore_case).map(
|(jentry, encoded, val_offset)| extract_by_jentry(&jentry, encoded, val_offset, value),
),
_ => None,
}
}

let mut offsets = None;
let mut key_offset = 8 * length + 4;
while let Some(key_jentry) = key_jentries.pop_front() {
let prev_key_offset = key_offset;
key_offset += key_jentry.length as usize;
let key =
unsafe { std::str::from_utf8_unchecked(&value[prev_key_offset..key_offset]) };
// first match the value with the same name, if not found,
// then match the value with the ignoring case name.
if name.eq(key) {
offsets = Some((jentry_offset, val_offset));
break;
} else if ignore_case && name.eq_ignore_ascii_case(key) && offsets.is_none() {
offsets = Some((jentry_offset, val_offset));
/// Extracts JSON sub-object at the specified path,
/// where path elements can be either field keys or array indexes encoded in utf-8 string.
pub fn get_by_keypath<'a, I: Iterator<Item = &'a [u8]>>(
value: &[u8],
keypath: I,
) -> Option<Vec<u8>> {
if !is_jsonb(value) {
return match parse_value(value) {
Ok(val) => {
let mut current_val = &val;
for key in keypath {
match from_utf8(key) {
Ok(k) => {
let res = match current_val {
Value::Array(arr) => match k.parse::<usize>() {
Ok(idx) => arr.get(idx),
Err(_) => None,
},
Value::Object(obj) => obj.get(k),
_ => None,
};
match res {
Some(v) => current_val = v,
None => return None,
};
}
Err(_) => return None,
}
}
let val_encoded = read_u32(value, jentry_offset).unwrap();
let val_jentry = JEntry::decode_jentry(val_encoded);
jentry_offset += 4;
val_offset += val_jentry.length as usize;
Some(current_val.to_vec())
}
if let Some((jentry_offset, val_offset)) = offsets {
let encoded = read_u32(value, jentry_offset).unwrap();
let jentry = JEntry::decode_jentry(encoded);
let val_length = jentry.length as usize;
let val = match jentry.type_code {
CONTAINER_TAG => value[val_offset..val_offset + val_length].to_vec(),
_ => {
let mut buf: Vec<u8> = Vec::with_capacity(val_length + 8);
let scalar_header = SCALAR_CONTAINER_TAG;
buf.extend_from_slice(&scalar_header.to_be_bytes());
buf.extend_from_slice(&encoded.to_be_bytes());
if val_length > 0 {
buf.extend_from_slice(&value[val_offset..val_offset + val_length]);
}
buf
Err(_) => None,
};
}

let mut curr_val_offset = 0;
let mut curr_jentry_encoded = 0;
let mut curr_jentry: Option<JEntry> = None;

for key in keypath {
match from_utf8(key) {
Ok(k) => {
if let Some(ref jentry) = curr_jentry {
if jentry.type_code != CONTAINER_TAG {
return None;
}
};
return Some(val);
let header = read_u32(value, curr_val_offset).unwrap();
match header & CONTAINER_HEADER_TYPE_MASK {
OBJECT_CONTAINER_TAG => {
match get_jentry_by_name(value, curr_val_offset, header, k, false) {
Some((jentry, encoded, value_offset)) => {
curr_jentry_encoded = encoded;
curr_jentry = Some(jentry);
curr_val_offset = value_offset;
}
None => return None,
};
}
ARRAY_CONTAINER_TAG => match k.parse::<usize>() {
Ok(idx) => match get_jentry_by_index(value, curr_val_offset, header, idx) {
Some((jentry, encoded, value_offset)) => {
curr_jentry_encoded = encoded;
curr_jentry = Some(jentry);
curr_val_offset = value_offset;
}
None => return None,
},
Err(_) => return None,
},
_ => return None,
}
}
None
Err(_) => return None,
}
_ => None,
}
curr_jentry
.map(|jentry| extract_by_jentry(&jentry, curr_jentry_encoded, curr_val_offset, value))
}

fn get_jentry_by_name(
value: &[u8],
offset: usize,
header: u32,
name: &str,
ignore_case: bool,
) -> Option<(JEntry, u32, usize)> {
let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
let mut jentry_offset = offset + 4;
let mut val_offset = offset + 8 * length + 4;

let mut key_jentries: VecDeque<JEntry> = VecDeque::with_capacity(length);
for _ in 0..length {
let encoded = read_u32(value, jentry_offset).unwrap();
let key_jentry = JEntry::decode_jentry(encoded);

jentry_offset += 4;
val_offset += key_jentry.length as usize;
key_jentries.push_back(key_jentry);
}

let mut result = None;
let mut key_offset = offset + 8 * length + 4;

while let Some(key_jentry) = key_jentries.pop_front() {
let prev_key_offset = key_offset;
key_offset += key_jentry.length as usize;
let key = unsafe { std::str::from_utf8_unchecked(&value[prev_key_offset..key_offset]) };

let val_encoded = read_u32(value, jentry_offset).unwrap();
let val_jentry = JEntry::decode_jentry(val_encoded);
let val_length = val_jentry.length as usize;

// first match the value with the same name, if not found,
// then match the value with the ignoring case name.
if name.eq(key) {
result = Some((val_jentry, val_encoded, val_offset));
break;
} else if ignore_case && name.eq_ignore_ascii_case(key) && result.is_none() {
result = Some((val_jentry, val_encoded, val_offset));
}

jentry_offset += 4;
val_offset += val_length;
}
result
}

fn get_jentry_by_index(
value: &[u8],
offset: usize,
header: u32,
index: usize,
) -> Option<(JEntry, u32, usize)> {
let length = (header & CONTAINER_HEADER_LEN_MASK) as usize;
if index >= length {
return None;
}
let mut jentry_offset = offset + 4;
let mut val_offset = offset + 4 * length + 4;

for i in 0..length {
let encoded = read_u32(value, jentry_offset).unwrap();
let jentry = JEntry::decode_jentry(encoded);
let val_length = jentry.length as usize;
if i < index {
jentry_offset += 4;
val_offset += val_length;
continue;
}
return Some((jentry, encoded, val_offset));
}
None
}

/// Get the keys of a `JSONB` object.
Expand Down Expand Up @@ -435,18 +511,7 @@ pub fn object_each(value: &[u8]) -> Option<Vec<(Vec<u8>, Vec<u8>)>> {
let (jentry, encoded) = jentries.pop_front().unwrap();
let key = keys.pop_front().unwrap();
let val_length = jentry.length as usize;
let val = match jentry.type_code {
CONTAINER_TAG => value[offset..offset + val_length].to_vec(),
_ => {
let mut buf = Vec::with_capacity(val_length + 8);
buf.extend_from_slice(&SCALAR_CONTAINER_TAG.to_be_bytes());
buf.extend_from_slice(&encoded.to_be_bytes());
if jentry.length > 0 {
buf.extend_from_slice(&value[offset..offset + val_length]);
}
buf
}
};
let val = extract_by_jentry(&jentry, encoded, offset, value);
offset += val_length;
items.push((key, val));
}
Expand Down Expand Up @@ -481,18 +546,7 @@ pub fn array_values(value: &[u8]) -> Option<Vec<Vec<u8>>> {
let encoded = read_u32(value, jentry_offset).unwrap();
let jentry = JEntry::decode_jentry(encoded);
let val_length = jentry.length as usize;
let item = match jentry.type_code {
CONTAINER_TAG => value[val_offset..val_offset + val_length].to_vec(),
_ => {
let mut buf = Vec::with_capacity(8 + val_length);
buf.extend_from_slice(&SCALAR_CONTAINER_TAG.to_be_bytes());
buf.extend_from_slice(&encoded.to_be_bytes());
if jentry.length > 0 {
buf.extend_from_slice(&value[val_offset..val_offset + val_length]);
}
buf
}
};
let item = extract_by_jentry(&jentry, encoded, val_offset, value);
items.push(item);

jentry_offset += 4;
Expand Down Expand Up @@ -591,6 +645,22 @@ pub fn compare(left: &[u8], right: &[u8]) -> Result<Ordering, Error> {
}
}

fn extract_by_jentry(jentry: &JEntry, encoded: u32, offset: usize, value: &[u8]) -> Vec<u8> {
let length = jentry.length as usize;
match jentry.type_code {
CONTAINER_TAG => value[offset..offset + length].to_vec(),
_ => {
let mut buf = Vec::with_capacity(8 + length);
buf.extend_from_slice(&SCALAR_CONTAINER_TAG.to_be_bytes());
buf.extend_from_slice(&encoded.to_be_bytes());
if jentry.length > 0 {
buf.extend_from_slice(&value[offset..offset + length]);
}
buf
}
}
}

// Different types of values have different levels and are definitely not equal
fn jentry_compare_level(jentry: &JEntry) -> u8 {
match jentry.type_code {
Expand Down
Loading