use std::collections::VecDeque;
use std::ops::Range;
use super::constants::*;
use super::jentry::JEntry;
use crate::core::databend::util::jentry_to_jsonb_item;
use crate::core::JsonbItem;
use crate::error::Result;
use crate::RawJsonb;
pub(crate) struct ArrayIterator<'a> {
raw_jsonb: RawJsonb<'a>,
jentry_offset: usize,
item_offset: usize,
length: usize,
index: usize,
}
impl<'a> ArrayIterator<'a> {
pub(crate) fn new(raw_jsonb: RawJsonb<'a>) -> Result<Option<Self>> {
let (header_type, header_len) = raw_jsonb.read_header(0)?;
if header_type == ARRAY_CONTAINER_TAG {
let jentry_offset = 4;
let item_offset = 4 + 4 * header_len;
Ok(Some(Self {
raw_jsonb,
jentry_offset,
item_offset,
length: header_len,
index: 0,
}))
} else {
Ok(None)
}
}
pub(crate) fn len(&self) -> usize {
self.length
}
}
impl<'a> Iterator for ArrayIterator<'a> {
type Item = Result<JsonbItem<'a>>;
fn next(&mut self) -> Option<Self::Item> {
if self.index >= self.length {
return None;
}
let jentry = match self.raw_jsonb.read_jentry(self.jentry_offset) {
Ok(jentry) => jentry,
Err(err) => return Some(Err(err)),
};
let item_length = jentry.length as usize;
let item_range = Range {
start: self.item_offset,
end: self.item_offset + item_length,
};
let data = match self.raw_jsonb.slice(item_range) {
Ok(data) => data,
Err(err) => return Some(Err(err)),
};
let item = jentry_to_jsonb_item(jentry, data);
self.index += 1;
self.jentry_offset += 4;
self.item_offset += item_length;
Some(Ok(item))
}
}
pub(crate) struct ObjectKeyIterator<'a> {
raw_jsonb: RawJsonb<'a>,
jentry_offset: usize,
item_offset: usize,
length: usize,
index: usize,
}
impl<'a> ObjectKeyIterator<'a> {
pub(crate) fn new(raw_jsonb: RawJsonb<'a>) -> Result<Option<Self>> {
let (header_type, header_len) = raw_jsonb.read_header(0)?;
if header_type == OBJECT_CONTAINER_TAG {
let jentry_offset = 4;
let item_offset = 4 + 8 * header_len;
Ok(Some(Self {
raw_jsonb,
jentry_offset,
item_offset,
length: header_len,
index: 0,
}))
} else {
Ok(None)
}
}
pub(crate) fn len(&self) -> usize {
self.length
}
}
impl<'a> Iterator for ObjectKeyIterator<'a> {
type Item = Result<JsonbItem<'a>>;
fn next(&mut self) -> Option<Self::Item> {
if self.index >= self.length {
return None;
}
let jentry = match self.raw_jsonb.read_jentry(self.jentry_offset) {
Ok(jentry) => jentry,
Err(err) => return Some(Err(err)),
};
let key_length = jentry.length as usize;
let key_range = Range {
start: self.item_offset,
end: self.item_offset + key_length,
};
let data = match self.raw_jsonb.slice(key_range) {
Ok(data) => data,
Err(err) => return Some(Err(err)),
};
let key_item = jentry_to_jsonb_item(jentry, data);
self.index += 1;
self.jentry_offset += 4;
self.item_offset += key_length;
Some(Ok(key_item))
}
}
pub(crate) struct ObjectValueIterator<'a> {
raw_jsonb: RawJsonb<'a>,
jentry_offset: usize,
item_offset: usize,
length: usize,
index: usize,
}
impl<'a> ObjectValueIterator<'a> {
pub(crate) fn new(raw_jsonb: RawJsonb<'a>) -> Result<Option<Self>> {
let (header_type, header_len) = raw_jsonb.read_header(0)?;
if header_type == OBJECT_CONTAINER_TAG {
let mut jentry_offset = 4;
let mut item_offset = 4 + 8 * header_len;
for _ in 0..header_len {
let key_jentry = raw_jsonb.read_jentry(jentry_offset)?;
jentry_offset += 4;
item_offset += key_jentry.length as usize;
}
Ok(Some(Self {
raw_jsonb,
jentry_offset,
item_offset,
length: header_len,
index: 0,
}))
} else {
Ok(None)
}
}
#[allow(dead_code)]
pub(crate) fn len(&self) -> usize {
self.length
}
}
impl<'a> Iterator for ObjectValueIterator<'a> {
type Item = Result<JsonbItem<'a>>;
fn next(&mut self) -> Option<Self::Item> {
if self.index >= self.length {
return None;
}
let jentry = match self.raw_jsonb.read_jentry(self.jentry_offset) {
Ok(jentry) => jentry,
Err(err) => return Some(Err(err)),
};
let val_length = jentry.length as usize;
let val_range = Range {
start: self.item_offset,
end: self.item_offset + val_length,
};
let data = match self.raw_jsonb.slice(val_range) {
Ok(data) => data,
Err(err) => return Some(Err(err)),
};
let val_item = jentry_to_jsonb_item(jentry, data);
self.index += 1;
self.jentry_offset += 4;
self.item_offset += val_length;
Some(Ok(val_item))
}
}
pub(crate) struct ObjectIterator<'a> {
raw_jsonb: RawJsonb<'a>,
key_jentries: VecDeque<JEntry>,
jentry_offset: usize,
key_offset: usize,
val_offset: usize,
length: usize,
}
impl<'a> ObjectIterator<'a> {
pub(crate) fn new(raw_jsonb: RawJsonb<'a>) -> Result<Option<Self>> {
let (header_type, header_len) = raw_jsonb.read_header(0)?;
if header_type == OBJECT_CONTAINER_TAG {
let mut jentry_offset = 4;
let mut key_jentries = VecDeque::with_capacity(header_len);
for _ in 0..header_len {
let key_jentry = raw_jsonb.read_jentry(jentry_offset)?;
jentry_offset += 4;
key_jentries.push_back(key_jentry);
}
let key_length: usize = key_jentries.iter().map(|j| j.length as usize).sum();
let key_offset = 4 + 8 * header_len;
let val_offset = key_offset + key_length;
Ok(Some(Self {
raw_jsonb,
key_jentries,
jentry_offset,
key_offset,
val_offset,
length: header_len,
}))
} else {
Ok(None)
}
}
pub(crate) fn len(&self) -> usize {
self.length
}
}
impl<'a> Iterator for ObjectIterator<'a> {
type Item = Result<(&'a str, JsonbItem<'a>)>;
fn next(&mut self) -> Option<Self::Item> {
match self.key_jentries.pop_front() {
Some(key_jentry) => {
let val_jentry = match self.raw_jsonb.read_jentry(self.jentry_offset) {
Ok(jentry) => jentry,
Err(err) => return Some(Err(err)),
};
let key_length = key_jentry.length as usize;
let val_length = val_jentry.length as usize;
let key_range = Range {
start: self.key_offset,
end: self.key_offset + key_length,
};
let key_data = match self.raw_jsonb.slice(key_range) {
Ok(data) => data,
Err(err) => return Some(Err(err)),
};
let key = unsafe { std::str::from_utf8_unchecked(key_data) };
let val_range = Range {
start: self.val_offset,
end: self.val_offset + val_length,
};
let val_data = match self.raw_jsonb.slice(val_range) {
Ok(data) => data,
Err(err) => return Some(Err(err)),
};
let val_item = jentry_to_jsonb_item(val_jentry, val_data);
self.jentry_offset += 4;
self.key_offset += key_length;
self.val_offset += val_length;
Some(Ok((key, val_item)))
}
None => None,
}
}
}